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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, strum::VariantArray)]
65pub enum Model {
66 ClaudeOpus46,
68 ClaudeOpus45,
69 ClaudeSonnet45,
70
71 ClaudeHaiku45,
73 Gemini3ProPreview,
74 Gpt52,
75 Gpt51CodexMax,
76
77 KimiK25,
79 Gemini3FlashPreview,
80 GLM47,
81 MinimaxM21,
82 Grok41Fast,
83 GrokCodeFast1,
84
85 Qwen3Coder,
87 GptOss120b,
88 OpenRouterAuto,
89
90 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 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 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 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}