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