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