cratefield_core/
classifier_routing.rs

1//! Two wrappers over a cheap and an expensive [`Classifier`] (issue #457):
2//! [`ShadowClassifier`] serves the expensive answers and measures the cheap
3//! one alongside, and [`RoutingClassifier`] serves the cheap answers only
4//! for question kinds where that measurement earned it.
5//!
6//! **Thresholds are per adapter.** [`Answer::confidence`] means a different
7//! thing per adapter: a trained classifier's calibrated probability, or a
8//! language model's token probability, which runs high whether or not it
9//! is right (see [`Calibration`]). A 0.8 from one is not a 0.8 from the
10//! other, so one threshold shared across adapters would be wrong for at
11//! least one of them. Each [`Thresholds`] names the [`Calibration`] it was
12//! tuned under, and [`RoutingClassifier::new`] refuses a threshold whose
13//! adapter reports a different one.
14
15use std::collections::BTreeMap;
16use std::sync::Arc;
17use std::task::Poll;
18
19use async_trait::async_trait;
20
21use crate::classifier_agreement::{AgreementLog, AgreementReport, answer_margin, compare};
22use crate::cost::{Accounting, AdapterId, CallRole, QuestionKind};
23use crate::ports::{
24    Answer, Calibration, Classifier, ClassifierError, ClassifierProfile, Question,
25    validate_questions,
26};
27
28/// When the cheap adapter's answer is good enough to serve: at least
29/// `min_confidence` and a top-two margin of at least `min_margin`, under
30/// the [`Calibration`] these numbers were tuned for.
31#[derive(Debug, Clone, Copy, PartialEq)]
32pub struct Thresholds {
33    pub calibration: Calibration,
34    pub min_confidence: f32,
35    pub min_margin: f32,
36}
37
38/// Whether routing happens at all.
39#[derive(Debug, Clone, PartialEq, Eq, Default)]
40pub enum RoutingMode {
41    /// Every question to the expensive adapter — the default.
42    #[default]
43    Off,
44    /// Every question to the named adapter, no escalation.
45    Pinned(AdapterId),
46    /// Per kind, on the policy's evidence.
47    Measured,
48}
49
50/// How [`RoutingClassifier`] decides. Defaults to [`RoutingMode::Off`]:
51/// a cheaper adapter is earned by measurement, never assumed.
52#[derive(Debug, Clone, PartialEq)]
53pub struct RoutingPolicy {
54    mode: RoutingMode,
55    evidence: Option<AgreementReport>,
56    thresholds: BTreeMap<AdapterId, Thresholds>,
57    min_questions: usize,
58    max_disagreement_rate: f64,
59}
60
61impl Default for RoutingPolicy {
62    fn default() -> Self {
63        Self::off()
64    }
65}
66
67impl RoutingPolicy {
68    /// Everything to the expensive adapter.
69    #[must_use]
70    pub fn off() -> Self {
71        Self {
72            mode: RoutingMode::Off,
73            evidence: None,
74            thresholds: BTreeMap::new(),
75            min_questions: 20,
76            max_disagreement_rate: 0.05,
77        }
78    }
79
80    /// Everything to `adapter`, which must be one the router holds.
81    #[must_use]
82    pub fn pinned(adapter: AdapterId) -> Self {
83        Self {
84            mode: RoutingMode::Pinned(adapter),
85            ..Self::off()
86        }
87    }
88
89    /// Route per kind on `evidence`, which must compare the router's two
90    /// adapters.
91    #[must_use]
92    pub fn measured(evidence: AgreementReport) -> Self {
93        Self {
94            mode: RoutingMode::Measured,
95            evidence: Some(evidence),
96            ..Self::off()
97        }
98    }
99
100    /// The thresholds for `adapter`. Without thresholds for the cheap
101    /// adapter nothing is routed to it.
102    #[must_use]
103    pub fn thresholds(mut self, adapter: AdapterId, thresholds: Thresholds) -> Self {
104        self.thresholds.insert(adapter, thresholds);
105        self
106    }
107
108    /// How many measured questions a kind needs, at the thresholds, before
109    /// the cheap adapter may answer it. Default 20.
110    #[must_use]
111    pub const fn min_questions(mut self, n: usize) -> Self {
112        self.min_questions = n;
113        self
114    }
115
116    /// The most disagreement a kind may show, at the thresholds, and still
117    /// go to the cheap adapter; equal to the cap passes. Default 0.05.
118    #[must_use]
119    pub const fn max_disagreement_rate(mut self, rate: f64) -> Self {
120        self.max_disagreement_rate = rate;
121        self
122    }
123
124    /// Why a kind is or is not trusted to the cheap adapter; `None` means
125    /// trusted.
126    fn distrust(
127        &self,
128        kind: &QuestionKind,
129        thresholds: Option<&Thresholds>,
130    ) -> Option<RouteReason> {
131        let Some(t) = thresholds else {
132            return Some(RouteReason::NoCalibration);
133        };
134        let Some(k) = self.evidence.as_ref().and_then(|e| e.kinds.get(kind)) else {
135            return Some(RouteReason::NoMeasurement);
136        };
137        let at = k.agreement_at(t.min_confidence, t.min_margin);
138        if at.questions < self.min_questions {
139            return Some(RouteReason::NotMeasuredEnough);
140        }
141        match at.disagreement_rate() {
142            Some(rate) if rate <= self.max_disagreement_rate => None,
143            _ => Some(RouteReason::DisagreementTooHigh),
144        }
145    }
146}
147
148/// Refusals from [`RoutingClassifier::new`]: a policy that cannot mean what
149/// it says for these adapters is rejected at wiring time, not at the first
150/// question.
151#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
152pub enum RoutingError {
153    /// Thresholds tuned under one calibration, for an adapter that reports
154    /// another: its confidence does not mean what the numbers assume.
155    #[error(
156        "thresholds for {adapter} declare {declared:?} calibration, the adapter reports {reported:?}"
157    )]
158    CalibrationMismatch {
159        adapter: AdapterId,
160        declared: Calibration,
161        reported: Calibration,
162    },
163    /// A pin or a threshold names an adapter the router does not hold.
164    #[error("the policy names {0}, which this router does not hold")]
165    UnknownAdapter(AdapterId),
166    /// Measured evidence about a different pair of adapters.
167    #[error("the evidence compares {cheap} with {expensive}, not this router's adapters")]
168    EvidenceForOtherAdapters {
169        cheap: AdapterId,
170        expensive: AdapterId,
171    },
172}
173
174/// Why a question went where it went.
175#[derive(Debug, Clone, Copy, PartialEq, Eq)]
176pub enum RouteReason {
177    /// [`RoutingMode::Off`].
178    RoutingOff,
179    /// [`RoutingMode::Pinned`].
180    Pinned,
181    /// No [`Thresholds`] for the cheap adapter.
182    NoCalibration,
183    /// The evidence has no such kind.
184    NoMeasurement,
185    /// Fewer measured questions than `min_questions` at the thresholds.
186    NotMeasuredEnough,
187    /// Measured disagreement above `max_disagreement_rate`.
188    DisagreementTooHigh,
189    /// Trusted, and the cheap answer cleared the thresholds.
190    CheapAccepted,
191    /// Trusted, but the cheap answer's confidence was below the minimum.
192    EscalatedLowConfidence,
193    /// Trusted, but the cheap answer's top-two margin was below the minimum.
194    EscalatedNarrowMargin,
195    /// Trusted, but the cheap adapter failed or left the question out.
196    EscalatedCheapFailed,
197}
198
199/// One answer, with who gave it and why.
200#[derive(Debug, Clone, PartialEq)]
201pub struct Routed {
202    pub adapter: AdapterId,
203    pub reason: RouteReason,
204    pub answer: Answer,
205}
206
207/// Serves the expensive adapter's answers, unchanged, and asks the cheap
208/// one the same questions to measure agreement. The cheap adapter cannot
209/// change or fail a result: its errors are logged, recorded and counted as
210/// skipped.
211///
212/// Both adapters are asked concurrently, on the caller's task, and the call
213/// returns when both have: its latency is the slower of the two, not the
214/// sum. A cheap call that stalls still holds the shadowed call, so the
215/// cheap adapter must enforce its own deadline. When the expensive adapter
216/// fails, its error is the call's; the cheap call is still recorded.
217///
218/// [`profile`](Classifier::profile) is the expensive adapter's: every
219/// answer returned is its.
220pub struct ShadowClassifier {
221    cheap: (AdapterId, Arc<dyn Classifier>),
222    expensive: (AdapterId, Arc<dyn Classifier>),
223    accounting: Accounting,
224    log: AgreementLog,
225}
226
227impl ShadowClassifier {
228    /// A shadow measuring `cheap` against `expensive`.
229    #[must_use]
230    pub fn new(
231        cheap: (AdapterId, Arc<dyn Classifier>),
232        expensive: (AdapterId, Arc<dyn Classifier>),
233        accounting: Accounting,
234    ) -> Self {
235        let log = AgreementLog::new(cheap.0.clone(), expensive.0.clone());
236        Self {
237            cheap,
238            expensive,
239            accounting,
240            log,
241        }
242    }
243
244    /// The agreement measured so far — ungraded, since live traffic has no
245    /// ground truth.
246    #[must_use]
247    pub fn report(&self) -> AgreementReport {
248        self.log.report()
249    }
250}
251
252/// Drives both futures on the caller's task until both finish — core
253/// cannot spawn — so the pair takes as long as the slower of the two.
254async fn join<A: Future, B: Future>(a: A, b: B) -> (A::Output, B::Output) {
255    let (mut a, mut b) = (std::pin::pin!(a), std::pin::pin!(b));
256    let (mut done_a, mut done_b) = (None, None);
257    std::future::poll_fn(|cx| {
258        if done_a.is_none()
259            && let Poll::Ready(out) = a.as_mut().poll(cx)
260        {
261            done_a = Some(out);
262        }
263        if done_b.is_none()
264            && let Poll::Ready(out) = b.as_mut().poll(cx)
265        {
266            done_b = Some(out);
267        }
268        match (done_a.take(), done_b.take()) {
269            (Some(a), Some(b)) => Poll::Ready((a, b)),
270            (a, b) => {
271                (done_a, done_b) = (a, b);
272                Poll::Pending
273            }
274        }
275    })
276    .await
277}
278
279#[async_trait]
280impl Classifier for ShadowClassifier {
281    fn profile(&self) -> ClassifierProfile {
282        self.expensive.1.profile()
283    }
284
285    async fn ask(
286        &self,
287        state: &str,
288        questions: &BTreeMap<String, Question>,
289    ) -> Result<BTreeMap<String, Answer>, ClassifierError> {
290        let (dear_id, dear) = &self.expensive;
291        let (cheap_id, cheap) = &self.cheap;
292        let serve = async {
293            let answers = dear.ask(state, questions).await;
294            let answered = answers.as_ref().ok();
295            let profile = dear.profile();
296            self.accounting.record(
297                dear_id,
298                &profile,
299                state,
300                questions,
301                answered,
302                CallRole::Served,
303            );
304            answers
305        };
306        let measure = async {
307            let answers = cheap.ask(state, questions).await;
308            let answered = answers.as_ref().ok();
309            let profile = cheap.profile();
310            self.accounting.record(
311                cheap_id,
312                &profile,
313                state,
314                questions,
315                answered,
316                CallRole::Shadow,
317            );
318            answers
319        };
320        let (served, shadow) = join(serve, measure).await;
321        if let Err(err) = &shadow {
322            tracing::debug!(adapter = %cheap_id, error = %err, "shadow classifier failed");
323        }
324        let served = match served {
325            Ok(answers) => answers,
326            Err(err) => {
327                compare(&self.log, questions, None, None, &BTreeMap::new(), None);
328                return Err(err);
329            }
330        };
331        compare(
332            &self.log,
333            questions,
334            shadow.as_ref().ok(),
335            Some(&served),
336            &BTreeMap::new(),
337            None,
338        );
339        Ok(served)
340    }
341}
342
343/// Sends each question to the cheap adapter when its kind has earned it
344/// and to the expensive adapter otherwise; see [`RoutingPolicy`].
345///
346/// Under [`RoutingMode::Measured`], one call:
347///
348/// 1. asks the cheap adapter the trusted kinds, in one batch;
349/// 2. accepts each cheap answer that clears its [`Thresholds`];
350/// 3. asks the expensive adapter, in one second batch, the untrusted kinds
351///    plus every cheap answer that did not clear — or all the trusted ones
352///    if the cheap adapter failed;
353/// 4. merges both under the ids asked.
354///
355/// An empty or malformed question set is rejected first, as the adapters
356/// reject it. An expensive failure is the call's error.
357/// [`ask_routed`](Self::ask_routed) says who answered each question and
358/// why; [`ask`](Classifier::ask) returns the answers alone.
359///
360/// **Measured answers mix confidence scales.** In measured mode one call's
361/// answers can come from both adapters, and each answer's
362/// [`Answer::confidence`] is on its own adapter's scale. A caller that
363/// thresholds confidence must use [`ask_routed`](Self::ask_routed) and
364/// threshold per [`Routed::adapter`], not the merged map `ask` returns.
365///
366/// [`profile`](Classifier::profile): off, the expensive adapter's; pinned,
367/// the pinned adapter's; measured, the expensive adapter's
368/// [`Calibration`] — a cheap answer is only served where it was measured
369/// to agree with the expensive one — and the smaller of the two
370/// `max_state_chars`, the least state any answer may have seen. In
371/// measured mode that calibration describes the expensive answers only;
372/// see above.
373pub struct RoutingClassifier {
374    cheap: (AdapterId, Arc<dyn Classifier>),
375    expensive: (AdapterId, Arc<dyn Classifier>),
376    policy: RoutingPolicy,
377    accounting: Accounting,
378}
379
380impl RoutingClassifier {
381    /// A router over `cheap` and `expensive`.
382    ///
383    /// # Errors
384    ///
385    /// [`RoutingError`] when the policy pins or thresholds an adapter the
386    /// router does not hold, declares a calibration an adapter does not
387    /// report, or (measured) carries evidence about other adapters.
388    pub fn new(
389        cheap: (AdapterId, Arc<dyn Classifier>),
390        expensive: (AdapterId, Arc<dyn Classifier>),
391        policy: RoutingPolicy,
392        accounting: Accounting,
393    ) -> Result<Self, RoutingError> {
394        let held = |id: &AdapterId| {
395            [&cheap, &expensive]
396                .into_iter()
397                .find(|(held, _)| held == id)
398        };
399        if let RoutingMode::Pinned(id) = &policy.mode
400            && held(id).is_none()
401        {
402            return Err(RoutingError::UnknownAdapter(id.clone()));
403        }
404        for (id, t) in &policy.thresholds {
405            let (_, adapter) = held(id).ok_or_else(|| RoutingError::UnknownAdapter(id.clone()))?;
406            let reported = adapter.profile().calibration;
407            if reported != t.calibration {
408                return Err(RoutingError::CalibrationMismatch {
409                    adapter: id.clone(),
410                    declared: t.calibration,
411                    reported,
412                });
413            }
414        }
415        if policy.mode == RoutingMode::Measured
416            && let Some(e) = &policy.evidence
417            && (e.cheap != cheap.0 || e.expensive != expensive.0)
418        {
419            return Err(RoutingError::EvidenceForOtherAdapters {
420                cheap: e.cheap.clone(),
421                expensive: e.expensive.clone(),
422            });
423        }
424        Ok(Self {
425            cheap,
426            expensive,
427            policy,
428            accounting,
429        })
430    }
431
432    /// Asks `questions` and says, per answer, which adapter gave it and why.
433    ///
434    /// # Errors
435    ///
436    /// The serving adapter's [`ClassifierError`]: the expensive one's, the
437    /// pinned one's, or the expensive escalation's. A cheap failure under
438    /// [`RoutingMode::Measured`] escalates instead.
439    pub async fn ask_routed(
440        &self,
441        state: &str,
442        questions: &BTreeMap<String, Question>,
443    ) -> Result<BTreeMap<String, Routed>, ClassifierError> {
444        let (target, reason) = match &self.policy.mode {
445            RoutingMode::Measured => return self.ask_measured(state, questions).await,
446            RoutingMode::Off => (&self.expensive, RouteReason::RoutingOff),
447            RoutingMode::Pinned(id) if *id == self.cheap.0 => (&self.cheap, RouteReason::Pinned),
448            RoutingMode::Pinned(_) => (&self.expensive, RouteReason::Pinned),
449        };
450        let answers = target.1.ask(state, questions).await;
451        self.accounting.record(
452            &target.0,
453            &target.1.profile(),
454            state,
455            questions,
456            answers.as_ref().ok(),
457            CallRole::Served,
458        );
459        Ok(answers?
460            .into_iter()
461            .map(|(id, answer)| {
462                (
463                    id,
464                    Routed {
465                        adapter: target.0.clone(),
466                        reason,
467                        answer,
468                    },
469                )
470            })
471            .collect())
472    }
473
474    async fn ask_measured(
475        &self,
476        state: &str,
477        questions: &BTreeMap<String, Question>,
478    ) -> Result<BTreeMap<String, Routed>, ClassifierError> {
479        validate_questions(questions)?;
480        let (cheap_id, cheap) = &self.cheap;
481        let (dear_id, dear) = &self.expensive;
482        let thresholds = self.policy.thresholds.get(cheap_id);
483
484        let mut to_cheap = BTreeMap::new();
485        let mut to_dear = BTreeMap::new();
486        let mut reasons = BTreeMap::new();
487        for (id, question) in questions {
488            let kind = QuestionKind::new(id.as_str());
489            if let Some(reason) = self.policy.distrust(&kind, thresholds) {
490                reasons.insert(id.clone(), reason);
491                to_dear.insert(id.clone(), question.clone());
492            } else {
493                to_cheap.insert(id.clone(), question.clone());
494            }
495        }
496
497        let mut routed = BTreeMap::new();
498        if let Some(t) = thresholds.filter(|_| !to_cheap.is_empty()) {
499            let answers = cheap.ask(state, &to_cheap).await;
500            if let Err(err) = &answers {
501                tracing::debug!(adapter = %cheap_id, error = %err, "cheap classifier failed; escalating");
502            }
503            for (id, question) in &to_cheap {
504                let answer = answers.as_ref().ok().and_then(|a| a.get(id));
505                match (judge(answer, t), answer) {
506                    (RouteReason::CheapAccepted, Some(answer)) => {
507                        let reason = RouteReason::CheapAccepted;
508                        let adapter = cheap_id.clone();
509                        let answer = answer.clone();
510                        routed.insert(
511                            id.clone(),
512                            Routed {
513                                adapter,
514                                reason,
515                                answer,
516                            },
517                        );
518                    }
519                    (reason, _) => {
520                        reasons.insert(id.clone(), reason);
521                        to_dear.insert(id.clone(), question.clone());
522                    }
523                }
524            }
525            // Recorded as soon as it returns, so a caller that drops the
526            // future during the escalation does not lose what it cost.
527            let role = if routed.is_empty() {
528                CallRole::Discarded
529            } else {
530                CallRole::Served
531            };
532            let answered = answers.as_ref().ok();
533            let profile = cheap.profile();
534            self.accounting
535                .record(cheap_id, &profile, state, &to_cheap, answered, role);
536        }
537
538        if !to_dear.is_empty() {
539            let answers = dear.ask(state, &to_dear).await;
540            self.accounting.record(
541                dear_id,
542                &dear.profile(),
543                state,
544                &to_dear,
545                answers.as_ref().ok(),
546                CallRole::Served,
547            );
548            for (id, answer) in answers? {
549                if let Some(&reason) = reasons.get(&id) {
550                    let adapter = dear_id.clone();
551                    routed.insert(
552                        id,
553                        Routed {
554                            adapter,
555                            reason,
556                            answer,
557                        },
558                    );
559                }
560            }
561        }
562        Ok(routed)
563    }
564}
565
566/// Whether a trusted kind's cheap answer is served, or why it escalates.
567/// The comparisons are in the positive, so a NaN never clears.
568fn judge(answer: Option<&Answer>, t: &Thresholds) -> RouteReason {
569    let Some(answer) = answer else {
570        return RouteReason::EscalatedCheapFailed;
571    };
572    let confident = answer.confidence >= t.min_confidence;
573    let decisive = answer_margin(answer) >= t.min_margin;
574    if !confident {
575        RouteReason::EscalatedLowConfidence
576    } else if !decisive {
577        RouteReason::EscalatedNarrowMargin
578    } else {
579        RouteReason::CheapAccepted
580    }
581}
582
583#[async_trait]
584impl Classifier for RoutingClassifier {
585    /// See [`RoutingClassifier`]: in measured mode the calibration is the
586    /// expensive adapter's, and cheap answers are on the cheap adapter's
587    /// scale.
588    fn profile(&self) -> ClassifierProfile {
589        match &self.policy.mode {
590            RoutingMode::Pinned(id) if *id == self.cheap.0 => self.cheap.1.profile(),
591            RoutingMode::Off | RoutingMode::Pinned(_) => self.expensive.1.profile(),
592            RoutingMode::Measured => {
593                let (cheap, dear) = (self.cheap.1.profile(), self.expensive.1.profile());
594                ClassifierProfile::new(
595                    dear.calibration,
596                    cheap.max_state_chars.min(dear.max_state_chars),
597                )
598            }
599        }
600    }
601
602    /// The answers of [`ask_routed`](RoutingClassifier::ask_routed) without
603    /// who gave them. In measured mode they mix adapters' confidence
604    /// scales: threshold confidence through `ask_routed`, per adapter.
605    async fn ask(
606        &self,
607        state: &str,
608        questions: &BTreeMap<String, Question>,
609    ) -> Result<BTreeMap<String, Answer>, ClassifierError> {
610        let routed = self.ask_routed(state, questions).await?;
611        Ok(routed.into_iter().map(|(id, r)| (id, r.answer)).collect())
612    }
613}
614
615#[cfg(test)]
616mod tests {
617    use super::*;
618
619    #[test]
620    fn a_bare_noul_probability_escalates_on_its_unreported_complement() {
621        // {"true": 0.55} leaves 0.45 unreported: the margin is 0.10, not
622        // 0.55, so a 0.3 minimum margin escalates it.
623        let answer = Answer::noul(true, BTreeMap::from([("true".to_owned(), 0.55)]));
624        let t = Thresholds {
625            calibration: Calibration::Classifier,
626            min_confidence: 0.5,
627            min_margin: 0.3,
628        };
629        assert_eq!(judge(Some(&answer), &t), RouteReason::EscalatedNarrowMargin);
630        let lenient = Thresholds {
631            min_margin: 0.05,
632            ..t
633        };
634        assert_eq!(judge(Some(&answer), &lenient), RouteReason::CheapAccepted);
635        assert_eq!(judge(None, &t), RouteReason::EscalatedCheapFailed);
636    }
637}
cratefield_core/classifier_routing.rs

cratefield_core/
classifier_routing.rs
