models.rs

  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    MetaLlama38BInstructV1,
 75    MetaLlama370BInstructV1,
 76    MetaLlama318BInstructV1_128k,
 77    MetaLlama318BInstructV1,
 78    MetaLlama3170BInstructV1_128k,
 79    MetaLlama3170BInstructV1,
 80    MetaLlama31405BInstructV1,
 81    MetaLlama321BInstructV1,
 82    MetaLlama323BInstructV1,
 83    MetaLlama3211BInstructV1,
 84    MetaLlama3290BInstructV1,
 85    MetaLlama3370BInstructV1,
 86    #[allow(non_camel_case_types)]
 87    MetaLlama4Scout17BInstructV1,
 88    #[allow(non_camel_case_types)]
 89    MetaLlama4Maverick17BInstructV1,
 90    // Mistral models
 91    MistralMistral7BInstructV0,
 92    MistralMixtral8x7BInstructV0,
 93    MistralMistralLarge2402V1,
 94    MistralMistralSmall2402V1,
 95    MistralPixtralLarge2502V1,
 96    // Writer models
 97    PalmyraWriterX5,
 98    PalmyraWriterX4,
 99    #[serde(rename = "custom")]
100    Custom {
101        name: String,
102        max_tokens: u64,
103        /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
104        display_name: Option<String>,
105        max_output_tokens: Option<u64>,
106        default_temperature: Option<f32>,
107    },
108}
109
110impl Model {
111    pub fn default_fast(region: &str) -> Self {
112        if region.starts_with("us-") {
113            Self::Claude3_5Haiku
114        } else {
115            Self::Claude3Haiku
116        }
117    }
118
119    pub fn from_id(id: &str) -> anyhow::Result<Self> {
120        if id.starts_with("claude-3-5-sonnet-v2") {
121            Ok(Self::Claude3_5SonnetV2)
122        } else if id.starts_with("claude-3-opus") {
123            Ok(Self::Claude3Opus)
124        } else if id.starts_with("claude-3-sonnet") {
125            Ok(Self::Claude3Sonnet)
126        } else if id.starts_with("claude-3-5-haiku") {
127            Ok(Self::Claude3_5Haiku)
128        } else if id.starts_with("claude-3-7-sonnet") {
129            Ok(Self::Claude3_7Sonnet)
130        } else if id.starts_with("claude-3-7-sonnet-thinking") {
131            Ok(Self::Claude3_7SonnetThinking)
132        } else {
133            anyhow::bail!("invalid model id {id}");
134        }
135    }
136
137    pub fn id(&self) -> &str {
138        match self {
139            Model::ClaudeSonnet4 => "claude-4-sonnet",
140            Model::ClaudeSonnet4Thinking => "claude-4-sonnet-thinking",
141            Model::ClaudeOpus4 => "claude-4-opus",
142            Model::ClaudeOpus4Thinking => "claude-4-opus-thinking",
143            Model::Claude3_5SonnetV2 => "claude-3-5-sonnet-v2",
144            Model::Claude3_5Sonnet => "claude-3-5-sonnet",
145            Model::Claude3Opus => "claude-3-opus",
146            Model::Claude3Sonnet => "claude-3-sonnet",
147            Model::Claude3Haiku => "claude-3-haiku",
148            Model::Claude3_5Haiku => "claude-3-5-haiku",
149            Model::Claude3_7Sonnet => "claude-3-7-sonnet",
150            Model::Claude3_7SonnetThinking => "claude-3-7-sonnet-thinking",
151            Model::AmazonNovaLite => "amazon-nova-lite",
152            Model::AmazonNovaMicro => "amazon-nova-micro",
153            Model::AmazonNovaPro => "amazon-nova-pro",
154            Model::AmazonNovaPremier => "amazon-nova-premier",
155            Model::DeepSeekR1 => "deepseek-r1",
156            Model::AI21J2GrandeInstruct => "ai21-j2-grande-instruct",
157            Model::AI21J2JumboInstruct => "ai21-j2-jumbo-instruct",
158            Model::AI21J2Mid => "ai21-j2-mid",
159            Model::AI21J2MidV1 => "ai21-j2-mid-v1",
160            Model::AI21J2Ultra => "ai21-j2-ultra",
161            Model::AI21J2UltraV1_8k => "ai21-j2-ultra-v1-8k",
162            Model::AI21J2UltraV1 => "ai21-j2-ultra-v1",
163            Model::AI21JambaInstructV1 => "ai21-jamba-instruct-v1",
164            Model::AI21Jamba15LargeV1 => "ai21-jamba-1-5-large-v1",
165            Model::AI21Jamba15MiniV1 => "ai21-jamba-1-5-mini-v1",
166            Model::CohereCommandTextV14_4k => "cohere-command-text-v14-4k",
167            Model::CohereCommandRV1 => "cohere-command-r-v1",
168            Model::CohereCommandRPlusV1 => "cohere-command-r-plus-v1",
169            Model::CohereCommandLightTextV14_4k => "cohere-command-light-text-v14-4k",
170            Model::MetaLlama38BInstructV1 => "meta-llama3-8b-instruct-v1",
171            Model::MetaLlama370BInstructV1 => "meta-llama3-70b-instruct-v1",
172            Model::MetaLlama318BInstructV1_128k => "meta-llama3-1-8b-instruct-v1-128k",
173            Model::MetaLlama318BInstructV1 => "meta-llama3-1-8b-instruct-v1",
174            Model::MetaLlama3170BInstructV1_128k => "meta-llama3-1-70b-instruct-v1-128k",
175            Model::MetaLlama3170BInstructV1 => "meta-llama3-1-70b-instruct-v1",
176            Model::MetaLlama31405BInstructV1 => "meta-llama3-1-405b-instruct-v1",
177            Model::MetaLlama321BInstructV1 => "meta-llama3-2-1b-instruct-v1",
178            Model::MetaLlama323BInstructV1 => "meta-llama3-2-3b-instruct-v1",
179            Model::MetaLlama3211BInstructV1 => "meta-llama3-2-11b-instruct-v1",
180            Model::MetaLlama3290BInstructV1 => "meta-llama3-2-90b-instruct-v1",
181            Model::MetaLlama3370BInstructV1 => "meta-llama3-3-70b-instruct-v1",
182            Model::MetaLlama4Scout17BInstructV1 => "meta-llama4-scout-17b-instruct-v1",
183            Model::MetaLlama4Maverick17BInstructV1 => "meta-llama4-maverick-17b-instruct-v1",
184            Model::MistralMistral7BInstructV0 => "mistral-7b-instruct-v0",
185            Model::MistralMixtral8x7BInstructV0 => "mistral-mixtral-8x7b-instruct-v0",
186            Model::MistralMistralLarge2402V1 => "mistral-large-2402-v1",
187            Model::MistralMistralSmall2402V1 => "mistral-small-2402-v1",
188            Model::MistralPixtralLarge2502V1 => "mistral-pixtral-large-2502-v1",
189            Model::PalmyraWriterX4 => "palmyra-writer-x4",
190            Model::PalmyraWriterX5 => "palmyra-writer-x5",
191            Self::Custom { name, .. } => name,
192        }
193    }
194
195    pub fn request_id(&self) -> &str {
196        match self {
197            Model::ClaudeSonnet4 | Model::ClaudeSonnet4Thinking => {
198                "anthropic.claude-sonnet-4-20250514-v1:0"
199            }
200            Model::ClaudeOpus4 | Model::ClaudeOpus4Thinking => {
201                "anthropic.claude-opus-4-20250514-v1:0"
202            }
203            Model::Claude3_5SonnetV2 => "anthropic.claude-3-5-sonnet-20241022-v2:0",
204            Model::Claude3_5Sonnet => "anthropic.claude-3-5-sonnet-20240620-v1:0",
205            Model::Claude3Opus => "anthropic.claude-3-opus-20240229-v1:0",
206            Model::Claude3Sonnet => "anthropic.claude-3-sonnet-20240229-v1:0",
207            Model::Claude3Haiku => "anthropic.claude-3-haiku-20240307-v1:0",
208            Model::Claude3_5Haiku => "anthropic.claude-3-5-haiku-20241022-v1:0",
209            Model::Claude3_7Sonnet | Model::Claude3_7SonnetThinking => {
210                "anthropic.claude-3-7-sonnet-20250219-v1:0"
211            }
212            Model::AmazonNovaLite => "amazon.nova-lite-v1:0",
213            Model::AmazonNovaMicro => "amazon.nova-micro-v1:0",
214            Model::AmazonNovaPro => "amazon.nova-pro-v1:0",
215            Model::AmazonNovaPremier => "amazon.nova-premier-v1:0",
216            Model::DeepSeekR1 => "deepseek.r1-v1:0",
217            Model::AI21J2GrandeInstruct => "ai21.j2-grande-instruct",
218            Model::AI21J2JumboInstruct => "ai21.j2-jumbo-instruct",
219            Model::AI21J2Mid => "ai21.j2-mid",
220            Model::AI21J2MidV1 => "ai21.j2-mid-v1",
221            Model::AI21J2Ultra => "ai21.j2-ultra",
222            Model::AI21J2UltraV1_8k => "ai21.j2-ultra-v1:0:8k",
223            Model::AI21J2UltraV1 => "ai21.j2-ultra-v1",
224            Model::AI21JambaInstructV1 => "ai21.jamba-instruct-v1:0",
225            Model::AI21Jamba15LargeV1 => "ai21.jamba-1-5-large-v1:0",
226            Model::AI21Jamba15MiniV1 => "ai21.jamba-1-5-mini-v1:0",
227            Model::CohereCommandTextV14_4k => "cohere.command-text-v14:7:4k",
228            Model::CohereCommandRV1 => "cohere.command-r-v1:0",
229            Model::CohereCommandRPlusV1 => "cohere.command-r-plus-v1:0",
230            Model::CohereCommandLightTextV14_4k => "cohere.command-light-text-v14:7:4k",
231            Model::MetaLlama38BInstructV1 => "meta.llama3-8b-instruct-v1:0",
232            Model::MetaLlama370BInstructV1 => "meta.llama3-70b-instruct-v1:0",
233            Model::MetaLlama318BInstructV1_128k => "meta.llama3-1-8b-instruct-v1:0",
234            Model::MetaLlama318BInstructV1 => "meta.llama3-1-8b-instruct-v1:0",
235            Model::MetaLlama3170BInstructV1_128k => "meta.llama3-1-70b-instruct-v1:0",
236            Model::MetaLlama3170BInstructV1 => "meta.llama3-1-70b-instruct-v1:0",
237            Model::MetaLlama31405BInstructV1 => "meta.llama3-1-405b-instruct-v1:0",
238            Model::MetaLlama3211BInstructV1 => "meta.llama3-2-11b-instruct-v1:0",
239            Model::MetaLlama3290BInstructV1 => "meta.llama3-2-90b-instruct-v1:0",
240            Model::MetaLlama321BInstructV1 => "meta.llama3-2-1b-instruct-v1:0",
241            Model::MetaLlama323BInstructV1 => "meta.llama3-2-3b-instruct-v1:0",
242            Model::MetaLlama3370BInstructV1 => "meta.llama3-3-70b-instruct-v1:0",
243            Model::MetaLlama4Scout17BInstructV1 => "meta.llama4-scout-17b-instruct-v1:0",
244            Model::MetaLlama4Maverick17BInstructV1 => "meta.llama4-maverick-17b-instruct-v1:0",
245            Model::MistralMistral7BInstructV0 => "mistral.mistral-7b-instruct-v0:2",
246            Model::MistralMixtral8x7BInstructV0 => "mistral.mixtral-8x7b-instruct-v0:1",
247            Model::MistralMistralLarge2402V1 => "mistral.mistral-large-2402-v1:0",
248            Model::MistralMistralSmall2402V1 => "mistral.mistral-small-2402-v1:0",
249            Model::MistralPixtralLarge2502V1 => "mistral.pixtral-large-2502-v1:0",
250            Model::PalmyraWriterX4 => "writer.palmyra-x4-v1:0",
251            Model::PalmyraWriterX5 => "writer.palmyra-x5-v1:0",
252            Self::Custom { name, .. } => name,
253        }
254    }
255
256    pub fn display_name(&self) -> &str {
257        match self {
258            Self::ClaudeSonnet4 => "Claude Sonnet 4",
259            Self::ClaudeSonnet4Thinking => "Claude Sonnet 4 Thinking",
260            Self::ClaudeOpus4 => "Claude Opus 4",
261            Self::ClaudeOpus4Thinking => "Claude Opus 4 Thinking",
262            Self::Claude3_5SonnetV2 => "Claude 3.5 Sonnet v2",
263            Self::Claude3_5Sonnet => "Claude 3.5 Sonnet",
264            Self::Claude3Opus => "Claude 3 Opus",
265            Self::Claude3Sonnet => "Claude 3 Sonnet",
266            Self::Claude3Haiku => "Claude 3 Haiku",
267            Self::Claude3_5Haiku => "Claude 3.5 Haiku",
268            Self::Claude3_7Sonnet => "Claude 3.7 Sonnet",
269            Self::Claude3_7SonnetThinking => "Claude 3.7 Sonnet Thinking",
270            Self::AmazonNovaLite => "Amazon Nova Lite",
271            Self::AmazonNovaMicro => "Amazon Nova Micro",
272            Self::AmazonNovaPro => "Amazon Nova Pro",
273            Self::AmazonNovaPremier => "Amazon Nova Premier",
274            Self::DeepSeekR1 => "DeepSeek R1",
275            Self::AI21J2GrandeInstruct => "AI21 Jurassic2 Grande Instruct",
276            Self::AI21J2JumboInstruct => "AI21 Jurassic2 Jumbo Instruct",
277            Self::AI21J2Mid => "AI21 Jurassic2 Mid",
278            Self::AI21J2MidV1 => "AI21 Jurassic2 Mid V1",
279            Self::AI21J2Ultra => "AI21 Jurassic2 Ultra",
280            Self::AI21J2UltraV1_8k => "AI21 Jurassic2 Ultra V1 8K",
281            Self::AI21J2UltraV1 => "AI21 Jurassic2 Ultra V1",
282            Self::AI21JambaInstructV1 => "AI21 Jamba Instruct",
283            Self::AI21Jamba15LargeV1 => "AI21 Jamba 1.5 Large",
284            Self::AI21Jamba15MiniV1 => "AI21 Jamba 1.5 Mini",
285            Self::CohereCommandTextV14_4k => "Cohere Command Text V14 4K",
286            Self::CohereCommandRV1 => "Cohere Command R V1",
287            Self::CohereCommandRPlusV1 => "Cohere Command R Plus V1",
288            Self::CohereCommandLightTextV14_4k => "Cohere Command Light Text V14 4K",
289            Self::MetaLlama38BInstructV1 => "Meta Llama 3 8B Instruct",
290            Self::MetaLlama370BInstructV1 => "Meta Llama 3 70B Instruct",
291            Self::MetaLlama318BInstructV1_128k => "Meta Llama 3.1 8B Instruct 128K",
292            Self::MetaLlama318BInstructV1 => "Meta Llama 3.1 8B Instruct",
293            Self::MetaLlama3170BInstructV1_128k => "Meta Llama 3.1 70B Instruct 128K",
294            Self::MetaLlama3170BInstructV1 => "Meta Llama 3.1 70B Instruct",
295            Self::MetaLlama31405BInstructV1 => "Meta Llama 3.1 405B Instruct",
296            Self::MetaLlama3211BInstructV1 => "Meta Llama 3.2 11B Instruct",
297            Self::MetaLlama3290BInstructV1 => "Meta Llama 3.2 90B Instruct",
298            Self::MetaLlama321BInstructV1 => "Meta Llama 3.2 1B Instruct",
299            Self::MetaLlama323BInstructV1 => "Meta Llama 3.2 3B Instruct",
300            Self::MetaLlama3370BInstructV1 => "Meta Llama 3.3 70B Instruct",
301            Self::MetaLlama4Scout17BInstructV1 => "Meta Llama 4 Scout 17B Instruct",
302            Self::MetaLlama4Maverick17BInstructV1 => "Meta Llama 4 Maverick 17B Instruct",
303            Self::MistralMistral7BInstructV0 => "Mistral 7B Instruct V0",
304            Self::MistralMixtral8x7BInstructV0 => "Mistral Mixtral 8x7B Instruct V0",
305            Self::MistralMistralLarge2402V1 => "Mistral Large 2402 V1",
306            Self::MistralMistralSmall2402V1 => "Mistral Small 2402 V1",
307            Self::MistralPixtralLarge2502V1 => "Pixtral Large 25.02 V1",
308            Self::PalmyraWriterX5 => "Writer Palmyra X5",
309            Self::PalmyraWriterX4 => "Writer Palmyra X4",
310            Self::Custom {
311                display_name, name, ..
312            } => display_name.as_deref().unwrap_or(name),
313        }
314    }
315
316    pub fn max_token_count(&self) -> u64 {
317        match self {
318            Self::Claude3_5SonnetV2
319            | Self::Claude3Opus
320            | Self::Claude3Sonnet
321            | Self::Claude3_5Haiku
322            | Self::Claude3_7Sonnet
323            | Self::ClaudeSonnet4
324            | Self::ClaudeOpus4
325            | Self::ClaudeSonnet4Thinking
326            | Self::ClaudeOpus4Thinking => 200_000,
327            Self::AmazonNovaPremier => 1_000_000,
328            Self::PalmyraWriterX5 => 1_000_000,
329            Self::PalmyraWriterX4 => 128_000,
330            Self::Custom { max_tokens, .. } => *max_tokens,
331            _ => 128_000,
332        }
333    }
334
335    pub fn max_output_tokens(&self) -> u64 {
336        match self {
337            Self::Claude3Opus | Self::Claude3Sonnet | Self::Claude3_5Haiku => 4_096,
338            Self::Claude3_7Sonnet
339            | Self::Claude3_7SonnetThinking
340            | Self::ClaudeSonnet4
341            | Self::ClaudeSonnet4Thinking
342            | Self::ClaudeOpus4
343            | Model::ClaudeOpus4Thinking => 128_000,
344            Self::Claude3_5SonnetV2 | Self::PalmyraWriterX4 | Self::PalmyraWriterX5 => 8_192,
345            Self::Custom {
346                max_output_tokens, ..
347            } => max_output_tokens.unwrap_or(4_096),
348            _ => 4_096,
349        }
350    }
351
352    pub fn default_temperature(&self) -> f32 {
353        match self {
354            Self::Claude3_5SonnetV2
355            | Self::Claude3Opus
356            | Self::Claude3Sonnet
357            | Self::Claude3_5Haiku
358            | Self::Claude3_7Sonnet
359            | Self::ClaudeOpus4
360            | Self::ClaudeOpus4Thinking
361            | Self::ClaudeSonnet4
362            | Self::ClaudeSonnet4Thinking => 1.0,
363            Self::Custom {
364                default_temperature,
365                ..
366            } => default_temperature.unwrap_or(1.0),
367            _ => 1.0,
368        }
369    }
370
371    pub fn supports_tool_use(&self) -> bool {
372        match self {
373            // Anthropic Claude 3 models (all support tool use)
374            Self::Claude3Opus
375            | Self::Claude3Sonnet
376            | Self::Claude3_5Sonnet
377            | Self::Claude3_5SonnetV2
378            | Self::Claude3_7Sonnet
379            | Self::Claude3_7SonnetThinking
380            | Self::ClaudeOpus4
381            | Self::ClaudeOpus4Thinking
382            | Self::ClaudeSonnet4
383            | Self::ClaudeSonnet4Thinking
384            | Self::Claude3_5Haiku => true,
385
386            // Amazon Nova models (all support tool use)
387            Self::AmazonNovaPremier
388            | Self::AmazonNovaPro
389            | Self::AmazonNovaLite
390            | Self::AmazonNovaMicro => true,
391
392            // AI21 Jamba 1.5 models support tool use
393            Self::AI21Jamba15LargeV1 | Self::AI21Jamba15MiniV1 => true,
394
395            // Cohere Command R models support tool use
396            Self::CohereCommandRV1 | Self::CohereCommandRPlusV1 => true,
397
398            // All other models don't support tool use
399            // Including Meta Llama 3.2, AI21 Jurassic, and others
400            _ => false,
401        }
402    }
403
404    pub fn mode(&self) -> BedrockModelMode {
405        match self {
406            Model::Claude3_7SonnetThinking => BedrockModelMode::Thinking {
407                budget_tokens: Some(4096),
408            },
409            Model::ClaudeSonnet4Thinking => BedrockModelMode::Thinking {
410                budget_tokens: Some(4096),
411            },
412            Model::ClaudeOpus4Thinking => BedrockModelMode::Thinking {
413                budget_tokens: Some(4096),
414            },
415            _ => BedrockModelMode::Default,
416        }
417    }
418
419    pub fn cross_region_inference_id(&self, region: &str) -> anyhow::Result<String> {
420        let region_group = if region.starts_with("us-gov-") {
421            "us-gov"
422        } else if region.starts_with("us-") {
423            "us"
424        } else if region.starts_with("eu-") {
425            "eu"
426        } else if region.starts_with("ap-") || region == "me-central-1" || region == "me-south-1" {
427            "apac"
428        } else if region.starts_with("ca-") || region.starts_with("sa-") {
429            // Canada and South America regions - default to US profiles
430            "us"
431        } else {
432            anyhow::bail!("Unsupported Region {region}");
433        };
434
435        let model_id = self.request_id();
436
437        match (self, region_group) {
438            // Custom models can't have CRI IDs
439            (Model::Custom { .. }, _) => Ok(self.request_id().into()),
440
441            // Models with US Gov only
442            (Model::Claude3_5Sonnet, "us-gov") | (Model::Claude3Haiku, "us-gov") => {
443                Ok(format!("{}.{}", region_group, model_id))
444            }
445
446            // Available everywhere
447            (Model::AmazonNovaLite | Model::AmazonNovaMicro | Model::AmazonNovaPro, _) => {
448                Ok(format!("{}.{}", region_group, model_id))
449            }
450
451            // Models in US
452            (
453                Model::AmazonNovaPremier
454                | Model::Claude3_5Haiku
455                | Model::Claude3_5Sonnet
456                | Model::Claude3_5SonnetV2
457                | Model::Claude3_7Sonnet
458                | Model::Claude3_7SonnetThinking
459                | Model::ClaudeSonnet4
460                | Model::ClaudeSonnet4Thinking
461                | Model::ClaudeOpus4
462                | Model::ClaudeOpus4Thinking
463                | Model::Claude3Haiku
464                | Model::Claude3Opus
465                | Model::Claude3Sonnet
466                | Model::DeepSeekR1
467                | Model::MetaLlama31405BInstructV1
468                | Model::MetaLlama3170BInstructV1_128k
469                | Model::MetaLlama3170BInstructV1
470                | Model::MetaLlama318BInstructV1_128k
471                | Model::MetaLlama318BInstructV1
472                | Model::MetaLlama3211BInstructV1
473                | Model::MetaLlama321BInstructV1
474                | Model::MetaLlama323BInstructV1
475                | Model::MetaLlama3290BInstructV1
476                | Model::MetaLlama3370BInstructV1
477                | Model::MetaLlama4Maverick17BInstructV1
478                | Model::MetaLlama4Scout17BInstructV1
479                | Model::MistralPixtralLarge2502V1
480                | Model::PalmyraWriterX4
481                | Model::PalmyraWriterX5,
482                "us",
483            ) => Ok(format!("{}.{}", region_group, model_id)),
484
485            // Models available in EU
486            (
487                Model::Claude3_5Sonnet
488                | Model::Claude3_7Sonnet
489                | Model::Claude3_7SonnetThinking
490                | Model::ClaudeSonnet4
491                | Model::ClaudeSonnet4Thinking
492                | Model::Claude3Haiku
493                | Model::Claude3Sonnet
494                | Model::MetaLlama321BInstructV1
495                | Model::MetaLlama323BInstructV1
496                | Model::MistralPixtralLarge2502V1,
497                "eu",
498            ) => Ok(format!("{}.{}", region_group, model_id)),
499
500            // Models available in APAC
501            (
502                Model::Claude3_5Sonnet
503                | Model::Claude3_5SonnetV2
504                | Model::Claude3Haiku
505                | Model::Claude3Sonnet
506                | Model::Claude3_7Sonnet
507                | Model::Claude3_7SonnetThinking
508                | Model::ClaudeSonnet4
509                | Model::ClaudeSonnet4Thinking,
510                "apac",
511            ) => Ok(format!("{}.{}", region_group, model_id)),
512
513            // Any other combination is not supported
514            _ => Ok(self.request_id().into()),
515        }
516    }
517}
518
519#[cfg(test)]
520mod tests {
521    use super::*;
522
523    #[test]
524    fn test_us_region_inference_ids() -> anyhow::Result<()> {
525        // Test US regions
526        assert_eq!(
527            Model::Claude3_5SonnetV2.cross_region_inference_id("us-east-1")?,
528            "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
529        );
530        assert_eq!(
531            Model::Claude3_5SonnetV2.cross_region_inference_id("us-west-2")?,
532            "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
533        );
534        assert_eq!(
535            Model::AmazonNovaPro.cross_region_inference_id("us-east-2")?,
536            "us.amazon.nova-pro-v1:0"
537        );
538        Ok(())
539    }
540
541    #[test]
542    fn test_eu_region_inference_ids() -> anyhow::Result<()> {
543        // Test European regions
544        assert_eq!(
545            Model::ClaudeSonnet4.cross_region_inference_id("eu-west-1")?,
546            "eu.anthropic.claude-sonnet-4-20250514-v1:0"
547        );
548        assert_eq!(
549            Model::Claude3Sonnet.cross_region_inference_id("eu-west-1")?,
550            "eu.anthropic.claude-3-sonnet-20240229-v1:0"
551        );
552        assert_eq!(
553            Model::AmazonNovaMicro.cross_region_inference_id("eu-north-1")?,
554            "eu.amazon.nova-micro-v1:0"
555        );
556        Ok(())
557    }
558
559    #[test]
560    fn test_apac_region_inference_ids() -> anyhow::Result<()> {
561        // Test Asia-Pacific regions
562        assert_eq!(
563            Model::Claude3_5SonnetV2.cross_region_inference_id("ap-northeast-1")?,
564            "apac.anthropic.claude-3-5-sonnet-20241022-v2:0"
565        );
566        assert_eq!(
567            Model::Claude3_5SonnetV2.cross_region_inference_id("ap-southeast-2")?,
568            "apac.anthropic.claude-3-5-sonnet-20241022-v2:0"
569        );
570        assert_eq!(
571            Model::AmazonNovaLite.cross_region_inference_id("ap-south-1")?,
572            "apac.amazon.nova-lite-v1:0"
573        );
574        Ok(())
575    }
576
577    #[test]
578    fn test_gov_region_inference_ids() -> anyhow::Result<()> {
579        // Test Government regions
580        assert_eq!(
581            Model::Claude3_5Sonnet.cross_region_inference_id("us-gov-east-1")?,
582            "us-gov.anthropic.claude-3-5-sonnet-20240620-v1:0"
583        );
584        assert_eq!(
585            Model::Claude3Haiku.cross_region_inference_id("us-gov-west-1")?,
586            "us-gov.anthropic.claude-3-haiku-20240307-v1:0"
587        );
588        Ok(())
589    }
590
591    #[test]
592    fn test_meta_models_inference_ids() -> anyhow::Result<()> {
593        // Test Meta models
594        assert_eq!(
595            Model::MetaLlama370BInstructV1.cross_region_inference_id("us-east-1")?,
596            "meta.llama3-70b-instruct-v1:0"
597        );
598        assert_eq!(
599            Model::MetaLlama3170BInstructV1.cross_region_inference_id("us-east-1")?,
600            "us.meta.llama3-1-70b-instruct-v1:0"
601        );
602        assert_eq!(
603            Model::MetaLlama321BInstructV1.cross_region_inference_id("eu-west-1")?,
604            "eu.meta.llama3-2-1b-instruct-v1:0"
605        );
606        Ok(())
607    }
608
609    #[test]
610    fn test_mistral_models_inference_ids() -> anyhow::Result<()> {
611        // Mistral models don't follow the regional prefix pattern,
612        // so they should return their original IDs
613        assert_eq!(
614            Model::MistralMistralLarge2402V1.cross_region_inference_id("us-east-1")?,
615            "mistral.mistral-large-2402-v1:0"
616        );
617        assert_eq!(
618            Model::MistralMixtral8x7BInstructV0.cross_region_inference_id("eu-west-1")?,
619            "mistral.mixtral-8x7b-instruct-v0:1"
620        );
621        Ok(())
622    }
623
624    #[test]
625    fn test_ai21_models_inference_ids() -> anyhow::Result<()> {
626        // AI21 models don't follow the regional prefix pattern,
627        // so they should return their original IDs
628        assert_eq!(
629            Model::AI21J2UltraV1.cross_region_inference_id("us-east-1")?,
630            "ai21.j2-ultra-v1"
631        );
632        assert_eq!(
633            Model::AI21JambaInstructV1.cross_region_inference_id("eu-west-1")?,
634            "ai21.jamba-instruct-v1:0"
635        );
636        Ok(())
637    }
638
639    #[test]
640    fn test_cohere_models_inference_ids() -> anyhow::Result<()> {
641        // Cohere models don't follow the regional prefix pattern,
642        // so they should return their original IDs
643        assert_eq!(
644            Model::CohereCommandRV1.cross_region_inference_id("us-east-1")?,
645            "cohere.command-r-v1:0"
646        );
647        assert_eq!(
648            Model::CohereCommandTextV14_4k.cross_region_inference_id("ap-southeast-1")?,
649            "cohere.command-text-v14:7:4k"
650        );
651        Ok(())
652    }
653
654    #[test]
655    fn test_custom_model_inference_ids() -> anyhow::Result<()> {
656        // Test custom models
657        let custom_model = Model::Custom {
658            name: "custom.my-model-v1:0".to_string(),
659            max_tokens: 100000,
660            display_name: Some("My Custom Model".to_string()),
661            max_output_tokens: Some(8192),
662            default_temperature: Some(0.7),
663        };
664
665        // Custom model should return its name unchanged
666        assert_eq!(
667            custom_model.cross_region_inference_id("us-east-1")?,
668            "custom.my-model-v1:0"
669        );
670
671        Ok(())
672    }
673
674    #[test]
675    fn test_friendly_id_vs_request_id() {
676        // Test that id() returns friendly identifiers
677        assert_eq!(Model::Claude3_5SonnetV2.id(), "claude-3-5-sonnet-v2");
678        assert_eq!(Model::AmazonNovaLite.id(), "amazon-nova-lite");
679        assert_eq!(Model::DeepSeekR1.id(), "deepseek-r1");
680        assert_eq!(
681            Model::MetaLlama38BInstructV1.id(),
682            "meta-llama3-8b-instruct-v1"
683        );
684
685        // Test that request_id() returns actual backend model IDs
686        assert_eq!(
687            Model::Claude3_5SonnetV2.request_id(),
688            "anthropic.claude-3-5-sonnet-20241022-v2:0"
689        );
690        assert_eq!(Model::AmazonNovaLite.request_id(), "amazon.nova-lite-v1:0");
691        assert_eq!(Model::DeepSeekR1.request_id(), "deepseek.r1-v1:0");
692        assert_eq!(
693            Model::MetaLlama38BInstructV1.request_id(),
694            "meta.llama3-8b-instruct-v1:0"
695        );
696
697        // Test thinking models have different friendly IDs but same request IDs
698        assert_eq!(Model::ClaudeSonnet4.id(), "claude-4-sonnet");
699        assert_eq!(
700            Model::ClaudeSonnet4Thinking.id(),
701            "claude-4-sonnet-thinking"
702        );
703        assert_eq!(
704            Model::ClaudeSonnet4.request_id(),
705            Model::ClaudeSonnet4Thinking.request_id()
706        );
707    }
708}