atm_core/
model.rs

1//! Model identification and metadata.
2//!
3//! Supports known Anthropic models with pricing/context data,
4//! and gracefully handles unknown models (new Anthropic releases
5//! or non-Anthropic models) by preserving their raw ID for display.
6
7use serde::{Deserialize, Serialize};
8use std::fmt;
9
10/// Model identifier.
11///
12/// Parsed from status line JSON: `model.id` field.
13/// Uses prefix matching for forward compatibility with new date variants.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
15#[serde(rename_all = "kebab-case")]
16pub enum Model {
17    /// Claude Opus 4.6 (claude-opus-4-6-*)
18    #[serde(rename = "claude-opus-4-6")]
19    Opus46,
20
21    /// Claude Opus 4.5 (claude-opus-4-5-*)
22    #[serde(rename = "claude-opus-4-5-20251101")]
23    Opus45,
24
25    /// Claude Sonnet 4.5 (claude-sonnet-4-5-*)
26    #[serde(rename = "claude-sonnet-4-5-20250929")]
27    Sonnet45,
28
29    /// Claude Sonnet 4 (claude-sonnet-4-*)
30    #[serde(rename = "claude-sonnet-4-20250514")]
31    Sonnet4,
32
33    /// Claude Haiku 4.5 (claude-haiku-4-5-*)
34    #[serde(rename = "claude-haiku-4-5-20251001")]
35    Haiku45,
36
37    /// Claude Haiku 3.5 (claude-3-5-haiku-*)
38    #[serde(rename = "claude-3-5-haiku-20241022")]
39    Haiku35,
40
41    /// Claude Sonnet 3.5 v2 (claude-3-5-sonnet-*)
42    #[serde(rename = "claude-3-5-sonnet-20241022")]
43    Sonnet35V2,
44
45    /// Unknown or non-Anthropic model
46    #[serde(other)]
47    Unknown,
48}
49
50impl Model {
51    /// Returns a human-readable display name for known models.
52    ///
53    /// For `Unknown` models, callers should use [`derive_display_name`]
54    /// with the raw model ID for a better fallback.
55    pub fn display_name(&self) -> &'static str {
56        match self {
57            Self::Opus46 => "Opus 4.6",
58            Self::Opus45 => "Opus 4.5",
59            Self::Sonnet45 => "Sonnet 4.5",
60            Self::Sonnet4 => "Sonnet 4",
61            Self::Haiku45 => "Haiku 4.5",
62            Self::Haiku35 => "Haiku 3.5",
63            Self::Sonnet35V2 => "Sonnet 3.5 v2",
64            Self::Unknown => "Unknown",
65        }
66    }
67
68    /// Returns the context window size for this model.
69    pub fn context_window_size(&self) -> u32 {
70        match self {
71            Self::Opus46 => 200_000,
72            Self::Opus45 => 200_000,
73            Self::Sonnet45 => 200_000,
74            Self::Sonnet4 => 200_000,
75            Self::Haiku45 => 200_000,
76            Self::Haiku35 => 200_000,
77            Self::Sonnet35V2 => 200_000,
78            Self::Unknown => 200_000, // Default assumption
79        }
80    }
81
82    /// Returns approximate cost per million input tokens (USD).
83    pub fn input_cost_per_million(&self) -> f64 {
84        match self {
85            Self::Opus46 => 5.00,
86            Self::Opus45 => 15.00,
87            Self::Sonnet45 => 3.00,
88            Self::Sonnet4 => 3.00,
89            Self::Haiku45 => 1.00,
90            Self::Haiku35 => 0.80,
91            Self::Sonnet35V2 => 3.00,
92            Self::Unknown => 3.00, // Conservative default
93        }
94    }
95
96    /// Returns approximate cost per million output tokens (USD).
97    pub fn output_cost_per_million(&self) -> f64 {
98        match self {
99            Self::Opus46 => 25.00,
100            Self::Opus45 => 75.00,
101            Self::Sonnet45 => 15.00,
102            Self::Sonnet4 => 15.00,
103            Self::Haiku45 => 5.00,
104            Self::Haiku35 => 4.00,
105            Self::Sonnet35V2 => 15.00,
106            Self::Unknown => 15.00, // Conservative default
107        }
108    }
109
110    /// Parses a model from its string ID using prefix matching.
111    ///
112    /// More specific prefixes are checked first to avoid false matches
113    /// (e.g., "claude-sonnet-4-5" before "claude-sonnet-4-").
114    pub fn from_id(id: &str) -> Self {
115        // More specific (longer) prefixes first to avoid false matches
116        if id.starts_with("claude-opus-4-6") {
117            Self::Opus46
118        } else if id.starts_with("claude-opus-4-5") {
119            Self::Opus45
120        } else if id.starts_with("claude-sonnet-4-5") {
121            Self::Sonnet45
122        } else if id.starts_with("claude-sonnet-4") {
123            Self::Sonnet4
124        } else if id.starts_with("claude-haiku-4-5") {
125            Self::Haiku45
126        } else if id.starts_with("claude-3-5-haiku") {
127            Self::Haiku35
128        } else if id.starts_with("claude-3-5-sonnet") {
129            Self::Sonnet35V2
130        } else {
131            Self::Unknown
132        }
133    }
134
135    /// Returns true if this is an unknown/unrecognized model.
136    pub fn is_unknown(&self) -> bool {
137        matches!(self, Self::Unknown)
138    }
139}
140
141/// Derives a human-readable display name from a raw model ID string.
142///
143/// Used as a fallback when `Model::from_id` returns `Unknown`.
144/// Handles both Claude-style IDs and arbitrary model IDs:
145/// - `"claude-opus-4-7-20260501"` → `"claude-opus-4-7"`
146/// - `"gpt-4o"` → `"gpt-4o"`
147/// - `"gemini-1.5-pro"` → `"gemini-1.5-pro"`
148pub fn derive_display_name(id: &str) -> String {
149    // Strip date suffix if present (pattern: -YYYYMMDD at end)
150    if id.len() > 9 {
151        let potential_date = &id[id.len() - 8..];
152        if potential_date.chars().all(|c| c.is_ascii_digit()) {
153            if let Some(base) = id[..id.len() - 8].strip_suffix('-') {
154                return base.to_string();
155            }
156        }
157    }
158    id.to_string()
159}
160
161impl Default for Model {
162    fn default() -> Self {
163        Self::Unknown
164    }
165}
166
167impl fmt::Display for Model {
168    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
169        write!(f, "{}", self.display_name())
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    // ── Known model parsing ──
178
179    #[test]
180    fn test_model_parsing_opus45() {
181        let model: Model = serde_json::from_str("\"claude-opus-4-5-20251101\"").unwrap();
182        assert_eq!(model, Model::Opus45);
183        assert_eq!(model.display_name(), "Opus 4.5");
184    }
185
186    #[test]
187    fn test_model_parsing_opus46() {
188        // Opus 4.6 may come without date suffix
189        assert_eq!(Model::from_id("claude-opus-4-6"), Model::Opus46);
190        assert_eq!(Model::Opus46.display_name(), "Opus 4.6");
191    }
192
193    #[test]
194    fn test_model_parsing_sonnet45() {
195        assert_eq!(
196            Model::from_id("claude-sonnet-4-5-20250929"),
197            Model::Sonnet45
198        );
199        assert_eq!(Model::Sonnet45.display_name(), "Sonnet 4.5");
200    }
201
202    #[test]
203    fn test_model_parsing_haiku45() {
204        assert_eq!(Model::from_id("claude-haiku-4-5-20251001"), Model::Haiku45);
205        assert_eq!(Model::Haiku45.display_name(), "Haiku 4.5");
206    }
207
208    #[test]
209    fn test_model_unknown_serde() {
210        let model: Model = serde_json::from_str("\"gpt-4o\"").unwrap();
211        assert_eq!(model, Model::Unknown);
212    }
213
214    // ── Prefix matching ──
215
216    #[test]
217    fn test_from_id_prefix_exact() {
218        assert_eq!(Model::from_id("claude-opus-4-5-20251101"), Model::Opus45);
219        assert_eq!(Model::from_id("claude-sonnet-4-20250514"), Model::Sonnet4);
220    }
221
222    #[test]
223    fn test_from_id_prefix_with_different_date() {
224        // Future date variants should still match the right model family
225        assert_eq!(Model::from_id("claude-opus-4-6-20260301"), Model::Opus46);
226        assert_eq!(
227            Model::from_id("claude-sonnet-4-5-20261201"),
228            Model::Sonnet45
229        );
230        assert_eq!(Model::from_id("claude-haiku-4-5-20260601"), Model::Haiku45);
231        assert_eq!(Model::from_id("claude-opus-4-5-20260101"), Model::Opus45);
232    }
233
234    #[test]
235    fn test_from_id_prefix_no_date() {
236        // Model IDs without date suffix
237        assert_eq!(Model::from_id("claude-opus-4-6"), Model::Opus46);
238        assert_eq!(Model::from_id("claude-sonnet-4-5"), Model::Sonnet45);
239    }
240
241    #[test]
242    fn test_from_id_sonnet4_not_confused_with_sonnet45() {
243        // "claude-sonnet-4-5" should match Sonnet45, not Sonnet4
244        assert_eq!(
245            Model::from_id("claude-sonnet-4-5-20250929"),
246            Model::Sonnet45
247        );
248        // "claude-sonnet-4-20250514" should match Sonnet4
249        assert_eq!(Model::from_id("claude-sonnet-4-20250514"), Model::Sonnet4);
250    }
251
252    // ── Unknown / non-Anthropic models ──
253
254    #[test]
255    fn test_from_id_unknown() {
256        assert_eq!(Model::from_id("gpt-4o"), Model::Unknown);
257        assert_eq!(Model::from_id("gemini-1.5-pro"), Model::Unknown);
258        assert_eq!(Model::from_id("llama-3-70b"), Model::Unknown);
259        assert_eq!(Model::from_id("unknown-model"), Model::Unknown);
260    }
261
262    #[test]
263    fn test_is_unknown() {
264        assert!(Model::Unknown.is_unknown());
265        assert!(!Model::Opus46.is_unknown());
266    }
267
268    // ── Display name derivation for unknown models ──
269
270    #[test]
271    fn test_derive_display_name_strips_date() {
272        assert_eq!(
273            derive_display_name("claude-opus-4-7-20260501"),
274            "claude-opus-4-7"
275        );
276        assert_eq!(
277            derive_display_name("claude-sonnet-5-20270101"),
278            "claude-sonnet-5"
279        );
280    }
281
282    #[test]
283    fn test_derive_display_name_no_date() {
284        assert_eq!(derive_display_name("gpt-4o"), "gpt-4o");
285        assert_eq!(derive_display_name("gemini-1.5-pro"), "gemini-1.5-pro");
286    }
287
288    #[test]
289    fn test_derive_display_name_short_ids() {
290        assert_eq!(derive_display_name("gpt-4"), "gpt-4");
291        assert_eq!(derive_display_name("o1"), "o1");
292    }
293}
atm_core/model.rs

atm_core/
model.rs
