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