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