tokenmiser_quality/
scheduler.rs

1//! Shadow scheduler: samples a fraction of traffic and runs the frontier
2//! comparison + judge in the background.
3//!
4//! The proxy calls `ShadowScheduler::maybe_enqueue(...)` synchronously
5//! after returning the cheap-model response to the user. If the sample
6//! roll succeeds, the scheduler fires a `tokio::spawn` that runs the
7//! frontier call + judge + aggregator update. User-facing latency is
8//! unaffected.
9
10use std::sync::Arc;
11
12use tokenmiser_providers::{ChatResponse, ProviderRegistry};
13use tracing::warn;
14
15use crate::{
16    judge::judge, log_sample, ShadowConfig, ShadowEnqueue, ShadowSample, WinRateAggregator,
17};
18
19pub struct ShadowScheduler {
20    cfg: ShadowConfig,
21    registry: Arc<ProviderRegistry>,
22    aggregator: Arc<WinRateAggregator>,
23}
24
25impl ShadowScheduler {
26    pub fn new(
27        cfg: ShadowConfig,
28        registry: Arc<ProviderRegistry>,
29        aggregator: Arc<WinRateAggregator>,
30    ) -> Arc<Self> {
31        Arc::new(Self {
32            cfg,
33            registry,
34            aggregator,
35        })
36    }
37
38    pub fn aggregator(&self) -> &Arc<WinRateAggregator> {
39        &self.aggregator
40    }
41
42    /// Roll the sample-rate dice. If true, spawn a background shadow
43    /// comparison.
44    pub fn maybe_enqueue(self: &Arc<Self>, e: ShadowEnqueue) {
45        if !roll(self.cfg.sample_rate) {
46            return;
47        }
48        let me = Arc::clone(self);
49        tokio::spawn(async move {
50            if let Err(e) = me.run(e).await {
51                warn!(error = %e, "shadow comparison failed");
52            }
53        });
54    }
55
56    async fn run(self: Arc<Self>, e: ShadowEnqueue) -> anyhow::Result<()> {
57        let cheap_text = extract_text(&e.cheap_response);
58
59        // Call the frontier model with the same request.
60        let mut frontier_req = e.req.clone();
61        frontier_req.model = self.cfg.frontier_model.clone();
62        let (frontier_provider, frontier_real) = self
63            .registry
64            .resolve(&self.cfg.frontier_model)
65            .map_err(|err| anyhow::anyhow!("frontier resolve: {err}"))?;
66        frontier_req.model = frontier_real.clone();
67        let frontier_resp = frontier_provider
68            .complete(&frontier_req)
69            .await
70            .map_err(|err| anyhow::anyhow!("frontier call: {err}"))?;
71        let frontier_text = extract_text(&frontier_resp);
72
73        // User prompt = first user message.
74        let user_prompt = e
75            .req
76            .messages
77            .iter()
78            .find(|m| m.role == "user")
79            .and_then(|m| match &m.content {
80                serde_json::Value::String(s) => Some(s.clone()),
81                _ => None,
82            })
83            .unwrap_or_default();
84
85        let verdict = judge(
86            &self.registry,
87            &self.cfg.judge_model,
88            &user_prompt,
89            &cheap_text,
90            &frontier_text,
91        )
92        .await?;
93
94        let sample = ShadowSample {
95            segment: e.segment,
96            cheap_model: e.cheap_model,
97            frontier_model: self.cfg.frontier_model.clone(),
98            verdict,
99        };
100        log_sample(&sample);
101        self.aggregator.record(&sample);
102
103        Ok(())
104    }
105}
106
107fn roll(rate: f32) -> bool {
108    if rate <= 0.0 {
109        return false;
110    }
111    if rate >= 1.0 {
112        return true;
113    }
114    // Deterministic-ish using nanos; good enough for sampling. v0.9 can
115    // swap in `rand::random::<f32>()` if the bias profile matters.
116    let nanos = std::time::SystemTime::now()
117        .duration_since(std::time::UNIX_EPOCH)
118        .map(|d| d.subsec_nanos())
119        .unwrap_or(0);
120    let r = (nanos as f32) / (u32::MAX as f32);
121    r < rate
122}
123
124fn extract_text(resp: &ChatResponse) -> String {
125    resp.choices
126        .first()
127        .map(|c| match &c.message.content {
128            serde_json::Value::String(s) => s.clone(),
129            other => other.to_string(),
130        })
131        .unwrap_or_default()
132}
133
134/// Re-export so tokenmiser-proxy can pull just `tokenmiser_quality::Arc` if
135/// it wants symmetry with internal helpers.
136#[allow(dead_code)]
137fn _arc_typecheck(_: Arc<ShadowScheduler>) {}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn roll_extremes() {
145        assert!(roll(1.0));
146        assert!(!roll(0.0));
147    }
148}
tokenmiser_quality/scheduler.rs

tokenmiser_quality/
scheduler.rs
