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, EnumIter)]
7pub enum Model {
8 // Anthropic models (already included)
9 #[default]
10 #[serde(rename = "claude-3-5-sonnet-v2", alias = "claude-3-5-sonnet-latest")]
11 Claude3_5Sonnet,
12 #[serde(rename = "claude-3-7-sonnet", alias = "claude-3-7-sonnet-latest")]
13 Claude3_7Sonnet,
14 #[serde(rename = "claude-3-opus", alias = "claude-3-opus-latest")]
15 Claude3Opus,
16 #[serde(rename = "claude-3-sonnet", alias = "claude-3-sonnet-latest")]
17 Claude3Sonnet,
18 #[serde(rename = "claude-3-5-haiku", alias = "claude-3-5-haiku-latest")]
19 Claude3_5Haiku,
20 // Amazon Nova Models
21 AmazonNovaLite,
22 AmazonNovaMicro,
23 AmazonNovaPro,
24 // AI21 models
25 AI21J2GrandeInstruct,
26 AI21J2JumboInstruct,
27 AI21J2Mid,
28 AI21J2MidV1,
29 AI21J2Ultra,
30 AI21J2UltraV1_8k,
31 AI21J2UltraV1,
32 AI21JambaInstructV1,
33 AI21Jamba15LargeV1,
34 AI21Jamba15MiniV1,
35 // Cohere models
36 CohereCommandTextV14_4k,
37 CohereCommandRV1,
38 CohereCommandRPlusV1,
39 CohereCommandLightTextV14_4k,
40 // DeepSeek
41 DeepSeekR1,
42 // Meta models
43 MetaLlama38BInstructV1,
44 MetaLlama370BInstructV1,
45 MetaLlama318BInstructV1_128k,
46 MetaLlama318BInstructV1,
47 MetaLlama3170BInstructV1_128k,
48 MetaLlama3170BInstructV1,
49 MetaLlama3211BInstructV1,
50 MetaLlama3290BInstructV1,
51 MetaLlama321BInstructV1,
52 MetaLlama323BInstructV1,
53 // Mistral models
54 MistralMistral7BInstructV0,
55 MistralMixtral8x7BInstructV0,
56 MistralMistralLarge2402V1,
57 MistralMistralSmall2402V1,
58 #[serde(rename = "custom")]
59 Custom {
60 name: String,
61 max_tokens: usize,
62 /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
63 display_name: Option<String>,
64 max_output_tokens: Option<u32>,
65 default_temperature: Option<f32>,
66 },
67}
68
69impl Model {
70 pub fn from_id(id: &str) -> anyhow::Result<Self> {
71 if id.starts_with("claude-3-5-sonnet-v2") {
72 Ok(Self::Claude3_5Sonnet)
73 } else if id.starts_with("claude-3-opus") {
74 Ok(Self::Claude3Opus)
75 } else if id.starts_with("claude-3-sonnet") {
76 Ok(Self::Claude3Sonnet)
77 } else if id.starts_with("claude-3-5-haiku") {
78 Ok(Self::Claude3_5Haiku)
79 } else if id.starts_with("claude-3-7-sonnet") {
80 Ok(Self::Claude3_7Sonnet)
81 } else {
82 Err(anyhow!("invalid model id"))
83 }
84 }
85
86 pub fn id(&self) -> &str {
87 match self {
88 Model::Claude3_5Sonnet => "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
89 Model::Claude3Opus => "us.anthropic.claude-3-opus-20240229-v1:0",
90 Model::Claude3Sonnet => "us.anthropic.claude-3-sonnet-20240229-v1:0",
91 Model::Claude3_5Haiku => "us.anthropic.claude-3-5-haiku-20241022-v1:0",
92 Model::Claude3_7Sonnet => "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
93 Model::AmazonNovaLite => "us.amazon.nova-lite-v1:0",
94 Model::AmazonNovaMicro => "us.amazon.nova-micro-v1:0",
95 Model::AmazonNovaPro => "us.amazon.nova-pro-v1:0",
96 Model::DeepSeekR1 => "us.deepseek.r1-v1:0",
97 Model::AI21J2GrandeInstruct => "ai21.j2-grande-instruct",
98 Model::AI21J2JumboInstruct => "ai21.j2-jumbo-instruct",
99 Model::AI21J2Mid => "ai21.j2-mid",
100 Model::AI21J2MidV1 => "ai21.j2-mid-v1",
101 Model::AI21J2Ultra => "ai21.j2-ultra",
102 Model::AI21J2UltraV1_8k => "ai21.j2-ultra-v1:0:8k",
103 Model::AI21J2UltraV1 => "ai21.j2-ultra-v1",
104 Model::AI21JambaInstructV1 => "ai21.jamba-instruct-v1:0",
105 Model::AI21Jamba15LargeV1 => "ai21.jamba-1-5-large-v1:0",
106 Model::AI21Jamba15MiniV1 => "ai21.jamba-1-5-mini-v1:0",
107 Model::CohereCommandTextV14_4k => "cohere.command-text-v14:7:4k",
108 Model::CohereCommandRV1 => "cohere.command-r-v1:0",
109 Model::CohereCommandRPlusV1 => "cohere.command-r-plus-v1:0",
110 Model::CohereCommandLightTextV14_4k => "cohere.command-light-text-v14:7:4k",
111 Model::MetaLlama38BInstructV1 => "meta.llama3-8b-instruct-v1:0",
112 Model::MetaLlama370BInstructV1 => "meta.llama3-70b-instruct-v1:0",
113 Model::MetaLlama318BInstructV1_128k => "meta.llama3-1-8b-instruct-v1:0:128k",
114 Model::MetaLlama318BInstructV1 => "meta.llama3-1-8b-instruct-v1:0",
115 Model::MetaLlama3170BInstructV1_128k => "meta.llama3-1-70b-instruct-v1:0:128k",
116 Model::MetaLlama3170BInstructV1 => "meta.llama3-1-70b-instruct-v1:0",
117 Model::MetaLlama3211BInstructV1 => "meta.llama3-2-11b-instruct-v1:0",
118 Model::MetaLlama3290BInstructV1 => "meta.llama3-2-90b-instruct-v1:0",
119 Model::MetaLlama321BInstructV1 => "meta.llama3-2-1b-instruct-v1:0",
120 Model::MetaLlama323BInstructV1 => "meta.llama3-2-3b-instruct-v1:0",
121 Model::MistralMistral7BInstructV0 => "mistral.mistral-7b-instruct-v0:2",
122 Model::MistralMixtral8x7BInstructV0 => "mistral.mixtral-8x7b-instruct-v0:1",
123 Model::MistralMistralLarge2402V1 => "mistral.mistral-large-2402-v1:0",
124 Model::MistralMistralSmall2402V1 => "mistral.mistral-small-2402-v1:0",
125 Self::Custom { name, .. } => name,
126 }
127 }
128
129 pub fn display_name(&self) -> &str {
130 match self {
131 Self::Claude3_5Sonnet => "Claude 3.5 Sonnet v2",
132 Self::Claude3Opus => "Claude 3 Opus",
133 Self::Claude3Sonnet => "Claude 3 Sonnet",
134 Self::Claude3_5Haiku => "Claude 3.5 Haiku",
135 Self::Claude3_7Sonnet => "Claude 3.7 Sonnet",
136 Self::AmazonNovaLite => "Amazon Nova Lite",
137 Self::AmazonNovaMicro => "Amazon Nova Micro",
138 Self::AmazonNovaPro => "Amazon Nova Pro",
139 Self::DeepSeekR1 => "DeepSeek R1",
140 Self::AI21J2GrandeInstruct => "AI21 Jurassic2 Grande Instruct",
141 Self::AI21J2JumboInstruct => "AI21 Jurassic2 Jumbo Instruct",
142 Self::AI21J2Mid => "AI21 Jurassic2 Mid",
143 Self::AI21J2MidV1 => "AI21 Jurassic2 Mid V1",
144 Self::AI21J2Ultra => "AI21 Jurassic2 Ultra",
145 Self::AI21J2UltraV1_8k => "AI21 Jurassic2 Ultra V1 8K",
146 Self::AI21J2UltraV1 => "AI21 Jurassic2 Ultra V1",
147 Self::AI21JambaInstructV1 => "AI21 Jamba Instruct",
148 Self::AI21Jamba15LargeV1 => "AI21 Jamba 1.5 Large",
149 Self::AI21Jamba15MiniV1 => "AI21 Jamba 1.5 Mini",
150 Self::CohereCommandTextV14_4k => "Cohere Command Text V14 4K",
151 Self::CohereCommandRV1 => "Cohere Command R V1",
152 Self::CohereCommandRPlusV1 => "Cohere Command R Plus V1",
153 Self::CohereCommandLightTextV14_4k => "Cohere Command Light Text V14 4K",
154 Self::MetaLlama38BInstructV1 => "Meta Llama 3 8B Instruct V1",
155 Self::MetaLlama370BInstructV1 => "Meta Llama 3 70B Instruct V1",
156 Self::MetaLlama318BInstructV1_128k => "Meta Llama 3 1.8B Instruct V1 128K",
157 Self::MetaLlama318BInstructV1 => "Meta Llama 3 1.8B Instruct V1",
158 Self::MetaLlama3170BInstructV1_128k => "Meta Llama 3 1 70B Instruct V1 128K",
159 Self::MetaLlama3170BInstructV1 => "Meta Llama 3 1 70B Instruct V1",
160 Self::MetaLlama3211BInstructV1 => "Meta Llama 3 2 11B Instruct V1",
161 Self::MetaLlama3290BInstructV1 => "Meta Llama 3 2 90B Instruct V1",
162 Self::MetaLlama321BInstructV1 => "Meta Llama 3 2 1B Instruct V1",
163 Self::MetaLlama323BInstructV1 => "Meta Llama 3 2 3B Instruct V1",
164 Self::MistralMistral7BInstructV0 => "Mistral 7B Instruct V0",
165 Self::MistralMixtral8x7BInstructV0 => "Mistral Mixtral 8x7B Instruct V0",
166 Self::MistralMistralLarge2402V1 => "Mistral Large 2402 V1",
167 Self::MistralMistralSmall2402V1 => "Mistral Small 2402 V1",
168 Self::Custom {
169 display_name, name, ..
170 } => display_name.as_deref().unwrap_or(name),
171 }
172 }
173
174 pub fn max_token_count(&self) -> usize {
175 match self {
176 Self::Claude3_5Sonnet
177 | Self::Claude3Opus
178 | Self::Claude3Sonnet
179 | Self::Claude3_5Haiku
180 | Self::Claude3_7Sonnet => 200_000,
181 Self::Custom { max_tokens, .. } => *max_tokens,
182 _ => 200_000,
183 }
184 }
185
186 pub fn max_output_tokens(&self) -> u32 {
187 match self {
188 Self::Claude3Opus | Self::Claude3Sonnet | Self::Claude3_5Haiku => 4_096,
189 Self::Claude3_5Sonnet => 8_192,
190 Self::Custom {
191 max_output_tokens, ..
192 } => max_output_tokens.unwrap_or(4_096),
193 _ => 4_096,
194 }
195 }
196
197 pub fn default_temperature(&self) -> f32 {
198 match self {
199 Self::Claude3_5Sonnet
200 | Self::Claude3Opus
201 | Self::Claude3Sonnet
202 | Self::Claude3_5Haiku
203 | Self::Claude3_7Sonnet => 1.0,
204 Self::Custom {
205 default_temperature,
206 ..
207 } => default_temperature.unwrap_or(1.0),
208 _ => 1.0,
209 }
210 }
211}