Skip to main content

lean_ctx_proxy/
compression_savings.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Input-compression savings from the `[proxy] compression_holdout` (#1905).
3//!
4//! Shadow Mode simulates its baseline from the treatment's own numbers, so its
5//! quality delta is zero by construction. The holdout instead forwards a
6//! deterministic fraction of conversations uncompressed and meters both arms,
7//! which gives a real baseline for the prompt size:
8//!
9//! * **Measured** — enough turns in both arms: the reduction of whole-prompt
10//!   tokens per turn, with a Welch 95 % confidence interval.
11//! * **Pending** — the holdout runs, but an arm is below the minimum sample.
12//! * **Off** — no holdout data. There is deliberately no estimate here: an
13//!   unmeasured number would be exactly what #1905 was about.
14//!
15//! Answer quality is never inferred from token counts. The holdout does not
16//! measure it, so every report says `quality: unknown`.
17
18use std::collections::HashMap;
19
20use super::holdout::Arm;
21use super::output_savings::{Measured, Sample, welch_reduction};
22use super::usage_meter::{CohortUsage, compression_cohort_key};
23use lean_ctx::core::context_quality::EvidenceTier;
24
25/// Turns each arm needs before a reduction is reported: the 30-pair power floor
26/// the paired quality reports used, so "measured" keeps one meaning.
27const MIN_TURNS_PER_ARM: u64 = 30;
28
29/// Outcome of an input-compression savings query.
30#[derive(Debug, Clone, PartialEq)]
31pub enum CompressionSavings {
32    /// Real A/B result from the holdout.
33    Measured(Measured),
34    /// Holdout running but not enough turns in one of the arms yet.
35    Pending {
36        control_n: u64,
37        treatment_n: u64,
38        needed: u64,
39    },
40    /// No holdout data.
41    Off,
42}
43
44/// Computes the outcome from the persisted cohort totals.
45#[must_use]
46pub fn current() -> CompressionSavings {
47    from_cohorts(&super::usage_meter::persisted_cohorts())
48}
49
50/// Pure core: decide measured / pending / off from cohort totals.
51#[must_use]
52pub fn from_cohorts(cohorts: &HashMap<String, CohortUsage>) -> CompressionSavings {
53    let control = cohorts.get(&compression_cohort_key(Arm::Control));
54    let treatment = cohorts.get(&compression_cohort_key(Arm::Treatment));
55    let (control_n, treatment_n) = (
56        control.map_or(0, |c| c.requests),
57        treatment.map_or(0, |t| t.requests),
58    );
59    let (Some(control), Some(treatment)) = (control, treatment) else {
60        return CompressionSavings::Off;
61    };
62    if control_n < MIN_TURNS_PER_ARM || treatment_n < MIN_TURNS_PER_ARM {
63        return CompressionSavings::Pending {
64            control_n,
65            treatment_n,
66            needed: MIN_TURNS_PER_ARM,
67        };
68    }
69    welch_reduction(prompt_sample(control), prompt_sample(treatment))
70        .map_or(CompressionSavings::Off, CompressionSavings::Measured)
71}
72
73fn prompt_sample(c: &CohortUsage) -> Sample {
74    Sample::new(c.prompt_tokens, c.sum_sq_prompt, c.requests)
75}
76
77/// Stable JSON shape, shared by the CLI and any dashboard route.
78///
79/// A measured reduction carries [`EvidenceTier::ProductionOutcome`]: it was
80/// observed in a real deployment with a real control arm. The tier describes
81/// the token measurement only; quality stays `unknown` in every state.
82#[must_use]
83pub fn to_json(s: &CompressionSavings) -> serde_json::Value {
84    let mut v = match s {
85        CompressionSavings::Measured(m) => serde_json::json!({
86            "status": "measured",
87            "power": "powered",
88            "evidence_tier": EvidenceTier::ProductionOutcome,
89            "reduction_pct": round2(m.reduction_pct),
90            "ci95_low_pct": round2(m.ci95_low_pct),
91            "ci95_high_pct": round2(m.ci95_high_pct),
92            "control_avg_prompt_tokens": round2(m.control_avg),
93            "treatment_avg_prompt_tokens": round2(m.treatment_avg),
94            "tokens_saved_per_turn": round2(m.tokens_saved_per_turn),
95            "control_n": m.control_n,
96            "treatment_n": m.treatment_n,
97        }),
98        CompressionSavings::Pending {
99            control_n,
100            treatment_n,
101            needed,
102        } => serde_json::json!({
103            "status": "pending",
104            "power": "underpowered",
105            "control_n": control_n,
106            "treatment_n": treatment_n,
107            "needed_per_arm": needed,
108        }),
109        CompressionSavings::Off => serde_json::json!({ "status": "off" }),
110    };
111    v["quality"] = serde_json::json!("unknown");
112    v
113}
114
115fn round2(x: f64) -> f64 {
116    (x * 100.0).round() / 100.0
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    /// `n` turns in one arm, each with `prompt` whole-prompt tokens ± `spread`.
124    fn arm(n: u64, prompt: u64, spread: u64) -> CohortUsage {
125        let mut c = CohortUsage::default();
126        for i in 0..n {
127            let p = prompt + (i % 3) * spread;
128            c.requests += 1;
129            c.prompt_tokens += p;
130            c.sum_sq_prompt += p * p;
131        }
132        c
133    }
134
135    fn cohorts(control: CohortUsage, treatment: CohortUsage) -> HashMap<String, CohortUsage> {
136        HashMap::from([
137            (compression_cohort_key(Arm::Control), control),
138            (compression_cohort_key(Arm::Treatment), treatment),
139        ])
140    }
141
142    /// The contract of the holdout report: no estimate without both arms, no
143    /// number below the sample minimum, a measured reduction with an interval
144    /// above it, and output-savings arms never leak into this comparison.
145    #[test]
146    fn reports_only_what_the_holdout_measured() {
147        // Output-savings arms alone are not a compression experiment.
148        let only_output = HashMap::from([
149            ("control".to_string(), arm(50, 1000, 10)),
150            ("treatment".to_string(), arm(50, 600, 10)),
151        ]);
152        assert_eq!(from_cohorts(&only_output), CompressionSavings::Off);
153
154        let small = from_cohorts(&cohorts(arm(10, 1000, 10), arm(40, 600, 10)));
155        assert!(matches!(
156            small,
157            CompressionSavings::Pending {
158                control_n: 10,
159                treatment_n: 40,
160                ..
161            }
162        ));
163
164        let CompressionSavings::Measured(m) =
165            from_cohorts(&cohorts(arm(60, 1000, 30), arm(60, 600, 30)))
166        else {
167            panic!("enough turns in both arms must be measured");
168        };
169        // Means 1030 vs 630 (the +0/+30/+60 spread adds 30 to both arms).
170        let expected = 400.0 / 1030.0 * 100.0;
171        assert!((m.reduction_pct - expected).abs() < 1e-9, "{m:?}");
172        assert!(m.ci95_low_pct < m.reduction_pct && m.reduction_pct < m.ci95_high_pct);
173
174        // Only a measured result claims the production-outcome tier, and the
175        // tier never upgrades quality, which stays unknown in every state.
176        assert_eq!(to_json(&small)["power"], "underpowered");
177        assert!(to_json(&small).get("evidence_tier").is_none());
178        let measured = to_json(&CompressionSavings::Measured(m.clone()));
179        assert_eq!(measured["evidence_tier"], "production_outcome");
180
181        for s in [
182            CompressionSavings::Off,
183            small,
184            CompressionSavings::Measured(m),
185        ] {
186            assert_eq!(to_json(&s)["quality"], "unknown", "{s:?}");
187        }
188    }
189}