tenant.openagents/omega

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    #[default]
11    #[serde(rename = "grok-4.3", alias = "grok-4.3-latest")]
12    Grok43,
13    #[serde(rename = "grok-4.20-0309-reasoning")]
14    Grok420Reasoning,
15    #[serde(rename = "grok-4.20-0309-non-reasoning")]
16    Grok420NonReasoning,
17    #[serde(rename = "custom")]
18    Custom {
19        name: String,
20        /// The name displayed in the UI, such as in the agent panel model dropdown menu.
21        display_name: Option<String>,
22        max_tokens: u64,
23        max_output_tokens: Option<u64>,
24        max_completion_tokens: Option<u64>,
25        supports_images: Option<bool>,
26        supports_tools: Option<bool>,
27        parallel_tool_calls: Option<bool>,
28    },
29}
30
31impl Model {
32    pub fn default_fast() -> Self {
33        Self::Grok43
34    }
35
36    pub fn from_id(id: &str) -> Result<Self> {
37        match id {
38            "grok-4.3" => Ok(Self::Grok43),
39            "grok-4.20-0309-reasoning" => Ok(Self::Grok420Reasoning),
40            "grok-4.20-0309-non-reasoning" => Ok(Self::Grok420NonReasoning),
41            _ => anyhow::bail!("invalid model id '{id}'"),
42        }
43    }
44
45    pub fn id(&self) -> &str {
46        match self {
47            Self::Grok43 => "grok-4.3",
48            Self::Grok420Reasoning => "grok-4.20-0309-reasoning",
49            Self::Grok420NonReasoning => "grok-4.20-0309-non-reasoning",
50            Self::Custom { name, .. } => name,
51        }
52    }
53
54    pub fn display_name(&self) -> &str {
55        match self {
56            Self::Grok43 => "Grok 4.3",
57            Self::Grok420Reasoning => "Grok 4.20 Reasoning",
58            Self::Grok420NonReasoning => "Grok 4.20 (Non-Reasoning)",
59            Self::Custom {
60                name, display_name, ..
61            } => display_name.as_ref().unwrap_or(name),
62        }
63    }
64
65    pub fn max_token_count(&self) -> u64 {
66        match self {
67            Self::Grok43 => 1_000_000,
68            Self::Grok420Reasoning | Self::Grok420NonReasoning => 2_000_000,
69            Self::Custom { max_tokens, .. } => *max_tokens,
70        }
71    }
72
73    pub fn max_output_tokens(&self) -> Option<u64> {
74        match self {
75            Self::Grok43 | Self::Grok420Reasoning | Self::Grok420NonReasoning => Some(64_000),
76            Self::Custom {
77                max_output_tokens, ..
78            } => *max_output_tokens,
79        }
80    }
81
82    pub fn supports_parallel_tool_calls(&self) -> bool {
83        match self {
84            Self::Grok43 | Self::Grok420Reasoning | Self::Grok420NonReasoning => true,
85            Self::Custom {
86                parallel_tool_calls: Some(support),
87                ..
88            } => *support,
89            Model::Custom { .. } => false,
90        }
91    }
92
93    pub fn requires_json_schema_subset(&self) -> bool {
94        match self {
95            Self::Grok43 | Self::Grok420Reasoning | Self::Grok420NonReasoning => true,
96            Self::Custom { .. } => false,
97        }
98    }
99
100    pub fn supports_prompt_cache_key(&self) -> bool {
101        false
102    }
103
104    pub fn supports_tool(&self) -> bool {
105        match self {
106            Self::Grok43 | Self::Grok420Reasoning | Self::Grok420NonReasoning => true,
107            Self::Custom {
108                supports_tools: Some(support),
109                ..
110            } => *support,
111            Model::Custom { .. } => false,
112        }
113    }
114
115    pub fn supports_images(&self) -> bool {
116        match self {
117            Self::Grok43 | Self::Grok420Reasoning | Self::Grok420NonReasoning => true,
118            Self::Custom {
119                supports_images: Some(support),
120                ..
121            } => *support,
122            Self::Custom { .. } => false,
123        }
124    }
125
126    pub fn supports_reasoning_effort(&self) -> bool {
127        match self {
128            Self::Grok43 => true,
129            Self::Grok420Reasoning | Self::Grok420NonReasoning | Self::Custom { .. } => false,
130        }
131    }
132}
133