1use 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
25const MIN_TURNS_PER_ARM: u64 = 30;
28
29#[derive(Debug, Clone, PartialEq)]
31pub enum CompressionSavings {
32 Measured(Measured),
34 Pending {
36 control_n: u64,
37 treatment_n: u64,
38 needed: u64,
39 },
40 Off,
42}
43
44#[must_use]
46pub fn current() -> CompressionSavings {
47 from_cohorts(&super::usage_meter::persisted_cohorts())
48}
49
50#[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#[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 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 #[test]
146 fn reports_only_what_the_holdout_measured() {
147 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 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 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}