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