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::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        // The most expensive chat model lands in the top price tertile,
139        // the cheapest in the bottom one.
140        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        // 128k context sits between the Fast and Advanced context bounds.
148        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}