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