Skip to main content

vct_core/pricing/
tiers.rs

1//! Per-request context-tier classification support.
2//!
3//! LiteLLM publishes `*_above_Nk_tokens` price tiers whose real billing
4//! semantics are **per request**: a request is promoted to the tier rate only
5//! when its own prompt context exceeds the threshold. The usage aggregator,
6//! however, merges tokens across records, files, and sessions before pricing,
7//! so tier selection at pricing time would compare the threshold against a
8//! cumulative figure and promote a whole month of small requests to the
9//! elevated rate.
10//!
11//! [`TierThresholds`] is a `Send + Sync` snapshot of "model → lowest tier
12//! threshold" derived from a [`ModelPricingMap`](super::ModelPricingMap). The
13//! usage scan hands it to the session parsers, which classify each request as
14//! it is folded and accumulate the above-threshold slice into a separate
15//! `above_tier` bucket that `calculate_cost` bills at the tier rate. Parsers
16//! without the snapshot (the `analysis` paths, offline runs) classify nothing,
17//! which degrades to billing everything at base rates.
18
19use crate::constants::FastHashMap;
20use crate::pricing::normalize_model_name;
21use std::hash::{DefaultHasher, Hash, Hasher};
22
23/// Immutable "model → lowest context-tier threshold (tokens)" snapshot.
24///
25/// Keys are stored both as the LiteLLM key lowercased and as its normalized
26/// form, so the session-log model names (`gpt-5.4`, `azure/gpt-5.5`,
27/// `claude-sonnet-5`) resolve without re-implementing the full pricing match
28/// chain. A model that resolves to no entry simply has no tier.
29#[derive(Debug, Default)]
30pub struct TierThresholds {
31    thresholds: FastHashMap<Box<str>, i64>,
32    fingerprint: u64,
33}
34
35impl TierThresholds {
36    /// Builds the snapshot from `(model key, lowest threshold)` pairs.
37    ///
38    /// On key collisions (e.g. `openai/gpt-5.4` and `azure/gpt-5.4`
39    /// normalizing to the same name) the smallest threshold wins — the
40    /// conservative choice given the tier rate is the higher one.
41    pub(crate) fn from_entries<'a>(entries: impl Iterator<Item = (&'a str, i64)>) -> Self {
42        let mut thresholds: FastHashMap<Box<str>, i64> = FastHashMap::default();
43        let mut insert_min = |key: String, threshold: i64| {
44            thresholds
45                .entry(key.into_boxed_str())
46                .and_modify(|existing| *existing = (*existing).min(threshold))
47                .or_insert(threshold);
48        };
49        for (key, threshold) in entries {
50            insert_min(key.to_lowercase(), threshold);
51            insert_min(normalize_model_name(key), threshold);
52        }
53
54        // Order-independent fingerprint so scan caches can detect a changed
55        // snapshot (daily pricing reload) without hashing map iteration order.
56        let mut fingerprint = thresholds.len() as u64;
57        for (key, threshold) in &thresholds {
58            let mut hasher = DefaultHasher::new();
59            (key, threshold).hash(&mut hasher);
60            fingerprint ^= hasher.finish();
61        }
62
63        Self {
64            thresholds,
65            fingerprint,
66        }
67    }
68
69    /// Lowest tier threshold for `model`, or `None` when the model has no
70    /// context tier (or cannot be resolved).
71    pub fn threshold_for(&self, model: &str) -> Option<i64> {
72        if self.thresholds.is_empty() {
73            return None;
74        }
75        let lowered = model.to_lowercase();
76        if let Some(threshold) = self.thresholds.get(lowered.as_str()) {
77            return Some(*threshold);
78        }
79        self.thresholds
80            .get(normalize_model_name(model).as_str())
81            .copied()
82    }
83
84    /// Whether no model carries a tier (nothing will ever classify).
85    pub fn is_empty(&self) -> bool {
86        self.thresholds.is_empty()
87    }
88
89    /// Stable identity of this snapshot's contents, `0` only when empty.
90    pub fn fingerprint(&self) -> u64 {
91        self.fingerprint
92    }
93}
94
95/// Per-parse memoized classifier over a [`TierThresholds`] snapshot.
96///
97/// Memoizes the per-model resolution (which lowercases/normalizes the name)
98/// so the per-record hot path is one map hit plus an integer comparison.
99#[derive(Debug)]
100pub struct TierClassifier<'a> {
101    thresholds: &'a TierThresholds,
102    memo: FastHashMap<String, Option<i64>>,
103}
104
105impl<'a> TierClassifier<'a> {
106    pub fn new(thresholds: &'a TierThresholds) -> Self {
107        Self {
108            thresholds,
109            memo: FastHashMap::default(),
110        }
111    }
112
113    /// Whether a request for `model` with `request_context` prompt tokens
114    /// (input + cache read + cache creation) is billed at the tier rate.
115    pub fn is_above(&mut self, model: &str, request_context: i64) -> bool {
116        let threshold = match self.memo.get(model) {
117            Some(threshold) => *threshold,
118            None => {
119                let resolved = self.thresholds.threshold_for(model);
120                self.memo.insert(model.to_string(), resolved);
121                resolved
122            }
123        };
124        threshold.is_some_and(|threshold| request_context > threshold)
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    fn snapshot() -> TierThresholds {
133        TierThresholds::from_entries(
134            [("gpt-5.4", 272_000), ("gemini-3.1-pro-preview", 200_000)].into_iter(),
135        )
136    }
137
138    #[test]
139    fn resolves_exact_and_normalized_names() {
140        let tiers = snapshot();
141        assert_eq!(tiers.threshold_for("gpt-5.4"), Some(272_000));
142        assert_eq!(tiers.threshold_for("GPT-5.4"), Some(272_000));
143        assert_eq!(tiers.threshold_for("gpt-4o"), None);
144    }
145
146    #[test]
147    fn collision_keeps_the_smallest_threshold() {
148        let tiers = TierThresholds::from_entries(
149            [("openai/gpt-x", 272_000), ("gpt-x", 200_000)].into_iter(),
150        );
151        assert_eq!(tiers.threshold_for("gpt-x"), Some(200_000));
152    }
153
154    #[test]
155    fn classifier_compares_strictly_above() {
156        let tiers = snapshot();
157        let mut classifier = TierClassifier::new(&tiers);
158        assert!(!classifier.is_above("gpt-5.4", 272_000));
159        assert!(classifier.is_above("gpt-5.4", 272_001));
160        assert!(!classifier.is_above("no-tier-model", i64::MAX));
161        // Memoized second lookup takes the fast path.
162        assert!(classifier.is_above("gpt-5.4", 300_000));
163    }
164
165    #[test]
166    fn fingerprint_is_order_independent_and_content_sensitive() {
167        let a = TierThresholds::from_entries([("m1", 100), ("m2", 200)].into_iter());
168        let b = TierThresholds::from_entries([("m2", 200), ("m1", 100)].into_iter());
169        let c = TierThresholds::from_entries([("m1", 100), ("m2", 300)].into_iter());
170        assert_eq!(a.fingerprint(), b.fingerprint());
171        assert_ne!(a.fingerprint(), c.fingerprint());
172        assert_eq!(TierThresholds::default().fingerprint(), 0);
173    }
174}