1use anyhow::Result;
2use serde::{Deserialize, Serialize};
3use strum::EnumIter;
4
5pub const XAI_API_URL: &str = "https://api.x.ai/v1";
6
7#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
8#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)]
9pub enum Model {
10 #[serde(rename = "grok-2-vision-latest")]
11 Grok2Vision,
12 #[default]
13 #[serde(rename = "grok-3-latest")]
14 Grok3,
15 #[serde(rename = "grok-3-mini-latest")]
16 Grok3Mini,
17 #[serde(rename = "grok-3-fast-latest")]
18 Grok3Fast,
19 #[serde(rename = "grok-3-mini-fast-latest")]
20 Grok3MiniFast,
21 #[serde(rename = "grok-4", alias = "grok-4-latest")]
22 Grok4,
23 #[serde(
24 rename = "grok-4-fast-reasoning",
25 alias = "grok-4-fast-reasoning-latest"
26 )]
27 Grok4FastReasoning,
28 #[serde(
29 rename = "grok-4-fast-non-reasoning",
30 alias = "grok-4-fast-non-reasoning-latest"
31 )]
32 Grok4FastNonReasoning,
33 #[serde(
34 rename = "grok-4-1-fast-non-reasoning",
35 alias = "grok-4-1-fast-non-reasoning-latest"
36 )]
37 Grok41FastNonReasoning,
38 #[serde(
39 rename = "grok-4-1-fast-reasoning",
40 alias = "grok-4-1-fast-reasoning-latest",
41 alias = "grok-4-1-fast"
42 )]
43 Grok41FastReasoning,
44 #[serde(rename = "grok-code-fast-1", alias = "grok-code-fast-1-0825")]
45 GrokCodeFast1,
46 #[serde(rename = "custom")]
47 Custom {
48 name: String,
49 /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
50 display_name: Option<String>,
51 max_tokens: u64,
52 max_output_tokens: Option<u64>,
53 max_completion_tokens: Option<u64>,
54 supports_images: Option<bool>,
55 supports_tools: Option<bool>,
56 parallel_tool_calls: Option<bool>,
57 },
58}
59
60impl Model {
61 pub fn default_fast() -> Self {
62 Self::Grok3Fast
63 }
64
65 pub fn from_id(id: &str) -> Result<Self> {
66 match id {
67 "grok-4" => Ok(Self::Grok4),
68 "grok-4-fast-reasoning" => Ok(Self::Grok4FastReasoning),
69 "grok-4-fast-non-reasoning" => Ok(Self::Grok4FastNonReasoning),
70 "grok-4-1-fast-non-reasoning" => Ok(Self::Grok41FastNonReasoning),
71 "grok-4-1-fast-reasoning" => Ok(Self::Grok41FastReasoning),
72 "grok-4-1-fast" => Ok(Self::Grok41FastReasoning),
73 "grok-2-vision" => Ok(Self::Grok2Vision),
74 "grok-3" => Ok(Self::Grok3),
75 "grok-3-mini" => Ok(Self::Grok3Mini),
76 "grok-3-fast" => Ok(Self::Grok3Fast),
77 "grok-3-mini-fast" => Ok(Self::Grok3MiniFast),
78 "grok-code-fast-1" => Ok(Self::GrokCodeFast1),
79 _ => anyhow::bail!("invalid model id '{id}'"),
80 }
81 }
82
83 pub fn id(&self) -> &str {
84 match self {
85 Self::Grok2Vision => "grok-2-vision",
86 Self::Grok3 => "grok-3",
87 Self::Grok3Mini => "grok-3-mini",
88 Self::Grok3Fast => "grok-3-fast",
89 Self::Grok3MiniFast => "grok-3-mini-fast",
90 Self::Grok4 => "grok-4",
91 Self::Grok4FastReasoning => "grok-4-fast-reasoning",
92 Self::Grok4FastNonReasoning => "grok-4-fast-non-reasoning",
93 Self::Grok41FastNonReasoning => "grok-4-1-fast-non-reasoning",
94 Self::Grok41FastReasoning => "grok-4-1-fast-reasoning",
95 Self::GrokCodeFast1 => "grok-code-fast-1",
96 Self::Custom { name, .. } => name,
97 }
98 }
99
100 pub fn display_name(&self) -> &str {
101 match self {
102 Self::Grok2Vision => "Grok 2 Vision",
103 Self::Grok3 => "Grok 3",
104 Self::Grok3Mini => "Grok 3 Mini",
105 Self::Grok3Fast => "Grok 3 Fast",
106 Self::Grok3MiniFast => "Grok 3 Mini Fast",
107 Self::Grok4 => "Grok 4",
108 Self::Grok4FastReasoning => "Grok 4 Fast",
109 Self::Grok4FastNonReasoning => "Grok 4 Fast (Non-Reasoning)",
110 Self::Grok41FastNonReasoning => "Grok 4.1 Fast (Non-Reasoning)",
111 Self::Grok41FastReasoning => "Grok 4.1 Fast",
112 Self::GrokCodeFast1 => "Grok Code Fast 1",
113 Self::Custom {
114 name, display_name, ..
115 } => display_name.as_ref().unwrap_or(name),
116 }
117 }
118
119 pub fn max_token_count(&self) -> u64 {
120 match self {
121 Self::Grok3 | Self::Grok3Mini | Self::Grok3Fast | Self::Grok3MiniFast => 131_072,
122 Self::Grok4 | Self::GrokCodeFast1 => 256_000,
123 Self::Grok4FastReasoning
124 | Self::Grok4FastNonReasoning
125 | Self::Grok41FastNonReasoning
126 | Self::Grok41FastReasoning => 2_000_000,
127 Self::Grok2Vision => 8_192,
128 Self::Custom { max_tokens, .. } => *max_tokens,
129 }
130 }
131
132 pub fn max_output_tokens(&self) -> Option<u64> {
133 match self {
134 Self::Grok3 | Self::Grok3Mini | Self::Grok3Fast | Self::Grok3MiniFast => Some(8_192),
135 Self::Grok4
136 | Self::Grok4FastReasoning
137 | Self::Grok4FastNonReasoning
138 | Self::Grok41FastNonReasoning
139 | Self::Grok41FastReasoning
140 | Self::GrokCodeFast1 => Some(64_000),
141 Self::Grok2Vision => Some(4_096),
142 Self::Custom {
143 max_output_tokens, ..
144 } => *max_output_tokens,
145 }
146 }
147
148 pub fn supports_parallel_tool_calls(&self) -> bool {
149 match self {
150 Self::Grok2Vision
151 | Self::Grok3
152 | Self::Grok3Mini
153 | Self::Grok3Fast
154 | Self::Grok3MiniFast
155 | Self::Grok4
156 | Self::Grok4FastReasoning
157 | Self::Grok4FastNonReasoning
158 | Self::Grok41FastNonReasoning
159 | Self::Grok41FastReasoning => true,
160 Self::Custom {
161 parallel_tool_calls: Some(support),
162 ..
163 } => *support,
164 Self::GrokCodeFast1 | Model::Custom { .. } => false,
165 }
166 }
167
168 pub fn supports_prompt_cache_key(&self) -> bool {
169 false
170 }
171
172 pub fn supports_tool(&self) -> bool {
173 match self {
174 Self::Grok2Vision
175 | Self::Grok3
176 | Self::Grok3Mini
177 | Self::Grok3Fast
178 | Self::Grok3MiniFast
179 | Self::Grok4
180 | Self::Grok4FastReasoning
181 | Self::Grok4FastNonReasoning
182 | Self::Grok41FastNonReasoning
183 | Self::Grok41FastReasoning
184 | Self::GrokCodeFast1 => true,
185 Self::Custom {
186 supports_tools: Some(support),
187 ..
188 } => *support,
189 Model::Custom { .. } => false,
190 }
191 }
192
193 pub fn supports_images(&self) -> bool {
194 match self {
195 Self::Grok2Vision
196 | Self::Grok4
197 | Self::Grok4FastReasoning
198 | Self::Grok4FastNonReasoning
199 | Self::Grok41FastNonReasoning
200 | Self::Grok41FastReasoning => true,
201 Self::Custom {
202 supports_images: Some(support),
203 ..
204 } => *support,
205 _ => false,
206 }
207 }
208}