claudectl_core/
models.rs

1use std::collections::HashMap;
2use std::sync::{Mutex, OnceLock};
3
4#[derive(Debug, Clone, Copy, PartialEq)]
5pub struct ModelProfile {
6    pub input_per_m: f64,
7    pub output_per_m: f64,
8    pub cache_read_per_m: f64,
9    pub cache_write_per_m: f64,
10    pub context_max: u64,
11}
12
13#[derive(Debug, Clone, PartialEq)]
14pub struct ModelOverride {
15    pub name: String,
16    pub profile: ModelProfile,
17}
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ModelProfileSource {
21    BuiltIn,
22    Override,
23    Fallback,
24}
25
26impl ModelProfileSource {
27    pub fn label(self) -> &'static str {
28        match self {
29            Self::BuiltIn => "built-in",
30            Self::Override => "override",
31            Self::Fallback => "fallback",
32        }
33    }
34}
35
36#[derive(Debug, Clone, PartialEq)]
37pub struct ResolvedModelProfile {
38    pub key: String,
39    pub profile: ModelProfile,
40    pub source: ModelProfileSource,
41}
42
43static MODEL_OVERRIDES: OnceLock<Mutex<HashMap<String, ModelProfile>>> = OnceLock::new();
44
45pub fn shorten_model(model: &str) -> String {
46    if model.contains("opus") {
47        if model.contains("4-6") {
48            "opus-4.6".into()
49        } else {
50            "opus".into()
51        }
52    } else if model.contains("sonnet") {
53        if model.contains("4-6") {
54            "sonnet-4.6".into()
55        } else {
56            "sonnet".into()
57        }
58    } else if model.contains("haiku") {
59        "haiku".into()
60    } else {
61        model.to_string()
62    }
63}
64
65pub fn set_overrides(overrides: Vec<ModelOverride>) {
66    let store = MODEL_OVERRIDES.get_or_init(|| Mutex::new(HashMap::new()));
67    let Ok(mut guard) = store.lock() else {
68        return;
69    };
70    guard.clear();
71    for override_ in overrides {
72        let raw = override_.name.trim().to_lowercase();
73        let shortened = shorten_model(&override_.name).to_lowercase();
74        guard.insert(raw, override_.profile);
75        guard.insert(shortened, override_.profile);
76    }
77}
78
79pub fn resolve(model: &str) -> ResolvedModelProfile {
80    let empty = HashMap::new();
81    let store = MODEL_OVERRIDES.get_or_init(|| Mutex::new(HashMap::new()));
82    let guard = store.lock().ok();
83    let overrides = guard.as_deref().unwrap_or(&empty);
84    resolve_with_overrides(model, overrides)
85}
86
87pub(crate) fn resolve_with_overrides(
88    model: &str,
89    overrides: &HashMap<String, ModelProfile>,
90) -> ResolvedModelProfile {
91    let raw_key = model.trim().to_lowercase();
92    let short_key = shorten_model(model).to_lowercase();
93
94    if let Some(profile) = overrides
95        .get(&raw_key)
96        .or_else(|| overrides.get(&short_key))
97        .copied()
98    {
99        return ResolvedModelProfile {
100            key: if raw_key.is_empty() {
101                short_key
102            } else {
103                raw_key
104            },
105            profile,
106            source: ModelProfileSource::Override,
107        };
108    }
109
110    if let Some(profile) = built_in_profile(&short_key) {
111        return ResolvedModelProfile {
112            key: short_key,
113            profile,
114            source: ModelProfileSource::BuiltIn,
115        };
116    }
117
118    ResolvedModelProfile {
119        key: if short_key.is_empty() {
120            "unknown".into()
121        } else {
122            short_key
123        },
124        profile: fallback_profile(),
125        source: ModelProfileSource::Fallback,
126    }
127}
128
129fn built_in_profile(key: &str) -> Option<ModelProfile> {
130    match key {
131        "opus-4.6" | "opus" => Some(ModelProfile {
132            input_per_m: 15.0,
133            output_per_m: 75.0,
134            cache_read_per_m: 1.875,
135            cache_write_per_m: 18.75,
136            context_max: 1_000_000,
137        }),
138        "sonnet-4.6" | "sonnet" => Some(ModelProfile {
139            input_per_m: 3.0,
140            output_per_m: 15.0,
141            cache_read_per_m: 0.375,
142            cache_write_per_m: 3.75,
143            context_max: 200_000,
144        }),
145        "haiku" => Some(ModelProfile {
146            input_per_m: 0.80,
147            output_per_m: 4.0,
148            cache_read_per_m: 0.10,
149            cache_write_per_m: 1.0,
150            context_max: 200_000,
151        }),
152        _ => None,
153    }
154}
155
156fn fallback_profile() -> ModelProfile {
157    ModelProfile {
158        input_per_m: 15.0,
159        output_per_m: 75.0,
160        cache_read_per_m: 1.875,
161        cache_write_per_m: 18.75,
162        context_max: 200_000,
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    #[test]
171    fn resolve_builtin_profile() {
172        let resolved = resolve_with_overrides("claude-opus-4-6-20260401", &HashMap::new());
173        assert_eq!(resolved.source, ModelProfileSource::BuiltIn);
174        assert_eq!(resolved.profile.context_max, 1_000_000);
175    }
176
177    #[test]
178    fn resolve_override_profile() {
179        let mut overrides = HashMap::new();
180        overrides.insert(
181            "gpt-4o".into(),
182            ModelProfile {
183                input_per_m: 1.0,
184                output_per_m: 2.0,
185                cache_read_per_m: 0.5,
186                cache_write_per_m: 1.5,
187                context_max: 128_000,
188            },
189        );
190        let resolved = resolve_with_overrides("gpt-4o", &overrides);
191        assert_eq!(resolved.source, ModelProfileSource::Override);
192        assert_eq!(resolved.profile.context_max, 128_000);
193    }
194
195    #[test]
196    fn resolve_fallback_profile() {
197        let resolved = resolve_with_overrides("mystery-model", &HashMap::new());
198        assert_eq!(resolved.source, ModelProfileSource::Fallback);
199        assert_eq!(resolved.profile.context_max, 200_000);
200    }
201}
claudectl_core/models.rs

claudectl_core/
models.rs
