headroom/tests/test_output_shaping_rollup.py
gglucass 12a9710665
feat(stats): per-bucket output-shaping savings in /stats-history (#1819)
## Description

Adds per-bucket **output-shaping savings** to `/stats-history`. Today
output-shaping savings exist only as a single global aggregate
(`savings.by_layer.output_shaping`), so downstream consumers can't chart
them over time. This threads a per-request output-savings estimate into
the existing rollup so every `series` bucket carries
`output_tokens_saved_delta` + `output_savings_usd_delta`, symmetric with
the existing `compression_savings_usd_delta`.

Motivation: on Claude Code subscription traffic, input is ~99%
cache-discounted (the compressible live zone is a fraction of a
percent), while output shaping is a ~36% reduction on full-price output
tokens — so it's the dominant, honestly-attributable saving, and
currently the only one a dashboard can't render per day.

Closes #1816

## Type of Change

- [x] New feature (non-breaking change that adds functionality)

## Changes Made

- `output_savings.py`: new read-only
`SavingsRecorder.estimate_request_savings(labels, output_tokens)` →
per-request synthetic-control estimate `max(0, baseline_mean(stratum) -
output_tokens)` for treatment requests; 0 for control / unknown stratum
/ no label. Does **not** mutate the ledger, so it composes with
`record_from_labels` without double-counting. `record_from_labels`'s
`bool` contract is unchanged.
- `outcome.py`: in the funnel, capture that estimate and pass it to
`record_request(output_tokens_saved=...)`.
- `savings_tracker.py`: `record_request` gains `output_tokens_saved`;
accumulates lifetime cumulative `output_tokens_saved` /
`output_savings_usd` (priced via new `_estimate_output_savings_usd`,
output-rate), writes them into each checkpoint, and now checkpoints when
**either** compression **or** output savings occurred (so output-only
requests aren't dropped). `_build_rollup` diffs the cumulative into
`output_tokens_saved_delta` / `output_savings_usd_delta` per bucket;
`_normalize_history_entry` and the CSV export carry the fields.
- Additive + backward-compatible: checkpoints predating the feature
default the new fields to 0.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ uv run --extra dev pytest tests/test_output_shaping_rollup.py tests/test_output_savings.py \
    tests/test_output_savings_cli.py tests/test_proxy_savings_history.py tests/test_request_outcome.py -q
... 103 passed

$ uv run --extra dev ruff check headroom/proxy/savings_tracker.py headroom/proxy/output_savings.py \
    headroom/proxy/prometheus_metrics.py headroom/proxy/outcome.py tests/test_output_shaping_rollup.py
All checks passed!

$ uv run --extra dev mypy headroom/proxy/savings_tracker.py headroom/proxy/output_savings.py
Success: no issues found in 2 source files
```

New tests (`tests/test_output_shaping_rollup.py`): output savings bucket
into the daily series; an output-only request (no compression) still
checkpoints; pre-feature requests default to 0;
`estimate_request_savings` returns the baseline-relative saving for
treatment and 0 for control / unknown / over-baseline.

## Real Behavior Proof

- Environment: macOS, CPython 3.10.18, this branch (rebased on latest
`main`), litellm pricing available.
- Exact command / steps: seed a baseline (as `learn --verbosity` would),
then drive 3 requests through the real, unmocked chain
`SavingsRecorder.estimate_request_savings` →
`SavingsTracker.record_request` → `history_response()`, and print
`series.daily`. Full script + raw output:

```text
$ uv run python proof.py   # seeds baseline ~1000 out-tok; 3 treatment requests (out=600/550/700), one with no compression
[
  { "timestamp": "2026-07-05T00:00:00Z", "tokens_saved": 120,
    "compression_savings_usd_delta": 0.0006,
    "output_tokens_saved_delta": 850, "output_savings_usd_delta": 0.02125 },
  { "timestamp": "2026-07-06T00:00:00Z", "tokens_saved": 80,
    "compression_savings_usd_delta": 0.0004,
    "output_tokens_saved_delta": 300, "output_savings_usd_delta": 0.0075 }
]
```

- Observed result: output-shaping savings appear per day and independent
of the compression axis. 2026-07-05 = 850 (400+450 saved by two
treatment requests vs the ~1000-token baseline, including one request
with zero compression — proving the output-only checkpoint path),
2026-07-06 = 300, each priced at the model's output rate. Matches
expectations.
- Not tested: the full live proxy over HTTP with a real learned baseline
and organic traffic — I exercised the same code path minus the
HTTP/streaming layer. The measured-vs-estimated `method` gating is
unchanged by this PR.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Screenshots (if applicable)

N/A — backend-only change (no UI surface in this repo). The runtime
effect is the `/stats-history` `series.daily` JSON with the new
`output_tokens_saved_delta` / `output_savings_usd_delta` fields, shown
under **Real Behavior Proof** above. The downstream chart that renders
them lives in the separate Headroom desktop app.

## Additional Notes

- Per CONTRIBUTING's issue-first policy for features, I opened #1816
first with the spec; happy to adjust the API surface (field names /
gating) to whatever you prefer. A downstream consumer (Headroom desktop
chart) is already implemented against this exact contract and stacks the
segment only when `output_reduction.method == "measured"`.
- Docs checkbox left unchecked: I didn't find a `/stats-history` schema
doc to update; point me at one if it exists.

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
2026-07-15 19:58:24 +00:00

88 lines
3.2 KiB
Python

"""Per-bucket output-shaping savings in the /stats-history rollup.
Covers the feature that lets a downstream dashboard stack output-shaping
savings as a distinct daily segment: SavingsTracker.record_request accepts a
per-request output_tokens_saved, accumulates it into each time bucket as
output_tokens_saved_delta / output_savings_usd_delta, and the read-only
SavingsRecorder.estimate_request_savings supplies that per-request number.
"""
from __future__ import annotations
from headroom.proxy.output_savings import (
SavingsRecorder,
stratum_key,
stratum_label,
)
from headroom.proxy.savings_tracker import SavingsTracker
def test_record_request_buckets_output_shaping_savings(tmp_path):
tracker = SavingsTracker(path=str(tmp_path / "s.json"))
# Request with both compression and output-shaping savings.
tracker.record_request(
model="claude-opus-4-8",
input_tokens=1000,
tokens_saved=100,
output_tokens_saved=5000,
timestamp="2026-03-27T09:00:00Z",
)
# Output-shaping-ONLY request (no compression) must still checkpoint, else
# its output savings would be dropped from the rollup.
tracker.record_request(
model="claude-opus-4-8",
input_tokens=1000,
tokens_saved=0,
output_tokens_saved=3000,
timestamp="2026-03-27T09:30:00Z",
)
daily = tracker.history_response()["series"]["daily"]
assert len(daily) == 1
assert daily[0]["output_tokens_saved_delta"] == 8000
assert daily[0]["output_savings_usd_delta"] > 0.0
# Compression axis stays independent.
assert daily[0]["tokens_saved"] == 100
def test_record_request_without_output_savings_is_backward_compatible(tmp_path):
tracker = SavingsTracker(path=str(tmp_path / "s.json"))
tracker.record_request(
model="gpt-4o",
input_tokens=8192,
tokens_saved=4096,
timestamp="2026-03-27T09:00:00Z",
)
daily = tracker.history_response()["series"]["daily"]
assert daily[0]["output_tokens_saved_delta"] == 0
assert daily[0]["output_savings_usd_delta"] == 0.0
def _key() -> str:
return stratum_key(turn_kind="code", input_tokens=8000, model="claude-opus-4-8", has_tools=True)
def test_estimate_request_savings_treatment_uses_baseline(tmp_path):
rec = SavingsRecorder(str(tmp_path / "o.json"), flush_every=1)
key = _key()
for _ in range(5):
rec._ledger.baseline.observe(key, 1000) # baseline mean ~1000
# Treatment request that emitted 600 -> saved ~400 vs the baseline.
saved = rec.estimate_request_savings([stratum_label("treatment", key)], 600)
assert saved == 400
def test_estimate_request_savings_zero_for_control_and_unknown(tmp_path):
rec = SavingsRecorder(str(tmp_path / "o.json"), flush_every=1)
key = _key()
for _ in range(5):
rec._ledger.baseline.observe(key, 1000)
# Control arm is unshaped -> no attributable saving.
assert rec.estimate_request_savings([stratum_label("control", key)], 600) == 0
# No shaping label at all.
assert rec.estimate_request_savings(["something-else"], 600) == 0
# Treatment but output exceeded the baseline -> clamped to 0, never negative.
assert rec.estimate_request_savings([stratum_label("treatment", key)], 5000) == 0