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

atm_core/
model.rs
