1use serde::{Deserialize, Serialize};
9
10use crate::types::{ModelEntry, ModelMode};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, Hash)]
14#[serde(rename_all = "lowercase")]
15pub enum ModelTier {
16 Advanced,
18 #[default]
20 Balanced,
21 Fast,
23 ImageGeneration,
25}
26
27impl ModelTier {
28 pub const ALL: &'static [Self] = &[
30 Self::Advanced,
31 Self::Balanced,
32 Self::Fast,
33 Self::ImageGeneration,
34 ];
35
36 #[must_use]
38 pub const fn display_name(self) -> &'static str {
39 match self {
40 Self::Advanced => "Advanced (Flagship)",
41 Self::Balanced => "Balanced",
42 Self::Fast => "Fast",
43 Self::ImageGeneration => "Image Generation",
44 }
45 }
46
47 #[must_use]
49 pub const fn short_name(self) -> &'static str {
50 match self {
51 Self::Advanced => "advanced",
52 Self::Balanced => "balanced",
53 Self::Fast => "fast",
54 Self::ImageGeneration => "image",
55 }
56 }
57}
58
59#[derive(Debug, Clone, Copy)]
65pub struct ChatPriceBreakpoints {
66 fast_ceiling: f64,
67 advanced_floor: f64,
68}
69
70impl ChatPriceBreakpoints {
71 pub fn compute<'a>(entries: impl Iterator<Item = &'a ModelEntry>) -> Option<Self> {
76 let mut prices: Vec<f64> = entries
77 .filter(|info| info.mode() == ModelMode::Chat)
78 .map(|info| info.pricing().output_per_token())
79 .filter(|price| price.is_finite() && *price > 0.0)
80 .collect();
81 if prices.len() < 3 {
82 return None;
83 }
84
85 prices.sort_by(f64::total_cmp);
86 let low_index = prices.len() / 3;
87 let high_index = (prices.len() * 2) / 3;
88 Some(Self {
89 fast_ceiling: prices[low_index],
90 advanced_floor: prices[high_index],
91 })
92 }
93
94 pub fn classify(self, info: &ModelEntry) -> Option<ModelTier> {
95 let output_price = info.pricing().output_per_token();
96 if !output_price.is_finite() || output_price <= 0.0 {
97 return None;
98 }
99 if output_price <= self.fast_ceiling {
100 return Some(ModelTier::Fast);
101 }
102 if output_price >= self.advanced_floor {
103 return Some(ModelTier::Advanced);
104 }
105 Some(ModelTier::Balanced)
106 }
107}
108
109pub fn classify_entry(info: &ModelEntry, breakpoints: Option<ChatPriceBreakpoints>) -> ModelTier {
111 match info.mode() {
112 ModelMode::ImageGeneration => return ModelTier::ImageGeneration,
113 ModelMode::Chat => {}
114 _ => return ModelTier::Balanced,
115 }
116
117 if let Some(tier) = breakpoints.and_then(|b| b.classify(info)) {
118 return tier;
119 }
120
121 match info.max_input_tokens() {
122 Some(tokens) if tokens <= 32_768 => ModelTier::Fast,
123 Some(tokens) if tokens >= 200_000 => ModelTier::Advanced,
124 _ => ModelTier::Balanced,
125 }
126}
127
128#[cfg(test)]
129mod tests {
130 use super::*;
131 use crate::ModelRegistry;
132
133 #[test]
134 fn bundled_registry_classifies_known_models() {
135 let registry = ModelRegistry::bundled();
136 assert_eq!(registry.tier("gpt-image-1"), ModelTier::ImageGeneration);
137 assert_ne!(registry.tier("claude-opus-4-1"), ModelTier::Fast);
139 assert_ne!(registry.tier("gpt-4o-mini"), ModelTier::Advanced);
141 }
142
143 #[test]
144 fn unknown_model_defaults_to_balanced() {
145 let registry = ModelRegistry::bundled();
146 assert_eq!(
147 registry.tier("definitely-not-a-real-model-xyz"),
148 ModelTier::Balanced
149 );
150 }
151
152 #[test]
153 fn serde_roundtrip_matches_lowercase() {
154 let json = serde_json::to_string(&ModelTier::Advanced).unwrap();
155 assert_eq!(json, "\"advanced\"");
156 let tier: ModelTier = serde_json::from_str("\"fast\"").unwrap();
157 assert_eq!(tier, ModelTier::Fast);
158 }
159}