lean_ctx_proxy/

output_savings.rs

// SPDX-License-Identifier: Apache-2.0
//! Output-token savings reporting (#895 Track B).
//!
//! lean-ctx shapes *output* tokens (cache-safe effort control #834 and the
//! optional verbosity steer). To report that honestly we distinguish two cases:
//!
//! * **Measured** — when a holdout ([`lean_ctx::core::config::ProxyConfig::output_holdout_fraction`])
//!   has produced enough paired turns, the reduction is
//!   `(control_avg − treatment_avg) / control_avg`, with a 95 % confidence
//!   interval from a two-sample (Welch) standard error on the difference of mean
//!   output tokens. This is a real A/B result, not a guess.
//! * **Estimated** — with no holdout (or not enough data yet) we fall back to the
//!   shared model-based heuristic and present it as a *band*, never a hard
//!   number, clearly labelled so it is never mistaken for a measurement.

use std::collections::HashMap;

use super::usage_meter::CohortUsage;

/// Minimum turns *per arm* before a measured reduction is reported. Below this
/// the interval is too wide to be informative, so we report [`Savings::Pending`].
pub const MIN_SAMPLES_PER_ARM: u64 = 30;

/// Standard-normal 97.5th percentile (two-sided 95 % CI multiplier).
const Z95: f64 = 1.959_964;

/// Relative half-width applied to the model-based point estimate to express its
/// inherent uncertainty as a band. This is an *estimate* parameter, not a
/// measurement — enabling `output_holdout` replaces it with a real CI.
const ESTIMATE_REL_UNCERTAINTY: f64 = 0.35;

/// Outcome of an output-savings query, ready for the CLI / dashboard to render.
#[derive(Debug, Clone, PartialEq)]
pub enum Savings {
    /// Real A/B result from the holdout.
    Measured(Measured),
    /// Holdout running but not enough paired turns yet.
    Pending {
        control_n: u64,
        treatment_n: u64,
        needed: u64,
    },
    /// No holdout: model-based estimate, presented as a band.
    Estimated {
        point_pct: f64,
        low_pct: f64,
        high_pct: f64,
    },
}

/// A measured output-token reduction with its 95 % confidence interval.
#[derive(Debug, Clone, PartialEq)]
pub struct Measured {
    pub control_avg: f64,
    pub treatment_avg: f64,
    pub tokens_saved_per_turn: f64,
    pub reduction_pct: f64,
    pub ci95_low_pct: f64,
    pub ci95_high_pct: f64,
    pub control_n: u64,
    pub treatment_n: u64,
}

/// Computes the output-savings outcome from the persisted cohort totals.
#[must_use]
pub fn current() -> Savings {
    from_cohorts(&super::usage_meter::persisted_cohorts())
}

/// Stable JSON shape for a [`Savings`], shared by the `lean-ctx output-savings`
/// CLI and the dashboard `/api/roi` route so both speak one schema.
#[must_use]
pub fn to_json(s: &Savings) -> serde_json::Value {
    match s {
        Savings::Measured(m) => serde_json::json!({
            "status": "measured",
            "reduction_pct": round2(m.reduction_pct),
            "ci95_low_pct": round2(m.ci95_low_pct),
            "ci95_high_pct": round2(m.ci95_high_pct),
            "control_avg_output": round2(m.control_avg),
            "treatment_avg_output": round2(m.treatment_avg),
            "tokens_saved_per_turn": round2(m.tokens_saved_per_turn),
            "control_n": m.control_n,
            "treatment_n": m.treatment_n,
        }),
        Savings::Pending {
            control_n,
            treatment_n,
            needed,
        } => serde_json::json!({
            "status": "pending",
            "control_n": control_n,
            "treatment_n": treatment_n,
            "needed_per_arm": needed,
        }),
        Savings::Estimated {
            point_pct,
            low_pct,
            high_pct,
        } => serde_json::json!({
            "status": "estimated",
            "point_pct": round2(*point_pct),
            "low_pct": round2(*low_pct),
            "high_pct": round2(*high_pct),
        }),
    }
}

fn round2(x: f64) -> f64 {
    (x * 100.0).round() / 100.0
}

/// Pure core: decide measured / pending / estimated from cohort totals.
#[must_use]
pub fn from_cohorts(cohorts: &HashMap<String, CohortUsage>) -> Savings {
    let control = cohorts.get("control");
    let treatment = cohorts.get("treatment");
    let (control_n, treatment_n) = (
        control.map_or(0, |c| c.requests),
        treatment.map_or(0, |t| t.requests),
    );

    // Both arms must exist for any comparison; otherwise there is no experiment.
    if control_n == 0 || treatment_n == 0 {
        return estimated();
    }
    if control_n < MIN_SAMPLES_PER_ARM || treatment_n < MIN_SAMPLES_PER_ARM {
        return Savings::Pending {
            control_n,
            treatment_n,
            needed: MIN_SAMPLES_PER_ARM,
        };
    }
    // Safe: presence + counts checked above.
    let (control, treatment) = (
        control.expect("control cohort present when control_n > 0"),
        treatment.expect("treatment cohort present when treatment_n > 0"),
    );
    match measured(control, treatment) {
        Some(m) => Savings::Measured(m),
        None => estimated(),
    }
}

/// Two-sample reduction with a Welch 95 % CI, or `None` (not measured) if the
/// sample is too small or degenerate: control mean 0, or a variance that is
/// zero, undefined or inconsistent with the sums (a zero-width CI is not a
/// measurement).
fn measured(control: &CohortUsage, treatment: &CohortUsage) -> Option<Measured> {
    if control.requests < MIN_SAMPLES_PER_ARM || treatment.requests < MIN_SAMPLES_PER_ARM {
        return None;
    }
    welch_reduction(
        Sample::new(
            control.output_tokens,
            control.sum_sq_output,
            control.requests,
        ),
        Sample::new(
            treatment.output_tokens,
            treatment.sum_sq_output,
            treatment.requests,
        ),
    )
}

/// Running totals of one per-turn quantity in one arm.
#[derive(Debug, Clone, Copy)]
pub(crate) struct Sample {
    sum: u64,
    sum_sq: u64,
    n: u64,
}

impl Sample {
    pub(crate) fn new(sum: u64, sum_sq: u64, n: u64) -> Self {
        Self { sum, sum_sq, n }
    }

    #[allow(clippy::cast_precision_loss)]
    fn mean(self) -> Option<f64> {
        (self.n > 0).then(|| self.sum as f64 / self.n as f64)
    }

    /// Unbiased sample variance, clamped at 0 against floating-point error.
    /// `None` with < 2 observations, or when the sum of squares cannot belong
    /// to the sum (`n·Σx² < (Σx)²`) — e.g. totals whose squares were lost.
    #[allow(clippy::cast_precision_loss)]
    fn variance(self) -> Option<f64> {
        if self.n < 2 {
            return None;
        }
        let sum = u128::from(self.sum);
        if u128::from(self.sum_sq) * u128::from(self.n) < sum * sum {
            return None;
        }
        let (n, sum, sum_sq) = (self.n as f64, self.sum as f64, self.sum_sq as f64);
        Some(((sum_sq - sum * sum / n) / (n - 1.0)).max(0.0))
    }
}

/// Relative reduction of the per-turn mean from `control` to `treatment`, with
/// a 95 % CI from the Welch standard error of the difference of means. Shared
/// by the output-savings and input-compression holdouts. `None` if degenerate.
pub(crate) fn welch_reduction(control: Sample, treatment: Sample) -> Option<Measured> {
    let control_avg = control.mean()?;
    let treatment_avg = treatment.mean()?;
    if control_avg <= 0.0 {
        return None;
    }
    let var_c = control.variance()?;
    let var_t = treatment.variance()?;
    #[allow(clippy::cast_precision_loss)]
    let (n_c, n_t) = (control.n as f64, treatment.n as f64);

    let diff = control_avg - treatment_avg; // tokens saved per turn (may be < 0)
    let se = (var_c / n_c + var_t / n_t).sqrt();
    // No spread means no interval: real per-turn token counts always vary, so
    // a zero standard error signals lost or synthetic sums, not certainty.
    if !(se.is_finite() && se > 0.0) {
        return None;
    }
    let margin = Z95 * se;

    Some(Measured {
        control_avg,
        treatment_avg,
        tokens_saved_per_turn: diff,
        reduction_pct: diff / control_avg * 100.0,
        ci95_low_pct: (diff - margin) / control_avg * 100.0,
        ci95_high_pct: (diff + margin) / control_avg * 100.0,
        control_n: control.n,
        treatment_n: treatment.n,
    })
}

/// Model-based estimate band from the shared verbose/concise heuristic.
fn estimated() -> Savings {
    let model = lean_ctx::core::stats::CostModel::default();
    #[allow(clippy::cast_precision_loss)]
    let verbose = model.avg_verbose_output_per_call as f64;
    #[allow(clippy::cast_precision_loss)]
    let concise = model.avg_concise_output_per_call as f64;
    let point = if verbose > 0.0 {
        (verbose - concise) / verbose * 100.0
    } else {
        0.0
    };
    let half = point * ESTIMATE_REL_UNCERTAINTY;
    Savings::Estimated {
        point_pct: point,
        low_pct: (point - half).max(0.0),
        high_pct: point + half,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn cohort(requests: u64, output_tokens: u64, sum_sq_output: u64) -> CohortUsage {
        CohortUsage {
            requests,
            input_tokens: 0,
            output_tokens,
            sum_sq_output,
            ..CohortUsage::default()
        }
    }

    /// Builds a cohort from explicit per-turn samples (exact sums).
    fn from_samples(samples: &[u64]) -> CohortUsage {
        let mut c = CohortUsage::default();
        for &s in samples {
            c.requests += 1;
            c.output_tokens += s;
            c.sum_sq_output += s * s;
        }
        c
    }

    #[test]
    fn no_cohorts_falls_back_to_estimate() {
        let s = from_cohorts(&HashMap::new());
        match s {
            Savings::Estimated {
                point_pct,
                low_pct,
                high_pct,
            } => {
                // Default heuristic 180→120 ⇒ 33.3% point.
                assert!((point_pct - 33.333).abs() < 0.1, "point {point_pct}");
                assert!(low_pct < point_pct && point_pct < high_pct);
                assert!(low_pct >= 0.0);
            }
            other => panic!("expected Estimated, got {other:?}"),
        }
    }

    #[test]
    fn single_arm_only_is_estimate_not_measured() {
        let mut m = HashMap::new();
        m.insert("control".to_string(), cohort(100, 18_000, 3_240_000));
        assert!(matches!(from_cohorts(&m), Savings::Estimated { .. }));
    }

    #[test]
    fn small_paired_sample_is_pending() {
        let mut m = HashMap::new();
        m.insert("control".to_string(), from_samples(&[180, 190, 170]));
        m.insert("treatment".to_string(), from_samples(&[120, 130, 110]));
        match from_cohorts(&m) {
            Savings::Pending {
                control_n,
                treatment_n,
                needed,
            } => {
                assert_eq!(control_n, 3);
                assert_eq!(treatment_n, 3);
                assert_eq!(needed, MIN_SAMPLES_PER_ARM);
            }
            other => panic!("expected Pending, got {other:?}"),
        }
    }

    #[test]
    fn large_paired_sample_is_measured_with_ci() {
        // 40 turns/arm: control ~180, treatment ~120 → ~33% reduction.
        let control = from_samples(&[170, 180, 190, 185].repeat(10));
        let treatment = from_samples(&[110, 120, 130, 125].repeat(10));
        let mut m = HashMap::new();
        m.insert("control".to_string(), control);
        m.insert("treatment".to_string(), treatment);
        match from_cohorts(&m) {
            Savings::Measured(measured) => {
                assert_eq!(measured.control_n, 40);
                assert_eq!(measured.treatment_n, 40);
                assert!(
                    (measured.reduction_pct - 33.0).abs() < 3.0,
                    "reduction {}",
                    measured.reduction_pct
                );
                assert!(
                    measured.ci95_low_pct < measured.reduction_pct
                        && measured.reduction_pct < measured.ci95_high_pct,
                    "CI must bracket the point estimate"
                );
                assert!(measured.tokens_saved_per_turn > 0.0);
            }
            other => panic!("expected Measured, got {other:?}"),
        }
    }

    #[test]
    fn zero_variance_is_not_reported_as_measured() {
        // Constant per-arm output → variance 0 → a zero-width CI, which is no
        // measurement at all.
        let control = from_samples(&[180; 40]);
        let treatment = from_samples(&[120; 40]);
        let mut m = HashMap::new();
        m.insert("control".to_string(), control);
        m.insert("treatment".to_string(), treatment);
        assert!(
            !matches!(from_cohorts(&m), Savings::Measured(_)),
            "zero variance must not be reported as measured"
        );
    }

    #[test]
    fn lost_sums_of_squares_are_not_reported_as_measured() {
        // Pre-fix restarts dropped `sum_sq_output`: the totals then imply a
        // negative variance. Those cohorts must read as not measured.
        let mut control = from_samples(&[170, 180, 190, 185].repeat(10));
        let mut treatment = from_samples(&[110, 120, 130, 125].repeat(10));
        control.sum_sq_output = 0;
        treatment.sum_sq_output = 0;
        let mut m = HashMap::new();
        m.insert("control".to_string(), control);
        m.insert("treatment".to_string(), treatment);
        assert!(!matches!(from_cohorts(&m), Savings::Measured(_)));
    }

    #[test]
    fn measured_requires_minimum_samples_on_its_own() {
        let control = from_samples(&[170, 180, 190]);
        let treatment = from_samples(&[110, 120, 130]);
        assert!(measured(&control, &treatment).is_none());
    }
}