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(rename = "grok-code-fast-1", alias = "grok-code-fast-1-0825")]
34 GrokCodeFast1,
35 #[serde(rename = "custom")]
36 Custom {
37 name: String,
38 /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
39 display_name: Option<String>,
40 max_tokens: u64,
41 max_output_tokens: Option<u64>,
42 max_completion_tokens: Option<u64>,
43 supports_images: Option<bool>,
44 supports_tools: Option<bool>,
45 parallel_tool_calls: Option<bool>,
46 },
47}
48
49impl Model {
50 pub fn default_fast() -> Self {
51 Self::Grok3Fast
52 }
53
54 pub fn from_id(id: &str) -> Result<Self> {
55 match id {
56 "grok-4" => Ok(Self::Grok4),
57 "grok-4-fast-reasoning" => Ok(Self::Grok4FastReasoning),
58 "grok-4-fast-non-reasoning" => Ok(Self::Grok4FastNonReasoning),
59 "grok-2-vision" => Ok(Self::Grok2Vision),
60 "grok-3" => Ok(Self::Grok3),
61 "grok-3-mini" => Ok(Self::Grok3Mini),
62 "grok-3-fast" => Ok(Self::Grok3Fast),
63 "grok-3-mini-fast" => Ok(Self::Grok3MiniFast),
64 "grok-code-fast-1" => Ok(Self::GrokCodeFast1),
65 _ => anyhow::bail!("invalid model id '{id}'"),
66 }
67 }
68
69 pub fn id(&self) -> &str {
70 match self {
71 Self::Grok2Vision => "grok-2-vision",
72 Self::Grok3 => "grok-3",
73 Self::Grok3Mini => "grok-3-mini",
74 Self::Grok3Fast => "grok-3-fast",
75 Self::Grok3MiniFast => "grok-3-mini-fast",
76 Self::Grok4 => "grok-4",
77 Self::Grok4FastReasoning => "grok-4-fast-reasoning",
78 Self::Grok4FastNonReasoning => "grok-4-fast-non-reasoning",
79 Self::GrokCodeFast1 => "grok-code-fast-1",
80 Self::Custom { name, .. } => name,
81 }
82 }
83
84 pub fn display_name(&self) -> &str {
85 match self {
86 Self::Grok2Vision => "Grok 2 Vision",
87 Self::Grok3 => "Grok 3",
88 Self::Grok3Mini => "Grok 3 Mini",
89 Self::Grok3Fast => "Grok 3 Fast",
90 Self::Grok3MiniFast => "Grok 3 Mini Fast",
91 Self::Grok4 => "Grok 4",
92 Self::Grok4FastReasoning => "Grok 4 Fast",
93 Self::Grok4FastNonReasoning => "Grok 4 Fast (Non-Reasoning)",
94 Self::GrokCodeFast1 => "Grok Code Fast 1",
95 Self::Custom {
96 name, display_name, ..
97 } => display_name.as_ref().unwrap_or(name),
98 }
99 }
100
101 pub fn max_token_count(&self) -> u64 {
102 match self {
103 Self::Grok3 | Self::Grok3Mini | Self::Grok3Fast | Self::Grok3MiniFast => 131_072,
104 Self::Grok4 | Self::GrokCodeFast1 => 256_000,
105 Self::Grok4FastReasoning | Self::Grok4FastNonReasoning => 128_000,
106 Self::Grok2Vision => 8_192,
107 Self::Custom { max_tokens, .. } => *max_tokens,
108 }
109 }
110
111 pub fn max_output_tokens(&self) -> Option<u64> {
112 match self {
113 Self::Grok3 | Self::Grok3Mini | Self::Grok3Fast | Self::Grok3MiniFast => Some(8_192),
114 Self::Grok4
115 | Self::Grok4FastReasoning
116 | Self::Grok4FastNonReasoning
117 | Self::GrokCodeFast1 => Some(64_000),
118 Self::Grok2Vision => Some(4_096),
119 Self::Custom {
120 max_output_tokens, ..
121 } => *max_output_tokens,
122 }
123 }
124
125 pub fn supports_parallel_tool_calls(&self) -> bool {
126 match self {
127 Self::Grok2Vision
128 | Self::Grok3
129 | Self::Grok3Mini
130 | Self::Grok3Fast
131 | Self::Grok3MiniFast
132 | Self::Grok4
133 | Self::Grok4FastReasoning
134 | Self::Grok4FastNonReasoning => true,
135 Self::Custom {
136 parallel_tool_calls: Some(support),
137 ..
138 } => *support,
139 Self::GrokCodeFast1 | Model::Custom { .. } => false,
140 }
141 }
142
143 pub fn supports_prompt_cache_key(&self) -> bool {
144 false
145 }
146
147 pub fn supports_tool(&self) -> bool {
148 match self {
149 Self::Grok2Vision
150 | Self::Grok3
151 | Self::Grok3Mini
152 | Self::Grok3Fast
153 | Self::Grok3MiniFast
154 | Self::Grok4
155 | Self::Grok4FastReasoning
156 | Self::Grok4FastNonReasoning
157 | Self::GrokCodeFast1 => true,
158 Self::Custom {
159 supports_tools: Some(support),
160 ..
161 } => *support,
162 Model::Custom { .. } => false,
163 }
164 }
165
166 pub fn supports_images(&self) -> bool {
167 match self {
168 Self::Grok2Vision => true,
169 Self::Custom {
170 supports_images: Some(support),
171 ..
172 } => *support,
173 _ => false,
174 }
175 }
176}