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