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