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::registry::tests::fixture;
132
133 #[test]
134 fn fixture_registry_classifies_by_mode_and_price() {
135 let registry = fixture();
136 assert_eq!(registry.tier("fixture-image"), ModelTier::ImageGeneration);
137 assert_eq!(registry.tier("fixture-embedding"), ModelTier::Balanced);
138 assert_eq!(registry.tier("fixture-flagship"), ModelTier::Advanced);
141 assert_eq!(registry.tier("fixture-gpt-mini"), ModelTier::Fast);
142 }
143
144 #[test]
145 fn unpriced_chat_model_falls_back_to_context_window() {
146 let registry = fixture();
147 assert_eq!(
149 registry.tier("github_copilot/fixture-hosted"),
150 ModelTier::Balanced
151 );
152 }
153
154 #[test]
155 fn unknown_model_defaults_to_balanced() {
156 assert_eq!(
157 fixture().tier("definitely-not-a-real-model-xyz"),
158 ModelTier::Balanced
159 );
160 }
161
162 #[test]
163 fn serde_roundtrip_matches_lowercase() {
164 let json = serde_json::to_string(&ModelTier::Advanced).unwrap();
165 assert_eq!(json, "\"advanced\"");
166 let tier: ModelTier = serde_json::from_str("\"fast\"").unwrap();
167 assert_eq!(tier, ModelTier::Fast);
168 }
169}