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 const 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 const 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 const 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 const 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 const fn supports_prompt_cache_key(&self) -> bool {
144        false
145    }
146
147    pub const 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 const 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}