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