firstpass_core/
predictor.rs

1//! Per-query gate-pass predictor (ADR 0008, Phase 2).
2//!
3//! A hand-rolled **online logistic regression** that estimates `P(gate-pass | rung, query
4//! features)` — the per-query, per-rung success probability the start-rung bandit's coarse
5//! context buckets can't express. It is trained incrementally from the deployment's own
6//! receipts (each attempt is a labeled example) and, in this phase, its prediction is recorded
7//! on the receipt in **shadow** — it never changes routing. Whether the prediction is good
8//! enough to *act* on is decided later, offline, via `firstpass predictor-eval` (AUC / Brier).
9//!
10//! No ML/linalg dependency — the model is a weight vector and two closed-form updates. The
11//! feature encoding is fixed-length and deterministic so predictions are reproducible and the
12//! receipt is stable.
13
14use crate::features::{Features, TaskKind};
15
16/// Number of `TaskKind` variants (one-hot width). Keep in sync with [`TaskKind`].
17const N_TASK_KINDS: usize = 7;
18/// Rungs one-hot width (rungs beyond this fold into the last slot).
19const N_RUNGS: usize = 8;
20/// Fixed feature-vector length. Layout (see [`encode`]):
21/// `[bias, task_kind×7, prompt_bucket_norm, has_tools, tool_count_norm, has_images,
22///   session_fail_norm, rung×8]` = 1 + 7 + 1 + 1 + 1 + 1 + 1 + 8 = 21.
23pub const FEATURE_DIM: usize = 1 + N_TASK_KINDS + 5 + N_RUNGS;
24
25/// Index of the (unregularized) bias term.
26const BIAS: usize = 0;
27
28/// One-hot slot for a task kind, in the block that starts right after the bias.
29fn task_kind_slot(k: TaskKind) -> usize {
30    let i = match k {
31        TaskKind::CodeEdit => 0,
32        TaskKind::TestGen => 1,
33        TaskKind::Explore => 2,
34        TaskKind::Review => 3,
35        TaskKind::Extract => 4,
36        TaskKind::Chat => 5,
37        TaskKind::Other => 6,
38    };
39    1 + i
40}
41
42/// Numerically safe logistic sigmoid.
43fn sigmoid(z: f64) -> f64 {
44    let z = z.clamp(-30.0, 30.0);
45    1.0 / (1.0 + (-z).exp())
46}
47
48/// Encode `(features, rung)` into the fixed-length feature vector (see [`FEATURE_DIM`]).
49///
50/// Deterministic and privacy-preserving — it reads only the already-coarsened [`Features`]
51/// (buckets and flags, never raw prompt text).
52#[must_use]
53pub fn encode(features: &Features, rung: u32) -> [f64; FEATURE_DIM] {
54    let mut x = [0.0_f64; FEATURE_DIM];
55    x[BIAS] = 1.0;
56    x[task_kind_slot(features.task_kind)] = 1.0;
57    let base = 1 + N_TASK_KINDS;
58    // prompt token bucket is already floor(log2 tokens); normalize by a generous ceiling.
59    x[base] = f64::from(features.prompt_token_bucket).min(16.0) / 16.0;
60    x[base + 1] = f64::from(u8::from(features.tool_count > 0));
61    x[base + 2] = f64::from(features.tool_count).min(10.0) / 10.0;
62    x[base + 3] = f64::from(u8::from(features.has_images));
63    x[base + 4] = f64::from(features.session_failure_count).min(5.0) / 5.0;
64    let rung_slot = 1 + N_TASK_KINDS + 5 + (rung as usize).min(N_RUNGS - 1);
65    x[rung_slot] = 1.0;
66    x
67}
68
69/// An online logistic-regression predictor of gate-pass probability.
70///
71/// Cheap to update (one SGD step per observed attempt) and cheap to query (one dot product).
72/// Wrap in `Arc<Mutex<_>>` for the proxy, like the bandit — per-process, in-memory; the
73/// durable state is the receipts it warm-starts from.
74#[derive(Debug, Clone)]
75pub struct PassPredictor {
76    weights: [f64; FEATURE_DIM],
77    lr: f64,
78    l2: f64,
79}
80
81impl PassPredictor {
82    /// Create a predictor with learning rate `lr` (0, 1] and L2 penalty `l2` (>= 0), zero-init.
83    #[must_use]
84    pub fn new(lr: f64, l2: f64) -> Self {
85        Self {
86            weights: [0.0; FEATURE_DIM],
87            lr,
88            l2,
89        }
90    }
91
92    /// Predicted `P(gate-pass)` for `(features, rung)`, in `(0, 1)`.
93    #[must_use]
94    pub fn predict(&self, features: &Features, rung: u32) -> f64 {
95        let x = encode(features, rung);
96        let z: f64 = self
97            .weights
98            .iter()
99            .zip(x.iter())
100            .map(|(w, xi)| w * xi)
101            .sum();
102        sigmoid(z)
103    }
104
105    /// One SGD step against the observed outcome `passed` for `(features, rung)`.
106    ///
107    /// Gradient of the logistic loss: `(p - y)·x`, plus L2 shrinkage on every weight **except**
108    /// the bias (regularizing the intercept would bias the base rate).
109    pub fn update(&mut self, features: &Features, rung: u32, passed: bool) {
110        let x = encode(features, rung);
111        let p = {
112            let z: f64 = self
113                .weights
114                .iter()
115                .zip(x.iter())
116                .map(|(w, xi)| w * xi)
117                .sum();
118            sigmoid(z)
119        };
120        let err = p - f64::from(u8::from(passed));
121        for (i, w) in self.weights.iter_mut().enumerate() {
122            let reg = if i == BIAS { 0.0 } else { self.l2 * *w };
123            *w -= self.lr * (err * x[i] + reg);
124        }
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    fn feats(kind: TaskKind) -> Features {
133        Features::new(kind)
134    }
135
136    #[test]
137    fn encode_layout_is_fixed_and_onehot_correct() {
138        let x = encode(&feats(TaskKind::CodeEdit), 0);
139        assert_eq!(x.len(), FEATURE_DIM);
140        assert_eq!(x[BIAS], 1.0);
141        assert_eq!(x[task_kind_slot(TaskKind::CodeEdit)], 1.0);
142        // exactly one task-kind slot set
143        let tk_set: usize = (1..1 + N_TASK_KINDS).filter(|&i| x[i] == 1.0).count();
144        assert_eq!(tk_set, 1);
145        // exactly one rung slot set
146        let rbase = 1 + N_TASK_KINDS + 5;
147        let r_set: usize = (rbase..rbase + N_RUNGS).filter(|&i| x[i] == 1.0).count();
148        assert_eq!(r_set, 1);
149        // rung beyond the width folds into the last slot
150        let xr = encode(&feats(TaskKind::Other), 99);
151        assert_eq!(xr[rbase + N_RUNGS - 1], 1.0);
152    }
153
154    #[test]
155    fn predict_stays_in_unit_interval() {
156        let mut p = PassPredictor::new(0.5, 0.0);
157        for _ in 0..1000 {
158            p.update(&feats(TaskKind::Chat), 0, true);
159        }
160        let v = p.predict(&feats(TaskKind::Chat), 0);
161        assert!(v > 0.0 && v < 1.0, "prediction must stay in (0,1): {v}");
162    }
163
164    #[test]
165    fn converges_to_separate_easy_from_hard_by_rung() {
166        // Synthetic separable pattern: rung 0 always FAILS, rung 2 always PASSES for CodeEdit.
167        let mut p = PassPredictor::new(0.2, 1e-4);
168        let f = feats(TaskKind::CodeEdit);
169        for _ in 0..2000 {
170            p.update(&f, 0, false);
171            p.update(&f, 2, true);
172        }
173        let low = p.predict(&f, 0);
174        let high = p.predict(&f, 2);
175        assert!(
176            low < 0.2,
177            "hopeless rung should predict low pass, got {low}"
178        );
179        assert!(
180            high > 0.8,
181            "reliable rung should predict high pass, got {high}"
182        );
183        assert!(high - low > 0.6, "predictor must separate the two rungs");
184    }
185
186    #[test]
187    fn l2_keeps_weights_bounded() {
188        // Contradictory labels + L2 → weights must not blow up.
189        let mut p = PassPredictor::new(0.5, 0.05);
190        let f = feats(TaskKind::Review);
191        for i in 0..5000 {
192            p.update(&f, 1, i % 2 == 0);
193        }
194        let maxw = p
195            .weights
196            .iter()
197            .cloned()
198            .fold(0.0_f64, |a, w| a.max(w.abs()));
199        assert!(
200            maxw.is_finite() && maxw < 50.0,
201            "L2 must bound weights: {maxw}"
202        );
203        // ambiguous 50/50 label → prediction near 0.5
204        let v = p.predict(&f, 1);
205        assert!((v - 0.5).abs() < 0.15, "50/50 labels → ~0.5, got {v}");
206    }
207}
firstpass_core/predictor.rs

firstpass_core/
predictor.rs
