Skip to main content

tycode_core/ai/
model.rs

1use crate::ai::provider::AiProvider;
2use crate::ai::tweaks::{ModelTweaks, RegistryFileModificationApi};
3use crate::ai::types::ReasoningBudget;
4use crate::ai::ModelSettings;
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7use strum::VariantArray;
8
9#[derive(
10    Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord, Default, JsonSchema,
11)]
12#[serde(rename_all = "snake_case")]
13pub enum ModelCost {
14    Free,
15    Low,
16    #[default]
17    Medium,
18    High,
19    Unlimited,
20}
21
22impl ModelCost {
23    pub const fn all_levels() -> [Self; 5] {
24        [
25            Self::Free,
26            Self::Low,
27            Self::Medium,
28            Self::High,
29            Self::Unlimited,
30        ]
31    }
32
33    pub const fn description(self) -> &'static str {
34        match self {
35            Self::Free => "Restrict to free models only. Your data will likely used for training.",
36            Self::Low => "Under $1/million tokens",
37            Self::Medium => "Under $5/million tokens",
38            Self::High => "Under $15/million tokens",
39            Self::Unlimited => "No restrictions",
40        }
41    }
42}
43
44impl TryFrom<&str> for ModelCost {
45    type Error = String;
46
47    fn try_from(value: &str) -> Result<Self, Self::Error> {
48        let lower = value.to_lowercase();
49        match lower.as_str() {
50            "free" => Ok(Self::Free),
51            "low" => Ok(Self::Low),
52            "medium" => Ok(Self::Medium),
53            "high" => Ok(Self::High),
54            "unlimited" => Ok(Self::Unlimited),
55            _ => Err(format!(
56                "Invalid model cost level: {}. Valid options: free, low, medium, high, unlimited",
57                value
58            )),
59        }
60    }
61}
62
63/// The supported models, subjectively ranked by quality
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, strum::VariantArray)]
65pub enum Model {
66    // The default models for unlimited/high budget
67    ClaudeOpus46,
68    ClaudeOpus45,
69    ClaudeSonnet45,
70
71    // Medium cost tier
72    ClaudeHaiku45,
73    Gemini3ProPreview,
74    Gpt52,
75    Gpt51CodexMax,
76
77    // Low cost models
78    KimiK25,
79    Gemini3FlashPreview,
80    GLM47,
81    MinimaxM21,
82    Grok41Fast,
83    GrokCodeFast1,
84
85    // Even lower cost models
86    Qwen3Coder,
87    GptOss120b,
88    OpenRouterAuto,
89
90    /// This allows code to match all models, but still match _ => to
91    /// avoid being *required* to match all models.
92    None,
93}
94
95impl Model {
96    pub fn tweaks(self) -> ModelTweaks {
97        match self {
98            Self::Gpt52 | Self::Gpt51CodexMax => ModelTweaks {
99                file_modification_api: Some(RegistryFileModificationApi::Patch),
100                ..Default::default()
101            },
102            _ => ModelTweaks {
103                file_modification_api: Some(RegistryFileModificationApi::FindReplace),
104                ..Default::default()
105            },
106        }
107    }
108
109    pub const fn name(self) -> &'static str {
110        match self {
111            Self::ClaudeSonnet45 => "claude-sonnet-45",
112            Self::ClaudeOpus46 => "claude-opus-4-6",
113            Self::ClaudeOpus45 => "claude-opus-4-5",
114            Self::ClaudeHaiku45 => "claude-haiku-45",
115
116            Self::Gemini3ProPreview => "gemini-3-pro-preview",
117            Self::Gemini3FlashPreview => "gemini-3-flash-preview",
118
119            Self::Gpt52 => "gpt-5-2",
120            Self::Gpt51CodexMax => "gpt-5-1-codex-max",
121            Self::GptOss120b => "gpt-oss-120b",
122
123            Self::GLM47 => "glm-4-7",
124            Self::MinimaxM21 => "minimax-m2-1",
125
126            Self::Grok41Fast => "grok-4-1-fast",
127            Self::GrokCodeFast1 => "grok-code-fast-1",
128            Self::KimiK25 => "kimi-k2-5",
129
130            Self::Qwen3Coder => "qwen3-coder",
131
132            Self::OpenRouterAuto => "openrouter/auto",
133            Self::None => "None",
134        }
135    }
136
137    pub fn from_name(s: &str) -> Option<Self> {
138        match s {
139            "claude-sonnet-45" => Some(Self::ClaudeSonnet45),
140            "claude-opus-4-6" => Some(Self::ClaudeOpus46),
141            "claude-opus-4-5" => Some(Self::ClaudeOpus45),
142            "claude-haiku-45" => Some(Self::ClaudeHaiku45),
143            "gemini-3-pro-preview" => Some(Self::Gemini3ProPreview),
144            "gemini-3-flash-preview" => Some(Self::Gemini3FlashPreview),
145            "gpt-5-2" => Some(Self::Gpt52),
146            "gpt-5-1-codex-max" => Some(Self::Gpt51CodexMax),
147            "gpt-oss-120b" => Some(Self::GptOss120b),
148            "glm-4-7" => Some(Self::GLM47),
149            "minimax-m2-1" => Some(Self::MinimaxM21),
150            "grok-4-1-fast" => Some(Self::Grok41Fast),
151            "grok-code-fast-1" => Some(Self::GrokCodeFast1),
152            "kimi-k2-5" => Some(Self::KimiK25),
153            "qwen3-coder" => Some(Self::Qwen3Coder),
154            "openrouter/auto" => Some(Self::OpenRouterAuto),
155            _ => None,
156        }
157    }
158
159    pub const fn supports_prompt_caching(self) -> bool {
160        match self {
161            Self::ClaudeSonnet45
162            | Self::ClaudeOpus46
163            | Self::ClaudeOpus45
164            | Self::ClaudeHaiku45 => true,
165            Self::OpenRouterAuto => false,
166            _ => false,
167        }
168    }
169
170    // Return default model settings for the model
171    pub fn default_settings(self) -> ModelSettings {
172        ModelSettings {
173            model: self,
174            max_tokens: Some(32000),
175            temperature: Some(1.0),
176            top_p: None,
177            reasoning_budget: ReasoningBudget::High,
178        }
179    }
180
181    /// Select the highest quality model supported by the provider that fits the cost threshold.
182    /// For Unlimited, returns the highest supported model. For others, filters by max(input/output cost per million tokens) <= threshold.
183    /// Free requires exact 0.0 match. Ranked highest-to-lowest to prefer premium within budget.
184    /// Returns None if no fit (surfaces error to caller—no fallback).
185    pub fn select_for_cost(provider: &dyn AiProvider, quality: ModelCost) -> Option<ModelSettings> {
186        let supported = provider.supported_models();
187        let models: Vec<&'static Model> = Model::VARIANTS
188            .into_iter()
189            .filter(|m| supported.contains(m))
190            .collect();
191
192        let threshold = match quality {
193            ModelCost::Free => 0.0,
194            ModelCost::Low => 1.0,
195            ModelCost::Medium => 3.0,
196            ModelCost::High => 10.0,
197            ModelCost::Unlimited => f64::MAX,
198        };
199
200        for model in models {
201            let cost = provider.get_cost(model);
202            // assume 5 is to 1 input to output
203            let cost = (cost.input_cost_per_million_tokens * 5.0
204                + cost.output_cost_per_million_tokens)
205                / 6.0;
206            if cost <= threshold {
207                return Some(model.default_settings());
208            }
209        }
210
211        None
212    }
213}