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