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