Skip to main content

aither_models/
tier.rs

1//! Model tier classification inferred from registry metadata.
2//!
3//! Tiers group models by capability/cost so applications can offer
4//! "advanced / balanced / fast" selection without hardcoding model lists.
5//! Classification uses the model mode, the output-price distribution of all
6//! chat models in the registry, and the context window as fallback signal.
7
8use serde::{Deserialize, Serialize};
9
10use crate::types::{ModelEntry, ModelMode};
11
12/// Capability/cost tier of a model.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default, Hash)]
14#[serde(rename_all = "lowercase")]
15pub enum ModelTier {
16    /// Flagship models - highest capability
17    Advanced,
18    /// Balanced models - good capability/cost ratio
19    #[default]
20    Balanced,
21    /// Fast models - optimized for speed
22    Fast,
23    /// Image generation models
24    ImageGeneration,
25}
26
27impl ModelTier {
28    /// All tiers for iteration.
29    pub const ALL: &'static [Self] = &[
30        Self::Advanced,
31        Self::Balanced,
32        Self::Fast,
33        Self::ImageGeneration,
34    ];
35
36    /// Human-readable tier name.
37    #[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    /// Short machine-friendly tier name.
48    #[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/// Output-price tertile breakpoints over a registry's chat models.
60///
61/// Computed once per registry: prices at the lower and upper third of the
62/// distribution. Chat models priced at or below the lower breakpoint classify
63/// as [`ModelTier::Fast`], at or above the upper as [`ModelTier::Advanced`].
64#[derive(Debug, Clone, Copy)]
65pub struct ChatPriceBreakpoints {
66    fast_ceiling: f64,
67    advanced_floor: f64,
68}
69
70impl ChatPriceBreakpoints {
71    /// Derive breakpoints from the chat-model price distribution.
72    ///
73    /// Returns `None` when fewer than three priced chat models exist, in
74    /// which case classification falls back to context-window heuristics.
75    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
109/// Classify an entry given precomputed chat-price breakpoints.
110pub 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        // Flagship reasoning models must not classify as Fast.
138        assert_ne!(registry.tier("claude-opus-4-1"), ModelTier::Fast);
139        // Small/cheap models must not classify as Advanced.
140        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}