1use anyhow::anyhow;
2use serde::{Deserialize, Serialize};
3use strum::EnumIter;
4
5#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
6#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
7pub enum BedrockModelMode {
8 #[default]
9 Default,
10 Thinking {
11 budget_tokens: Option<u64>,
12 },
13}
14
15#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
16#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)]
17pub enum Model {
18 // Anthropic models (already included)
19 #[serde(rename = "claude-sonnet-4", alias = "claude-sonnet-4-latest")]
20 ClaudeSonnet4,
21 #[serde(
22 rename = "claude-sonnet-4-thinking",
23 alias = "claude-sonnet-4-thinking-latest"
24 )]
25 ClaudeSonnet4Thinking,
26 #[serde(rename = "claude-opus-4", alias = "claude-opus-4-latest")]
27 ClaudeOpus4,
28 #[serde(
29 rename = "claude-opus-4-thinking",
30 alias = "claude-opus-4-thinking-latest"
31 )]
32 ClaudeOpus4Thinking,
33 #[default]
34 #[serde(rename = "claude-3-5-sonnet-v2", alias = "claude-3-5-sonnet-latest")]
35 Claude3_5SonnetV2,
36 #[serde(rename = "claude-3-7-sonnet", alias = "claude-3-7-sonnet-latest")]
37 Claude3_7Sonnet,
38 #[serde(
39 rename = "claude-3-7-sonnet-thinking",
40 alias = "claude-3-7-sonnet-thinking-latest"
41 )]
42 Claude3_7SonnetThinking,
43 #[serde(rename = "claude-3-opus", alias = "claude-3-opus-latest")]
44 Claude3Opus,
45 #[serde(rename = "claude-3-sonnet", alias = "claude-3-sonnet-latest")]
46 Claude3Sonnet,
47 #[serde(rename = "claude-3-5-haiku", alias = "claude-3-5-haiku-latest")]
48 Claude3_5Haiku,
49 Claude3_5Sonnet,
50 Claude3Haiku,
51 // Amazon Nova Models
52 AmazonNovaLite,
53 AmazonNovaMicro,
54 AmazonNovaPro,
55 AmazonNovaPremier,
56 // AI21 models
57 AI21J2GrandeInstruct,
58 AI21J2JumboInstruct,
59 AI21J2Mid,
60 AI21J2MidV1,
61 AI21J2Ultra,
62 AI21J2UltraV1_8k,
63 AI21J2UltraV1,
64 AI21JambaInstructV1,
65 AI21Jamba15LargeV1,
66 AI21Jamba15MiniV1,
67 // Cohere models
68 CohereCommandTextV14_4k,
69 CohereCommandRV1,
70 CohereCommandRPlusV1,
71 CohereCommandLightTextV14_4k,
72 // DeepSeek
73 DeepSeekR1,
74 // Meta models
75 MetaLlama38BInstructV1,
76 MetaLlama370BInstructV1,
77 MetaLlama318BInstructV1_128k,
78 MetaLlama318BInstructV1,
79 MetaLlama3170BInstructV1_128k,
80 MetaLlama3170BInstructV1,
81 MetaLlama3211BInstructV1,
82 MetaLlama3290BInstructV1,
83 MetaLlama321BInstructV1,
84 MetaLlama323BInstructV1,
85 // Mistral models
86 MistralMistral7BInstructV0,
87 MistralMixtral8x7BInstructV0,
88 MistralMistralLarge2402V1,
89 MistralMistralSmall2402V1,
90 MistralPixtralLarge2502V1,
91 // Writer models
92 PalmyraWriterX5,
93 PalmyraWriterX4,
94 #[serde(rename = "custom")]
95 Custom {
96 name: String,
97 max_tokens: usize,
98 /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
99 display_name: Option<String>,
100 max_output_tokens: Option<u32>,
101 default_temperature: Option<f32>,
102 },
103}
104
105impl Model {
106 pub fn default_fast() -> Self {
107 Self::Claude3_5Haiku
108 }
109
110 pub fn from_id(id: &str) -> anyhow::Result<Self> {
111 if id.starts_with("claude-3-5-sonnet-v2") {
112 Ok(Self::Claude3_5SonnetV2)
113 } else if id.starts_with("claude-3-opus") {
114 Ok(Self::Claude3Opus)
115 } else if id.starts_with("claude-3-sonnet") {
116 Ok(Self::Claude3Sonnet)
117 } else if id.starts_with("claude-3-5-haiku") {
118 Ok(Self::Claude3_5Haiku)
119 } else if id.starts_with("claude-3-7-sonnet") {
120 Ok(Self::Claude3_7Sonnet)
121 } else if id.starts_with("claude-3-7-sonnet-thinking") {
122 Ok(Self::Claude3_7SonnetThinking)
123 } else {
124 Err(anyhow!("invalid model id"))
125 }
126 }
127
128 pub fn id(&self) -> &str {
129 match self {
130 Model::ClaudeSonnet4 | Model::ClaudeSonnet4Thinking => {
131 "anthropic.claude-sonnet-4-20250514-v1:0"
132 }
133 Model::ClaudeOpus4 | Model::ClaudeOpus4Thinking => {
134 "anthropic.claude-opus-4-20250514-v1:0"
135 }
136 Model::Claude3_5SonnetV2 => "anthropic.claude-3-5-sonnet-20241022-v2:0",
137 Model::Claude3_5Sonnet => "anthropic.claude-3-5-sonnet-20240620-v1:0",
138 Model::Claude3Opus => "anthropic.claude-3-opus-20240229-v1:0",
139 Model::Claude3Sonnet => "anthropic.claude-3-sonnet-20240229-v1:0",
140 Model::Claude3Haiku => "anthropic.claude-3-haiku-20240307-v1:0",
141 Model::Claude3_5Haiku => "anthropic.claude-3-5-haiku-20241022-v1:0",
142 Model::Claude3_7Sonnet | Model::Claude3_7SonnetThinking => {
143 "anthropic.claude-3-7-sonnet-20250219-v1:0"
144 }
145 Model::AmazonNovaLite => "amazon.nova-lite-v1:0",
146 Model::AmazonNovaMicro => "amazon.nova-micro-v1:0",
147 Model::AmazonNovaPro => "amazon.nova-pro-v1:0",
148 Model::AmazonNovaPremier => "amazon.nova-premier-v1:0",
149 Model::DeepSeekR1 => "us.deepseek.r1-v1:0",
150 Model::AI21J2GrandeInstruct => "ai21.j2-grande-instruct",
151 Model::AI21J2JumboInstruct => "ai21.j2-jumbo-instruct",
152 Model::AI21J2Mid => "ai21.j2-mid",
153 Model::AI21J2MidV1 => "ai21.j2-mid-v1",
154 Model::AI21J2Ultra => "ai21.j2-ultra",
155 Model::AI21J2UltraV1_8k => "ai21.j2-ultra-v1:0:8k",
156 Model::AI21J2UltraV1 => "ai21.j2-ultra-v1",
157 Model::AI21JambaInstructV1 => "ai21.jamba-instruct-v1:0",
158 Model::AI21Jamba15LargeV1 => "ai21.jamba-1-5-large-v1:0",
159 Model::AI21Jamba15MiniV1 => "ai21.jamba-1-5-mini-v1:0",
160 Model::CohereCommandTextV14_4k => "cohere.command-text-v14:7:4k",
161 Model::CohereCommandRV1 => "cohere.command-r-v1:0",
162 Model::CohereCommandRPlusV1 => "cohere.command-r-plus-v1:0",
163 Model::CohereCommandLightTextV14_4k => "cohere.command-light-text-v14:7:4k",
164 Model::MetaLlama38BInstructV1 => "meta.llama3-8b-instruct-v1:0",
165 Model::MetaLlama370BInstructV1 => "meta.llama3-70b-instruct-v1:0",
166 Model::MetaLlama318BInstructV1_128k => "meta.llama3-1-8b-instruct-v1:0:128k",
167 Model::MetaLlama318BInstructV1 => "meta.llama3-1-8b-instruct-v1:0",
168 Model::MetaLlama3170BInstructV1_128k => "meta.llama3-1-70b-instruct-v1:0:128k",
169 Model::MetaLlama3170BInstructV1 => "meta.llama3-1-70b-instruct-v1:0",
170 Model::MetaLlama3211BInstructV1 => "meta.llama3-2-11b-instruct-v1:0",
171 Model::MetaLlama3290BInstructV1 => "meta.llama3-2-90b-instruct-v1:0",
172 Model::MetaLlama321BInstructV1 => "meta.llama3-2-1b-instruct-v1:0",
173 Model::MetaLlama323BInstructV1 => "meta.llama3-2-3b-instruct-v1:0",
174 Model::MistralMistral7BInstructV0 => "mistral.mistral-7b-instruct-v0:2",
175 Model::MistralMixtral8x7BInstructV0 => "mistral.mixtral-8x7b-instruct-v0:1",
176 Model::MistralMistralLarge2402V1 => "mistral.mistral-large-2402-v1:0",
177 Model::MistralMistralSmall2402V1 => "mistral.mistral-small-2402-v1:0",
178 Model::MistralPixtralLarge2502V1 => "mistral.pixtral-large-2502-v1:0",
179 Model::PalmyraWriterX4 => "writer.palmyra-x4-v1:0",
180 Model::PalmyraWriterX5 => "writer.palmyra-x5-v1:0",
181 Self::Custom { name, .. } => name,
182 }
183 }
184
185 pub fn display_name(&self) -> &str {
186 match self {
187 Self::ClaudeSonnet4 => "Claude Sonnet 4",
188 Self::ClaudeSonnet4Thinking => "Claude Sonnet 4 Thinking",
189 Self::ClaudeOpus4 => "Claude Opus 4",
190 Self::ClaudeOpus4Thinking => "Claude Opus 4 Thinking",
191 Self::Claude3_5SonnetV2 => "Claude 3.5 Sonnet v2",
192 Self::Claude3_5Sonnet => "Claude 3.5 Sonnet",
193 Self::Claude3Opus => "Claude 3 Opus",
194 Self::Claude3Sonnet => "Claude 3 Sonnet",
195 Self::Claude3Haiku => "Claude 3 Haiku",
196 Self::Claude3_5Haiku => "Claude 3.5 Haiku",
197 Self::Claude3_7Sonnet => "Claude 3.7 Sonnet",
198 Self::Claude3_7SonnetThinking => "Claude 3.7 Sonnet Thinking",
199 Self::AmazonNovaLite => "Amazon Nova Lite",
200 Self::AmazonNovaMicro => "Amazon Nova Micro",
201 Self::AmazonNovaPro => "Amazon Nova Pro",
202 Self::AmazonNovaPremier => "Amazon Nova Premier",
203 Self::DeepSeekR1 => "DeepSeek R1",
204 Self::AI21J2GrandeInstruct => "AI21 Jurassic2 Grande Instruct",
205 Self::AI21J2JumboInstruct => "AI21 Jurassic2 Jumbo Instruct",
206 Self::AI21J2Mid => "AI21 Jurassic2 Mid",
207 Self::AI21J2MidV1 => "AI21 Jurassic2 Mid V1",
208 Self::AI21J2Ultra => "AI21 Jurassic2 Ultra",
209 Self::AI21J2UltraV1_8k => "AI21 Jurassic2 Ultra V1 8K",
210 Self::AI21J2UltraV1 => "AI21 Jurassic2 Ultra V1",
211 Self::AI21JambaInstructV1 => "AI21 Jamba Instruct",
212 Self::AI21Jamba15LargeV1 => "AI21 Jamba 1.5 Large",
213 Self::AI21Jamba15MiniV1 => "AI21 Jamba 1.5 Mini",
214 Self::CohereCommandTextV14_4k => "Cohere Command Text V14 4K",
215 Self::CohereCommandRV1 => "Cohere Command R V1",
216 Self::CohereCommandRPlusV1 => "Cohere Command R Plus V1",
217 Self::CohereCommandLightTextV14_4k => "Cohere Command Light Text V14 4K",
218 Self::MetaLlama38BInstructV1 => "Meta Llama 3 8B Instruct V1",
219 Self::MetaLlama370BInstructV1 => "Meta Llama 3 70B Instruct V1",
220 Self::MetaLlama318BInstructV1_128k => "Meta Llama 3 1.8B Instruct V1 128K",
221 Self::MetaLlama318BInstructV1 => "Meta Llama 3 1.8B Instruct V1",
222 Self::MetaLlama3170BInstructV1_128k => "Meta Llama 3 1 70B Instruct V1 128K",
223 Self::MetaLlama3170BInstructV1 => "Meta Llama 3 1 70B Instruct V1",
224 Self::MetaLlama3211BInstructV1 => "Meta Llama 3 2 11B Instruct V1",
225 Self::MetaLlama3290BInstructV1 => "Meta Llama 3 2 90B Instruct V1",
226 Self::MetaLlama321BInstructV1 => "Meta Llama 3 2 1B Instruct V1",
227 Self::MetaLlama323BInstructV1 => "Meta Llama 3 2 3B Instruct V1",
228 Self::MistralMistral7BInstructV0 => "Mistral 7B Instruct V0",
229 Self::MistralMixtral8x7BInstructV0 => "Mistral Mixtral 8x7B Instruct V0",
230 Self::MistralMistralLarge2402V1 => "Mistral Large 2402 V1",
231 Self::MistralMistralSmall2402V1 => "Mistral Small 2402 V1",
232 Self::MistralPixtralLarge2502V1 => "Pixtral Large 25.02 V1",
233 Self::PalmyraWriterX5 => "Writer Palmyra X5",
234 Self::PalmyraWriterX4 => "Writer Palmyra X4",
235 Self::Custom {
236 display_name, name, ..
237 } => display_name.as_deref().unwrap_or(name),
238 }
239 }
240
241 pub fn max_token_count(&self) -> usize {
242 match self {
243 Self::Claude3_5SonnetV2
244 | Self::Claude3Opus
245 | Self::Claude3Sonnet
246 | Self::Claude3_5Haiku
247 | Self::Claude3_7Sonnet
248 | Self::ClaudeSonnet4
249 | Self::ClaudeOpus4 => 200_000,
250 Self::AmazonNovaPremier => 1_000_000,
251 Self::PalmyraWriterX5 => 1_000_000,
252 Self::PalmyraWriterX4 => 128_000,
253 Self::Custom { max_tokens, .. } => *max_tokens,
254 _ => 128_000,
255 }
256 }
257
258 pub fn max_output_tokens(&self) -> u32 {
259 match self {
260 Self::Claude3Opus | Self::Claude3Sonnet | Self::Claude3_5Haiku => 4_096,
261 Self::Claude3_7Sonnet
262 | Self::Claude3_7SonnetThinking
263 | Self::ClaudeSonnet4
264 | Self::ClaudeSonnet4Thinking
265 | Self::ClaudeOpus4
266 | Model::ClaudeOpus4Thinking => 128_000,
267 Self::Claude3_5SonnetV2 | Self::PalmyraWriterX4 | Self::PalmyraWriterX5 => 8_192,
268 Self::Custom {
269 max_output_tokens, ..
270 } => max_output_tokens.unwrap_or(4_096),
271 _ => 4_096,
272 }
273 }
274
275 pub fn default_temperature(&self) -> f32 {
276 match self {
277 Self::Claude3_5SonnetV2
278 | Self::Claude3Opus
279 | Self::Claude3Sonnet
280 | Self::Claude3_5Haiku
281 | Self::Claude3_7Sonnet
282 | Self::ClaudeOpus4
283 | Self::ClaudeOpus4Thinking
284 | Self::ClaudeSonnet4
285 | Self::ClaudeSonnet4Thinking => 1.0,
286 Self::Custom {
287 default_temperature,
288 ..
289 } => default_temperature.unwrap_or(1.0),
290 _ => 1.0,
291 }
292 }
293
294 pub fn supports_tool_use(&self) -> bool {
295 match self {
296 // Anthropic Claude 3 models (all support tool use)
297 Self::Claude3Opus
298 | Self::Claude3Sonnet
299 | Self::Claude3_5Sonnet
300 | Self::Claude3_5SonnetV2
301 | Self::Claude3_7Sonnet
302 | Self::Claude3_7SonnetThinking
303 | Self::ClaudeOpus4
304 | Self::ClaudeOpus4Thinking
305 | Self::ClaudeSonnet4
306 | Self::ClaudeSonnet4Thinking
307 | Self::Claude3_5Haiku => true,
308
309 // Amazon Nova models (all support tool use)
310 Self::AmazonNovaPremier
311 | Self::AmazonNovaPro
312 | Self::AmazonNovaLite
313 | Self::AmazonNovaMicro => true,
314
315 // AI21 Jamba 1.5 models support tool use
316 Self::AI21Jamba15LargeV1 | Self::AI21Jamba15MiniV1 => true,
317
318 // Cohere Command R models support tool use
319 Self::CohereCommandRV1 | Self::CohereCommandRPlusV1 => true,
320
321 // All other models don't support tool use
322 // Including Meta Llama 3.2, AI21 Jurassic, and others
323 _ => false,
324 }
325 }
326
327 pub fn mode(&self) -> BedrockModelMode {
328 match self {
329 Model::Claude3_7SonnetThinking => BedrockModelMode::Thinking {
330 budget_tokens: Some(4096),
331 },
332 Model::ClaudeSonnet4Thinking => BedrockModelMode::Thinking {
333 budget_tokens: Some(4096),
334 },
335 Model::ClaudeOpus4Thinking => BedrockModelMode::Thinking {
336 budget_tokens: Some(4096),
337 },
338 _ => BedrockModelMode::Default,
339 }
340 }
341
342 pub fn cross_region_inference_id(&self, region: &str) -> Result<String, anyhow::Error> {
343 let region_group = if region.starts_with("us-gov-") {
344 "us-gov"
345 } else if region.starts_with("us-") {
346 "us"
347 } else if region.starts_with("eu-") {
348 "eu"
349 } else if region.starts_with("ap-") || region == "me-central-1" || region == "me-south-1" {
350 "apac"
351 } else if region.starts_with("ca-") || region.starts_with("sa-") {
352 // Canada and South America regions - default to US profiles
353 "us"
354 } else {
355 // Unknown region
356 return Err(anyhow!("Unsupported Region"));
357 };
358
359 let model_id = self.id();
360
361 match (self, region_group) {
362 // Custom models can't have CRI IDs
363 (Model::Custom { .. }, _) => Ok(self.id().into()),
364
365 // Models with US Gov only
366 (Model::Claude3_5Sonnet, "us-gov") | (Model::Claude3Haiku, "us-gov") => {
367 Ok(format!("{}.{}", region_group, model_id))
368 }
369
370 // Models available only in US
371 (Model::Claude3Opus, "us")
372 | (Model::Claude3_5Haiku, "us")
373 | (Model::Claude3_7Sonnet, "us")
374 | (Model::ClaudeSonnet4, "us")
375 | (Model::ClaudeOpus4, "us")
376 | (Model::ClaudeSonnet4Thinking, "us")
377 | (Model::ClaudeOpus4Thinking, "us")
378 | (Model::Claude3_7SonnetThinking, "us")
379 | (Model::AmazonNovaPremier, "us")
380 | (Model::MistralPixtralLarge2502V1, "us") => {
381 Ok(format!("{}.{}", region_group, model_id))
382 }
383
384 // Models available in US, EU, and APAC
385 (Model::Claude3_5SonnetV2, "us")
386 | (Model::Claude3_5SonnetV2, "apac")
387 | (Model::Claude3_5Sonnet, _)
388 | (Model::Claude3Haiku, _)
389 | (Model::Claude3Sonnet, _)
390 | (Model::AmazonNovaLite, _)
391 | (Model::AmazonNovaMicro, _)
392 | (Model::AmazonNovaPro, _) => Ok(format!("{}.{}", region_group, model_id)),
393
394 // Models with limited EU availability
395 (Model::MetaLlama321BInstructV1, "us")
396 | (Model::MetaLlama321BInstructV1, "eu")
397 | (Model::MetaLlama323BInstructV1, "us")
398 | (Model::MetaLlama323BInstructV1, "eu") => {
399 Ok(format!("{}.{}", region_group, model_id))
400 }
401
402 // US-only models (all remaining Meta models)
403 (Model::MetaLlama38BInstructV1, "us")
404 | (Model::MetaLlama370BInstructV1, "us")
405 | (Model::MetaLlama318BInstructV1, "us")
406 | (Model::MetaLlama318BInstructV1_128k, "us")
407 | (Model::MetaLlama3170BInstructV1, "us")
408 | (Model::MetaLlama3170BInstructV1_128k, "us")
409 | (Model::MetaLlama3211BInstructV1, "us")
410 | (Model::MetaLlama3290BInstructV1, "us") => {
411 Ok(format!("{}.{}", region_group, model_id))
412 }
413
414 // Writer models only available in the US
415 (Model::PalmyraWriterX4, "us") | (Model::PalmyraWriterX5, "us") => {
416 // They have some goofiness
417 Ok(format!("{}.{}", region_group, model_id))
418 }
419
420 // Any other combination is not supported
421 _ => Ok(self.id().into()),
422 }
423 }
424}
425
426#[cfg(test)]
427mod tests {
428 use super::*;
429
430 #[test]
431 fn test_us_region_inference_ids() -> anyhow::Result<()> {
432 // Test US regions
433 assert_eq!(
434 Model::Claude3_5SonnetV2.cross_region_inference_id("us-east-1")?,
435 "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
436 );
437 assert_eq!(
438 Model::Claude3_5SonnetV2.cross_region_inference_id("us-west-2")?,
439 "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
440 );
441 assert_eq!(
442 Model::AmazonNovaPro.cross_region_inference_id("us-east-2")?,
443 "us.amazon.nova-pro-v1:0"
444 );
445 Ok(())
446 }
447
448 #[test]
449 fn test_eu_region_inference_ids() -> anyhow::Result<()> {
450 // Test European regions
451 assert_eq!(
452 Model::Claude3Sonnet.cross_region_inference_id("eu-west-1")?,
453 "eu.anthropic.claude-3-sonnet-20240229-v1:0"
454 );
455 assert_eq!(
456 Model::AmazonNovaMicro.cross_region_inference_id("eu-north-1")?,
457 "eu.amazon.nova-micro-v1:0"
458 );
459 Ok(())
460 }
461
462 #[test]
463 fn test_apac_region_inference_ids() -> anyhow::Result<()> {
464 // Test Asia-Pacific regions
465 assert_eq!(
466 Model::Claude3_5SonnetV2.cross_region_inference_id("ap-northeast-1")?,
467 "apac.anthropic.claude-3-5-sonnet-20241022-v2:0"
468 );
469 assert_eq!(
470 Model::AmazonNovaLite.cross_region_inference_id("ap-south-1")?,
471 "apac.amazon.nova-lite-v1:0"
472 );
473 Ok(())
474 }
475
476 #[test]
477 fn test_gov_region_inference_ids() -> anyhow::Result<()> {
478 // Test Government regions
479 assert_eq!(
480 Model::Claude3_5Sonnet.cross_region_inference_id("us-gov-east-1")?,
481 "us-gov.anthropic.claude-3-5-sonnet-20240620-v1:0"
482 );
483 assert_eq!(
484 Model::Claude3Haiku.cross_region_inference_id("us-gov-west-1")?,
485 "us-gov.anthropic.claude-3-haiku-20240307-v1:0"
486 );
487 Ok(())
488 }
489
490 #[test]
491 fn test_meta_models_inference_ids() -> anyhow::Result<()> {
492 // Test Meta models
493 assert_eq!(
494 Model::MetaLlama370BInstructV1.cross_region_inference_id("us-east-1")?,
495 "us.meta.llama3-70b-instruct-v1:0"
496 );
497 assert_eq!(
498 Model::MetaLlama321BInstructV1.cross_region_inference_id("eu-west-1")?,
499 "eu.meta.llama3-2-1b-instruct-v1:0"
500 );
501 Ok(())
502 }
503
504 #[test]
505 fn test_mistral_models_inference_ids() -> anyhow::Result<()> {
506 // Mistral models don't follow the regional prefix pattern,
507 // so they should return their original IDs
508 assert_eq!(
509 Model::MistralMistralLarge2402V1.cross_region_inference_id("us-east-1")?,
510 "mistral.mistral-large-2402-v1:0"
511 );
512 assert_eq!(
513 Model::MistralMixtral8x7BInstructV0.cross_region_inference_id("eu-west-1")?,
514 "mistral.mixtral-8x7b-instruct-v0:1"
515 );
516 Ok(())
517 }
518
519 #[test]
520 fn test_ai21_models_inference_ids() -> anyhow::Result<()> {
521 // AI21 models don't follow the regional prefix pattern,
522 // so they should return their original IDs
523 assert_eq!(
524 Model::AI21J2UltraV1.cross_region_inference_id("us-east-1")?,
525 "ai21.j2-ultra-v1"
526 );
527 assert_eq!(
528 Model::AI21JambaInstructV1.cross_region_inference_id("eu-west-1")?,
529 "ai21.jamba-instruct-v1:0"
530 );
531 Ok(())
532 }
533
534 #[test]
535 fn test_cohere_models_inference_ids() -> anyhow::Result<()> {
536 // Cohere models don't follow the regional prefix pattern,
537 // so they should return their original IDs
538 assert_eq!(
539 Model::CohereCommandRV1.cross_region_inference_id("us-east-1")?,
540 "cohere.command-r-v1:0"
541 );
542 assert_eq!(
543 Model::CohereCommandTextV14_4k.cross_region_inference_id("ap-southeast-1")?,
544 "cohere.command-text-v14:7:4k"
545 );
546 Ok(())
547 }
548
549 #[test]
550 fn test_custom_model_inference_ids() -> anyhow::Result<()> {
551 // Test custom models
552 let custom_model = Model::Custom {
553 name: "custom.my-model-v1:0".to_string(),
554 max_tokens: 100000,
555 display_name: Some("My Custom Model".to_string()),
556 max_output_tokens: Some(8192),
557 default_temperature: Some(0.7),
558 };
559
560 // Custom model should return its name unchanged
561 assert_eq!(
562 custom_model.cross_region_inference_id("us-east-1")?,
563 "custom.my-model-v1:0"
564 );
565
566 Ok(())
567 }
568}