tokenmiser_quality/
lib.rs

1//! Shadow A/B with an LLM judge.
2//!
3//! A sampled fraction of routed traffic is replayed against the frontier model
4//! in the background, a judge scores the pair, and per-segment win rates are
5//! tallied so a regressed segment can raise an alert.
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use parking_lot::Mutex;
11use serde::{Deserialize, Serialize};
12use tokenmiser_providers::{ChatMessage, ChatRequest, ChatResponse};
13use tracing::{info, warn};
14
15pub mod judge;
16pub mod scheduler;
17
18pub use judge::JudgeVerdict;
19pub use scheduler::ShadowScheduler;
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct ShadowConfig {
23    /// Fraction of traffic to shadow-test, in [0.0, 1.0].
24    pub sample_rate: f32,
25    /// Which model to use as the frontier baseline in shadow tests.
26    pub frontier_model: String,
27    /// Model used as the judge.
28    pub judge_model: String,
29    /// Cheap-model win rate below this over `min_samples_per_segment` marks
30    /// the segment regressed.
31    pub auto_gate_floor: f32,
32    pub min_samples_per_segment: u32,
33}
34
35impl Default for ShadowConfig {
36    fn default() -> Self {
37        Self {
38            sample_rate: 0.01,
39            frontier_model: "claude-opus-4-7".into(),
40            judge_model: "claude-sonnet-4-6".into(),
41            auto_gate_floor: 0.45,
42            min_samples_per_segment: 30,
43        }
44    }
45}
46
47/// One completed shadow comparison.
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct ShadowSample {
50    pub segment: String,
51    pub cheap_model: String,
52    pub frontier_model: String,
53    pub verdict: JudgeVerdict,
54}
55
56/// Win-rate aggregator: per-segment running tallies, plus regression flag.
57pub struct WinRateAggregator {
58    by_segment: Mutex<HashMap<String, SegmentStats>>,
59    floor: f32,
60    min_samples: u32,
61}
62
63#[derive(Debug, Default, Clone, Serialize, Deserialize)]
64pub struct SegmentStats {
65    pub cheap_wins: u32,
66    pub frontier_wins: u32,
67    pub ties: u32,
68    pub regressed: bool,
69}
70
71impl SegmentStats {
72    pub fn cheap_win_rate(&self) -> f32 {
73        let total = self.cheap_wins + self.frontier_wins;
74        if total == 0 {
75            return 1.0;
76        }
77        self.cheap_wins as f32 / total as f32
78    }
79    pub fn total(&self) -> u32 {
80        self.cheap_wins + self.frontier_wins + self.ties
81    }
82}
83
84impl WinRateAggregator {
85    pub fn new(cfg: &ShadowConfig) -> Arc<Self> {
86        Arc::new(Self {
87            by_segment: Mutex::new(HashMap::new()),
88            floor: cfg.auto_gate_floor,
89            min_samples: cfg.min_samples_per_segment,
90        })
91    }
92
93    pub fn record(&self, sample: &ShadowSample) {
94        let mut map = self.by_segment.lock();
95        let s = map.entry(sample.segment.clone()).or_default();
96        match sample.verdict {
97            JudgeVerdict::A => s.cheap_wins += 1,
98            JudgeVerdict::B => s.frontier_wins += 1,
99            JudgeVerdict::Tie => s.ties += 1,
100        }
101        if s.total() >= self.min_samples && s.cheap_win_rate() < self.floor && !s.regressed {
102            s.regressed = true;
103            warn!(
104                segment = %sample.segment,
105                cheap_win_rate = s.cheap_win_rate(),
106                floor = self.floor,
107                samples = s.total(),
108                "AUTO-GATE: segment regressed; rerouting to frontier"
109            );
110        }
111    }
112
113    pub fn snapshot(&self) -> HashMap<String, SegmentStats> {
114        self.by_segment.lock().clone()
115    }
116
117    pub fn regressed_segments(&self) -> Vec<String> {
118        self.by_segment
119            .lock()
120            .iter()
121            .filter(|(_, s)| s.regressed)
122            .map(|(k, _)| k.clone())
123            .collect()
124    }
125}
126
127/// Bucket a request into a coarse aggregation segment, keyed on the opening
128/// words of its first user message.
129pub fn segment_of(req: &ChatRequest) -> String {
130    let user_text = req
131        .messages
132        .iter()
133        .find(|m: &&ChatMessage| m.role == "user")
134        .and_then(|m| match &m.content {
135            serde_json::Value::String(s) => Some(s.clone()),
136            _ => None,
137        })
138        .unwrap_or_else(|| "<empty>".into());
139    user_text
140        .split_whitespace()
141        .take(6)
142        .collect::<Vec<_>>()
143        .join(" ")
144        .to_lowercase()
145}
146
147/// Glue type the proxy uses to enqueue shadow work without blocking.
148pub struct ShadowEnqueue {
149    pub req: ChatRequest,
150    pub cheap_response: ChatResponse,
151    pub cheap_model: String,
152    pub segment: String,
153}
154
155impl ShadowEnqueue {
156    pub fn from_request(
157        req: &ChatRequest,
158        cheap_response: &ChatResponse,
159        cheap_model: &str,
160    ) -> Self {
161        Self {
162            req: req.clone(),
163            cheap_response: cheap_response.clone(),
164            cheap_model: cheap_model.to_string(),
165            segment: segment_of(req),
166        }
167    }
168}
169
170// Surfaced so consumers need no second import.
171pub use tokenmiser_providers::ProviderRegistry as Registry;
172
173/// Log a completed shadow sample.
174pub fn log_sample(sample: &ShadowSample) {
175    info!(
176        segment = %sample.segment,
177        cheap = %sample.cheap_model,
178        frontier = %sample.frontier_model,
179        verdict = ?sample.verdict,
180        "shadow_sample"
181    );
182}
183
184// Re-exported so callers can build the Arc without an extra import.
185pub use std::sync::Arc as ArcReexport;
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn segment_buckets_on_first_words() {
193        let req = ChatRequest {
194            model: "auto".into(),
195            messages: vec![ChatMessage {
196                role: "user".into(),
197                content: serde_json::Value::String(
198                    "What is the capital of france please answer briefly".into(),
199                ),
200                extra: Default::default(),
201            }],
202            temperature: None,
203            max_tokens: None,
204            top_p: None,
205            stream: None,
206            extra: Default::default(),
207        };
208        assert_eq!(segment_of(&req), "what is the capital of france");
209    }
210
211    #[test]
212    fn aggregator_auto_gates_on_low_win_rate() {
213        let agg = WinRateAggregator::new(&ShadowConfig {
214            min_samples_per_segment: 4,
215            auto_gate_floor: 0.45,
216            ..Default::default()
217        });
218        // Cheap win rate of 0.25 is below the 0.45 floor.
219        for v in [
220            JudgeVerdict::A,
221            JudgeVerdict::B,
222            JudgeVerdict::B,
223            JudgeVerdict::B,
224        ] {
225            agg.record(&ShadowSample {
226                segment: "what is".into(),
227                cheap_model: "cheap".into(),
228                frontier_model: "frontier".into(),
229                verdict: v,
230            });
231        }
232        assert_eq!(agg.regressed_segments(), vec!["what is".to_string()]);
233    }
234
235    #[test]
236    fn aggregator_does_not_gate_below_min_samples() {
237        let agg = WinRateAggregator::new(&ShadowConfig {
238            min_samples_per_segment: 100,
239            ..Default::default()
240        });
241        agg.record(&ShadowSample {
242            segment: "x".into(),
243            cheap_model: "c".into(),
244            frontier_model: "f".into(),
245            verdict: JudgeVerdict::B,
246        });
247        assert!(agg.regressed_segments().is_empty());
248    }
249}
tokenmiser_quality/lib.rs

tokenmiser_quality/
lib.rs
