Skip to main content

atm_core/
model.rs

1//! Claude model identification and metadata.
2
3use serde::{Deserialize, Serialize};
4use std::fmt;
5
6/// Claude model identifier.
7///
8/// Parsed from status line JSON: `model.id` field.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
10#[serde(rename_all = "kebab-case")]
11pub enum Model {
12    /// Claude Opus 4.5 (claude-opus-4-5-20251101)
13    #[serde(rename = "claude-opus-4-5-20251101")]
14    Opus45,
15
16    /// Claude Sonnet 4 (claude-sonnet-4-20250514)
17    #[serde(rename = "claude-sonnet-4-20250514")]
18    Sonnet4,
19
20    /// Claude Haiku 3.5 (claude-3-5-haiku-20241022)
21    #[serde(rename = "claude-3-5-haiku-20241022")]
22    Haiku35,
23
24    /// Claude Sonnet 3.5 v2 (claude-3-5-sonnet-20241022)
25    #[serde(rename = "claude-3-5-sonnet-20241022")]
26    Sonnet35V2,
27
28    /// Unknown or future model
29    #[serde(other)]
30    Unknown,
31}
32
33impl Model {
34    /// Returns a human-readable display name.
35    pub fn display_name(&self) -> &'static str {
36        match self {
37            Self::Opus45 => "Opus 4.5",
38            Self::Sonnet4 => "Sonnet 4",
39            Self::Haiku35 => "Haiku 3.5",
40            Self::Sonnet35V2 => "Sonnet 3.5 v2",
41            Self::Unknown => "Unknown",
42        }
43    }
44
45    /// Returns the context window size for this model.
46    pub fn context_window_size(&self) -> u32 {
47        match self {
48            Self::Opus45 => 200_000,
49            Self::Sonnet4 => 200_000,
50            Self::Haiku35 => 200_000,
51            Self::Sonnet35V2 => 200_000,
52            Self::Unknown => 200_000, // Default assumption
53        }
54    }
55
56    /// Returns approximate cost per million input tokens (USD).
57    pub fn input_cost_per_million(&self) -> f64 {
58        match self {
59            Self::Opus45 => 15.00,
60            Self::Sonnet4 => 3.00,
61            Self::Haiku35 => 0.80,
62            Self::Sonnet35V2 => 3.00,
63            Self::Unknown => 3.00, // Conservative default
64        }
65    }
66
67    /// Returns approximate cost per million output tokens (USD).
68    pub fn output_cost_per_million(&self) -> f64 {
69        match self {
70            Self::Opus45 => 75.00,
71            Self::Sonnet4 => 15.00,
72            Self::Haiku35 => 4.00,
73            Self::Sonnet35V2 => 15.00,
74            Self::Unknown => 15.00, // Conservative default
75        }
76    }
77
78    /// Parses a model from its string ID.
79    pub fn from_id(id: &str) -> Self {
80        match id {
81            "claude-opus-4-5-20251101" => Self::Opus45,
82            "claude-sonnet-4-20250514" => Self::Sonnet4,
83            "claude-3-5-haiku-20241022" => Self::Haiku35,
84            "claude-3-5-sonnet-20241022" => Self::Sonnet35V2,
85            _ => Self::Unknown,
86        }
87    }
88}
89
90impl Default for Model {
91    fn default() -> Self {
92        Self::Unknown
93    }
94}
95
96impl fmt::Display for Model {
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        write!(f, "{}", self.display_name())
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn test_model_parsing() {
108        let model: Model = serde_json::from_str("\"claude-opus-4-5-20251101\"").unwrap();
109        assert_eq!(model, Model::Opus45);
110        assert_eq!(model.display_name(), "Opus 4.5");
111    }
112
113    #[test]
114    fn test_model_unknown() {
115        let model: Model = serde_json::from_str("\"claude-future-model\"").unwrap();
116        assert_eq!(model, Model::Unknown);
117    }
118
119    #[test]
120    fn test_model_from_id() {
121        assert_eq!(Model::from_id("claude-opus-4-5-20251101"), Model::Opus45);
122        assert_eq!(Model::from_id("unknown-model"), Model::Unknown);
123    }
124}