1use anyhow::anyhow;
2use serde::{Deserialize, Serialize};
3use strum::EnumIter;
4
5#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
6#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
7pub enum BedrockModelMode {
8 #[default]
9 Default,
10 Thinking {
11 budget_tokens: Option<u64>,
12 },
13}
14
15#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
16#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)]
17pub enum Model {
18 // Anthropic models (already included)
19 #[default]
20 #[serde(rename = "claude-3-5-sonnet-v2", alias = "claude-3-5-sonnet-latest")]
21 Claude3_5SonnetV2,
22 #[serde(rename = "claude-3-7-sonnet", alias = "claude-3-7-sonnet-latest")]
23 Claude3_7Sonnet,
24 #[serde(
25 rename = "claude-3-7-sonnet-thinking",
26 alias = "claude-3-7-sonnet-thinking-latest"
27 )]
28 Claude3_7SonnetThinking,
29 #[serde(rename = "claude-3-opus", alias = "claude-3-opus-latest")]
30 Claude3Opus,
31 #[serde(rename = "claude-3-sonnet", alias = "claude-3-sonnet-latest")]
32 Claude3Sonnet,
33 #[serde(rename = "claude-3-5-haiku", alias = "claude-3-5-haiku-latest")]
34 Claude3_5Haiku,
35 Claude3_5Sonnet,
36 Claude3Haiku,
37 // Amazon Nova Models
38 AmazonNovaLite,
39 AmazonNovaMicro,
40 AmazonNovaPro,
41 // AI21 models
42 AI21J2GrandeInstruct,
43 AI21J2JumboInstruct,
44 AI21J2Mid,
45 AI21J2MidV1,
46 AI21J2Ultra,
47 AI21J2UltraV1_8k,
48 AI21J2UltraV1,
49 AI21JambaInstructV1,
50 AI21Jamba15LargeV1,
51 AI21Jamba15MiniV1,
52 // Cohere models
53 CohereCommandTextV14_4k,
54 CohereCommandRV1,
55 CohereCommandRPlusV1,
56 CohereCommandLightTextV14_4k,
57 // DeepSeek
58 DeepSeekR1,
59 // Meta models
60 MetaLlama38BInstructV1,
61 MetaLlama370BInstructV1,
62 MetaLlama318BInstructV1_128k,
63 MetaLlama318BInstructV1,
64 MetaLlama3170BInstructV1_128k,
65 MetaLlama3170BInstructV1,
66 MetaLlama3211BInstructV1,
67 MetaLlama3290BInstructV1,
68 MetaLlama321BInstructV1,
69 MetaLlama323BInstructV1,
70 // Mistral models
71 MistralMistral7BInstructV0,
72 MistralMixtral8x7BInstructV0,
73 MistralMistralLarge2402V1,
74 MistralMistralSmall2402V1,
75 #[serde(rename = "custom")]
76 Custom {
77 name: String,
78 max_tokens: usize,
79 /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
80 display_name: Option<String>,
81 max_output_tokens: Option<u32>,
82 default_temperature: Option<f32>,
83 },
84}
85
86impl Model {
87 pub fn default_fast() -> Self {
88 Self::Claude3_5Haiku
89 }
90
91 pub fn from_id(id: &str) -> anyhow::Result<Self> {
92 if id.starts_with("claude-3-5-sonnet-v2") {
93 Ok(Self::Claude3_5SonnetV2)
94 } else if id.starts_with("claude-3-opus") {
95 Ok(Self::Claude3Opus)
96 } else if id.starts_with("claude-3-sonnet") {
97 Ok(Self::Claude3Sonnet)
98 } else if id.starts_with("claude-3-5-haiku") {
99 Ok(Self::Claude3_5Haiku)
100 } else if id.starts_with("claude-3-7-sonnet") {
101 Ok(Self::Claude3_7Sonnet)
102 } else if id.starts_with("claude-3-7-sonnet-thinking") {
103 Ok(Self::Claude3_7SonnetThinking)
104 } else {
105 Err(anyhow!("invalid model id"))
106 }
107 }
108
109 pub fn id(&self) -> &str {
110 match self {
111 Model::Claude3_5SonnetV2 => "anthropic.claude-3-5-sonnet-20241022-v2:0",
112 Model::Claude3_5Sonnet => "anthropic.claude-3-5-sonnet-20240620-v1:0",
113 Model::Claude3Opus => "anthropic.claude-3-opus-20240229-v1:0",
114 Model::Claude3Sonnet => "anthropic.claude-3-sonnet-20240229-v1:0",
115 Model::Claude3Haiku => "anthropic.claude-3-haiku-20240307-v1:0",
116 Model::Claude3_5Haiku => "anthropic.claude-3-5-haiku-20241022-v1:0",
117 Model::Claude3_7Sonnet | Model::Claude3_7SonnetThinking => {
118 "anthropic.claude-3-7-sonnet-20250219-v1:0"
119 }
120 Model::AmazonNovaLite => "amazon.nova-lite-v1:0",
121 Model::AmazonNovaMicro => "amazon.nova-micro-v1:0",
122 Model::AmazonNovaPro => "amazon.nova-pro-v1:0",
123 Model::DeepSeekR1 => "us.deepseek.r1-v1:0",
124 Model::AI21J2GrandeInstruct => "ai21.j2-grande-instruct",
125 Model::AI21J2JumboInstruct => "ai21.j2-jumbo-instruct",
126 Model::AI21J2Mid => "ai21.j2-mid",
127 Model::AI21J2MidV1 => "ai21.j2-mid-v1",
128 Model::AI21J2Ultra => "ai21.j2-ultra",
129 Model::AI21J2UltraV1_8k => "ai21.j2-ultra-v1:0:8k",
130 Model::AI21J2UltraV1 => "ai21.j2-ultra-v1",
131 Model::AI21JambaInstructV1 => "ai21.jamba-instruct-v1:0",
132 Model::AI21Jamba15LargeV1 => "ai21.jamba-1-5-large-v1:0",
133 Model::AI21Jamba15MiniV1 => "ai21.jamba-1-5-mini-v1:0",
134 Model::CohereCommandTextV14_4k => "cohere.command-text-v14:7:4k",
135 Model::CohereCommandRV1 => "cohere.command-r-v1:0",
136 Model::CohereCommandRPlusV1 => "cohere.command-r-plus-v1:0",
137 Model::CohereCommandLightTextV14_4k => "cohere.command-light-text-v14:7:4k",
138 Model::MetaLlama38BInstructV1 => "meta.llama3-8b-instruct-v1:0",
139 Model::MetaLlama370BInstructV1 => "meta.llama3-70b-instruct-v1:0",
140 Model::MetaLlama318BInstructV1_128k => "meta.llama3-1-8b-instruct-v1:0:128k",
141 Model::MetaLlama318BInstructV1 => "meta.llama3-1-8b-instruct-v1:0",
142 Model::MetaLlama3170BInstructV1_128k => "meta.llama3-1-70b-instruct-v1:0:128k",
143 Model::MetaLlama3170BInstructV1 => "meta.llama3-1-70b-instruct-v1:0",
144 Model::MetaLlama3211BInstructV1 => "meta.llama3-2-11b-instruct-v1:0",
145 Model::MetaLlama3290BInstructV1 => "meta.llama3-2-90b-instruct-v1:0",
146 Model::MetaLlama321BInstructV1 => "meta.llama3-2-1b-instruct-v1:0",
147 Model::MetaLlama323BInstructV1 => "meta.llama3-2-3b-instruct-v1:0",
148 Model::MistralMistral7BInstructV0 => "mistral.mistral-7b-instruct-v0:2",
149 Model::MistralMixtral8x7BInstructV0 => "mistral.mixtral-8x7b-instruct-v0:1",
150 Model::MistralMistralLarge2402V1 => "mistral.mistral-large-2402-v1:0",
151 Model::MistralMistralSmall2402V1 => "mistral.mistral-small-2402-v1:0",
152 Self::Custom { name, .. } => name,
153 }
154 }
155
156 pub fn display_name(&self) -> &str {
157 match self {
158 Self::Claude3_5SonnetV2 => "Claude 3.5 Sonnet v2",
159 Self::Claude3_5Sonnet => "Claude 3.5 Sonnet",
160 Self::Claude3Opus => "Claude 3 Opus",
161 Self::Claude3Sonnet => "Claude 3 Sonnet",
162 Self::Claude3Haiku => "Claude 3 Haiku",
163 Self::Claude3_5Haiku => "Claude 3.5 Haiku",
164 Self::Claude3_7Sonnet => "Claude 3.7 Sonnet",
165 Self::Claude3_7SonnetThinking => "Claude 3.7 Sonnet Thinking",
166 Self::AmazonNovaLite => "Amazon Nova Lite",
167 Self::AmazonNovaMicro => "Amazon Nova Micro",
168 Self::AmazonNovaPro => "Amazon Nova Pro",
169 Self::DeepSeekR1 => "DeepSeek R1",
170 Self::AI21J2GrandeInstruct => "AI21 Jurassic2 Grande Instruct",
171 Self::AI21J2JumboInstruct => "AI21 Jurassic2 Jumbo Instruct",
172 Self::AI21J2Mid => "AI21 Jurassic2 Mid",
173 Self::AI21J2MidV1 => "AI21 Jurassic2 Mid V1",
174 Self::AI21J2Ultra => "AI21 Jurassic2 Ultra",
175 Self::AI21J2UltraV1_8k => "AI21 Jurassic2 Ultra V1 8K",
176 Self::AI21J2UltraV1 => "AI21 Jurassic2 Ultra V1",
177 Self::AI21JambaInstructV1 => "AI21 Jamba Instruct",
178 Self::AI21Jamba15LargeV1 => "AI21 Jamba 1.5 Large",
179 Self::AI21Jamba15MiniV1 => "AI21 Jamba 1.5 Mini",
180 Self::CohereCommandTextV14_4k => "Cohere Command Text V14 4K",
181 Self::CohereCommandRV1 => "Cohere Command R V1",
182 Self::CohereCommandRPlusV1 => "Cohere Command R Plus V1",
183 Self::CohereCommandLightTextV14_4k => "Cohere Command Light Text V14 4K",
184 Self::MetaLlama38BInstructV1 => "Meta Llama 3 8B Instruct V1",
185 Self::MetaLlama370BInstructV1 => "Meta Llama 3 70B Instruct V1",
186 Self::MetaLlama318BInstructV1_128k => "Meta Llama 3 1.8B Instruct V1 128K",
187 Self::MetaLlama318BInstructV1 => "Meta Llama 3 1.8B Instruct V1",
188 Self::MetaLlama3170BInstructV1_128k => "Meta Llama 3 1 70B Instruct V1 128K",
189 Self::MetaLlama3170BInstructV1 => "Meta Llama 3 1 70B Instruct V1",
190 Self::MetaLlama3211BInstructV1 => "Meta Llama 3 2 11B Instruct V1",
191 Self::MetaLlama3290BInstructV1 => "Meta Llama 3 2 90B Instruct V1",
192 Self::MetaLlama321BInstructV1 => "Meta Llama 3 2 1B Instruct V1",
193 Self::MetaLlama323BInstructV1 => "Meta Llama 3 2 3B Instruct V1",
194 Self::MistralMistral7BInstructV0 => "Mistral 7B Instruct V0",
195 Self::MistralMixtral8x7BInstructV0 => "Mistral Mixtral 8x7B Instruct V0",
196 Self::MistralMistralLarge2402V1 => "Mistral Large 2402 V1",
197 Self::MistralMistralSmall2402V1 => "Mistral Small 2402 V1",
198 Self::Custom {
199 display_name, name, ..
200 } => display_name.as_deref().unwrap_or(name),
201 }
202 }
203
204 pub fn max_token_count(&self) -> usize {
205 match self {
206 Self::Claude3_5SonnetV2
207 | Self::Claude3Opus
208 | Self::Claude3Sonnet
209 | Self::Claude3_5Haiku
210 | Self::Claude3_7Sonnet => 200_000,
211 Self::Custom { max_tokens, .. } => *max_tokens,
212 _ => 200_000,
213 }
214 }
215
216 pub fn max_output_tokens(&self) -> u32 {
217 match self {
218 Self::Claude3Opus | Self::Claude3Sonnet | Self::Claude3_5Haiku => 4_096,
219 Self::Claude3_7Sonnet | Self::Claude3_7SonnetThinking => 128_000,
220 Self::Claude3_5SonnetV2 => 8_192,
221 Self::Custom {
222 max_output_tokens, ..
223 } => max_output_tokens.unwrap_or(4_096),
224 _ => 4_096,
225 }
226 }
227
228 pub fn default_temperature(&self) -> f32 {
229 match self {
230 Self::Claude3_5SonnetV2
231 | Self::Claude3Opus
232 | Self::Claude3Sonnet
233 | Self::Claude3_5Haiku
234 | Self::Claude3_7Sonnet => 1.0,
235 Self::Custom {
236 default_temperature,
237 ..
238 } => default_temperature.unwrap_or(1.0),
239 _ => 1.0,
240 }
241 }
242
243 pub fn supports_tool_use(&self) -> bool {
244 match self {
245 // Anthropic Claude 3 models (all support tool use)
246 Self::Claude3Opus
247 | Self::Claude3Sonnet
248 | Self::Claude3_5Sonnet
249 | Self::Claude3_5SonnetV2
250 | Self::Claude3_7Sonnet
251 | Self::Claude3_7SonnetThinking
252 | Self::Claude3_5Haiku => true,
253
254 // Amazon Nova models (all support tool use)
255 Self::AmazonNovaPro | Self::AmazonNovaLite | Self::AmazonNovaMicro => true,
256
257 // AI21 Jamba 1.5 models support tool use
258 Self::AI21Jamba15LargeV1 | Self::AI21Jamba15MiniV1 => true,
259
260 // Cohere Command R models support tool use
261 Self::CohereCommandRV1 | Self::CohereCommandRPlusV1 => true,
262
263 // All other models don't support tool use
264 // Including Meta Llama 3.2, AI21 Jurassic, and others
265 _ => false,
266 }
267 }
268
269 pub fn mode(&self) -> BedrockModelMode {
270 match self {
271 Model::Claude3_7SonnetThinking => BedrockModelMode::Thinking {
272 budget_tokens: Some(4096),
273 },
274 _ => BedrockModelMode::Default,
275 }
276 }
277
278 pub fn cross_region_inference_id(&self, region: &str) -> Result<String, anyhow::Error> {
279 let region_group = if region.starts_with("us-gov-") {
280 "us-gov"
281 } else if region.starts_with("us-") {
282 "us"
283 } else if region.starts_with("eu-") {
284 "eu"
285 } else if region.starts_with("ap-") || region == "me-central-1" || region == "me-south-1" {
286 "apac"
287 } else if region.starts_with("ca-") || region.starts_with("sa-") {
288 // Canada and South America regions - default to US profiles
289 "us"
290 } else {
291 // Unknown region
292 return Err(anyhow!("Unsupported Region"));
293 };
294
295 let model_id = self.id();
296
297 match (self, region_group) {
298 // Custom models can't have CRI IDs
299 (Model::Custom { .. }, _) => Ok(self.id().into()),
300
301 // Models with US Gov only
302 (Model::Claude3_5Sonnet, "us-gov") | (Model::Claude3Haiku, "us-gov") => {
303 Ok(format!("{}.{}", region_group, model_id))
304 }
305
306 // Models available only in US
307 (Model::Claude3Opus, "us")
308 | (Model::Claude3_7Sonnet, "us")
309 | (Model::Claude3_7SonnetThinking, "us") => {
310 Ok(format!("{}.{}", region_group, model_id))
311 }
312
313 // Models available in US, EU, and APAC
314 (Model::Claude3_5SonnetV2, "us")
315 | (Model::Claude3_5SonnetV2, "apac")
316 | (Model::Claude3_5Sonnet, _)
317 | (Model::Claude3Haiku, _)
318 | (Model::Claude3Sonnet, _)
319 | (Model::AmazonNovaLite, _)
320 | (Model::AmazonNovaMicro, _)
321 | (Model::AmazonNovaPro, _) => Ok(format!("{}.{}", region_group, model_id)),
322
323 // Models with limited EU availability
324 (Model::MetaLlama321BInstructV1, "us")
325 | (Model::MetaLlama321BInstructV1, "eu")
326 | (Model::MetaLlama323BInstructV1, "us")
327 | (Model::MetaLlama323BInstructV1, "eu") => {
328 Ok(format!("{}.{}", region_group, model_id))
329 }
330
331 // US-only models (all remaining Meta models)
332 (Model::MetaLlama38BInstructV1, "us")
333 | (Model::MetaLlama370BInstructV1, "us")
334 | (Model::MetaLlama318BInstructV1, "us")
335 | (Model::MetaLlama318BInstructV1_128k, "us")
336 | (Model::MetaLlama3170BInstructV1, "us")
337 | (Model::MetaLlama3170BInstructV1_128k, "us")
338 | (Model::MetaLlama3211BInstructV1, "us")
339 | (Model::MetaLlama3290BInstructV1, "us") => {
340 Ok(format!("{}.{}", region_group, model_id))
341 }
342
343 // Any other combination is not supported
344 _ => Ok(self.id().into()),
345 }
346 }
347}
348
349#[cfg(test)]
350mod tests {
351 use super::*;
352
353 #[test]
354 fn test_us_region_inference_ids() -> anyhow::Result<()> {
355 // Test US regions
356 assert_eq!(
357 Model::Claude3_5SonnetV2.cross_region_inference_id("us-east-1")?,
358 "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
359 );
360 assert_eq!(
361 Model::Claude3_5SonnetV2.cross_region_inference_id("us-west-2")?,
362 "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
363 );
364 assert_eq!(
365 Model::AmazonNovaPro.cross_region_inference_id("us-east-2")?,
366 "us.amazon.nova-pro-v1:0"
367 );
368 Ok(())
369 }
370
371 #[test]
372 fn test_eu_region_inference_ids() -> anyhow::Result<()> {
373 // Test European regions
374 assert_eq!(
375 Model::Claude3Sonnet.cross_region_inference_id("eu-west-1")?,
376 "eu.anthropic.claude-3-sonnet-20240229-v1:0"
377 );
378 assert_eq!(
379 Model::AmazonNovaMicro.cross_region_inference_id("eu-north-1")?,
380 "eu.amazon.nova-micro-v1:0"
381 );
382 Ok(())
383 }
384
385 #[test]
386 fn test_apac_region_inference_ids() -> anyhow::Result<()> {
387 // Test Asia-Pacific regions
388 assert_eq!(
389 Model::Claude3_5SonnetV2.cross_region_inference_id("ap-northeast-1")?,
390 "apac.anthropic.claude-3-5-sonnet-20241022-v2:0"
391 );
392 assert_eq!(
393 Model::AmazonNovaLite.cross_region_inference_id("ap-south-1")?,
394 "apac.amazon.nova-lite-v1:0"
395 );
396 Ok(())
397 }
398
399 #[test]
400 fn test_gov_region_inference_ids() -> anyhow::Result<()> {
401 // Test Government regions
402 assert_eq!(
403 Model::Claude3_5Sonnet.cross_region_inference_id("us-gov-east-1")?,
404 "us-gov.anthropic.claude-3-5-sonnet-20240620-v1:0"
405 );
406 assert_eq!(
407 Model::Claude3Haiku.cross_region_inference_id("us-gov-west-1")?,
408 "us-gov.anthropic.claude-3-haiku-20240307-v1:0"
409 );
410 Ok(())
411 }
412
413 #[test]
414 fn test_meta_models_inference_ids() -> anyhow::Result<()> {
415 // Test Meta models
416 assert_eq!(
417 Model::MetaLlama370BInstructV1.cross_region_inference_id("us-east-1")?,
418 "us.meta.llama3-70b-instruct-v1:0"
419 );
420 assert_eq!(
421 Model::MetaLlama321BInstructV1.cross_region_inference_id("eu-west-1")?,
422 "eu.meta.llama3-2-1b-instruct-v1:0"
423 );
424 Ok(())
425 }
426
427 #[test]
428 fn test_mistral_models_inference_ids() -> anyhow::Result<()> {
429 // Mistral models don't follow the regional prefix pattern,
430 // so they should return their original IDs
431 assert_eq!(
432 Model::MistralMistralLarge2402V1.cross_region_inference_id("us-east-1")?,
433 "mistral.mistral-large-2402-v1:0"
434 );
435 assert_eq!(
436 Model::MistralMixtral8x7BInstructV0.cross_region_inference_id("eu-west-1")?,
437 "mistral.mixtral-8x7b-instruct-v0:1"
438 );
439 Ok(())
440 }
441
442 #[test]
443 fn test_ai21_models_inference_ids() -> anyhow::Result<()> {
444 // AI21 models don't follow the regional prefix pattern,
445 // so they should return their original IDs
446 assert_eq!(
447 Model::AI21J2UltraV1.cross_region_inference_id("us-east-1")?,
448 "ai21.j2-ultra-v1"
449 );
450 assert_eq!(
451 Model::AI21JambaInstructV1.cross_region_inference_id("eu-west-1")?,
452 "ai21.jamba-instruct-v1:0"
453 );
454 Ok(())
455 }
456
457 #[test]
458 fn test_cohere_models_inference_ids() -> anyhow::Result<()> {
459 // Cohere models don't follow the regional prefix pattern,
460 // so they should return their original IDs
461 assert_eq!(
462 Model::CohereCommandRV1.cross_region_inference_id("us-east-1")?,
463 "cohere.command-r-v1:0"
464 );
465 assert_eq!(
466 Model::CohereCommandTextV14_4k.cross_region_inference_id("ap-southeast-1")?,
467 "cohere.command-text-v14:7:4k"
468 );
469 Ok(())
470 }
471
472 #[test]
473 fn test_custom_model_inference_ids() -> anyhow::Result<()> {
474 // Test custom models
475 let custom_model = Model::Custom {
476 name: "custom.my-model-v1:0".to_string(),
477 max_tokens: 100000,
478 display_name: Some("My Custom Model".to_string()),
479 max_output_tokens: Some(8192),
480 default_temperature: Some(0.7),
481 };
482
483 // Custom model should return its name unchanged
484 assert_eq!(
485 custom_model.cross_region_inference_id("us-east-1")?,
486 "custom.my-model-v1:0"
487 );
488
489 Ok(())
490 }
491}