headroom/tests/test_verbosity_learn.py
Tejas Chopra a99dc61424
feat: output-token reduction — verbosity shaper, per-user learning, counterfactual savings (#965)
## Description

Adds the first levers that reduce the tokens the model **writes back**
(output), complementing Headroom's existing input compression. Output
costs 5× input on Opus-class models and is full of waste (ceremony,
restated code, deep "thinking" on routine steps). Two phases in one
self-contained PR off `main`: the request-side output shaper, then
per-user verbosity learning plus an honest counterfactual savings
estimator and dashboard surfacing.

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [x] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [x] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- **Output shaper** (`output_shaper.py`, opt-in
`HEADROOM_OUTPUT_SHAPER=1`): cache-safe verbosity steering appended to
the system-prompt tail (5 levels); effort routing that lowers
`output_config.effort` on mechanical tool-result continuations; legacy
`thinking.budget_tokens` clamp. Never injects effort where absent, never
toggles `thinking.type`.
- **`headroom learn --verbosity`**: mines Claude Code transcripts for
behavioral signals (interrupts, length-adaptive fast-skips, echo ratio),
recommends a verbosity level (heuristic + optional `--llm-judge`), and
seeds the savings baseline.
- **Counterfactual estimator** (`output_savings.py`): per-stratum
synthetic-control (estimated) + A/B holdout (measured) with a propagated
95% CI; conversation-stable arm assignment for A/B validity and
prefix-cache safety.
- **AIMD verbosity controller** (`verbosity_controller.py`):
additive-increase / fast-back-off state machine; live signal emission
gated off by default.
- **Wiring + surfaces**: shaper resolves the learned level; recording
rides the existing `transforms_applied` channel through the outcome
funnel (no `RequestOutcome` changes); `headroom output-savings` CLI;
dashboard "Output Tokens Saved" card.
- **Docs**: simple-words user guide + design doc with the counterfactual
methodology.

## 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
$ pytest tests/test_output_savings.py tests/test_output_savings_cli.py \
        tests/test_verbosity_learn.py tests/test_verbosity_controller.py \
        tests/test_output_shaper.py -q
94 passed in 0.54s

$ pytest tests/test_request_outcome.py tests/test_handler_outcome_tag_invariant.py \
        tests/test_proxy_dashboard_stats_cache.py -q
44 passed

$ ruff format --check .
831 files already formatted

$ mypy headroom --ignore-missing-imports
Success: no issues found in 361 source files
```

## Real Behavior Proof

- Environment: macOS, Python 3.12 (`.venv`), `anthropic` 0.76, live API
model `claude-opus-4-8`.
- Exact command / steps: `HEADROOM_OUTPUT_SHAPER=1`; `headroom learn
--verbosity --apply` (seeds level + baseline); `python
scripts/eval_output_shaper.py A` (live before/after); simulate holdout
traffic then `headroom output-savings`.
- Observed result: code-review ask — baseline 1,750 output tokens → L2
1,354 (−22.7%) → L3 599 (−65.8%), same bugs found. `learn --verbosity`
on 24 real sessions → 11% interrupt / 26% fast-skip → L3 (high
confidence). Measured A/B path → 31.7% reduction (95% CI 27.7%–35.7%).
94 new tests + 44 existing outcome/dashboard tests green; ruff + mypy
clean.
- Not tested: live streaming-path recording exercised only via unit
tests (the `transforms_applied` funnel is shared across paths); runtime
AIMD signal emission is gated off by default and not exercised live.

## 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
- [x] 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
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

Output savings are counterfactual (we never observe what the model
*would* have written), so the estimator separates **estimated** (vs a
learned baseline) from **measured** (A/B holdout via
`HEADROOM_OUTPUT_HOLDOUT`) and always reports a confidence band — never
a single made-up number. CHANGELOG left unchecked (release-please
manages it). Runtime AIMD self-tuning is intentionally a TODO
(controller built/tested; live signal emission gated behind
`HEADROOM_VERBOSITY_AUTOTUNE`).
2026-06-16 21:06:43 -07:00

251 lines
8 KiB
Python

"""Tests for headroom.learn.verbosity — behavioral signal extraction."""
from __future__ import annotations
import json
from pathlib import Path
from headroom.learn.verbosity import (
VerbositySignals,
analyze,
extract_signals,
recommend_level,
)
def _write_session(tmp_path: Path, name: str, lines: list[dict]) -> Path:
p = tmp_path / f"{name}.jsonl"
p.write_text("\n".join(json.dumps(line) for line in lines))
return p
def _assistant(
text: str,
*,
ts: str,
out_tokens: int = 100,
model: str = "claude-opus-4-8",
in_tokens: int = 5000,
) -> dict:
return {
"type": "assistant",
"timestamp": ts,
"message": {
"model": model,
"content": [{"type": "text", "text": text}],
"usage": {"input_tokens": in_tokens, "output_tokens": out_tokens},
},
}
def _user(text: str, *, ts: str) -> dict:
return {"type": "user", "timestamp": ts, "message": {"role": "user", "content": text}}
def _tool_result(*, ts: str, content: str = "ok") -> dict:
return {
"type": "user",
"timestamp": ts,
"message": {
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": content}],
},
}
LONG = " ".join(["word"] * 400) # well above the long-output floor
class TestSignalExtraction:
def test_interrupt_counted(self, tmp_path):
p = _write_session(
tmp_path,
"s",
[
_user("do a thing", ts="2026-01-01T00:00:00Z"),
_assistant(LONG, ts="2026-01-01T00:00:10Z"),
_user("[Request interrupted by user]", ts="2026-01-01T00:00:12Z"),
],
)
sig, _ = extract_signals([p])
assert sig.interrupts == 1
assert sig.human_msgs == 1 # the initial ask
def test_fast_skip_detected_length_adaptive(self, tmp_path):
# 400-word answer needs ~96s to read; reply after 5s = fast skip.
p = _write_session(
tmp_path,
"s",
[
_user("explain", ts="2026-01-01T00:00:00Z"),
_assistant(LONG, ts="2026-01-01T00:00:00Z"),
_user("ok next", ts="2026-01-01T00:00:05Z"),
],
)
sig, _ = extract_signals([p])
assert sig.skip_eligible == 1
assert sig.fast_skips == 1
def test_slow_reply_is_not_a_skip(self, tmp_path):
# Reply 120s after a 400-word answer (>read time) = read, not skipped.
p = _write_session(
tmp_path,
"s",
[
_user("explain", ts="2026-01-01T00:00:00Z"),
_assistant(LONG, ts="2026-01-01T00:00:00Z"),
_user("ok next", ts="2026-01-01T00:02:00Z"),
],
)
sig, _ = extract_signals([p])
assert sig.skip_eligible == 1
assert sig.fast_skips == 0
def test_short_answer_not_skip_eligible(self, tmp_path):
p = _write_session(
tmp_path,
"s",
[
_user("hi", ts="2026-01-01T00:00:00Z"),
_assistant("short reply", ts="2026-01-01T00:00:00Z"),
_user("ok", ts="2026-01-01T00:00:01Z"),
],
)
sig, _ = extract_signals([p])
assert sig.skip_eligible == 0
def test_baseline_captures_output_tokens_by_stratum(self, tmp_path):
p = _write_session(
tmp_path,
"s",
[
_user("task", ts="2026-01-01T00:00:00Z"),
_assistant("a reply", ts="2026-01-01T00:00:01Z", out_tokens=420, in_tokens=5000),
],
)
_, baseline = extract_signals([p])
assert baseline.total_samples == 1
# new_user_ask, input bucket "s" (5000), opus, no tools in this session
mean, _, n = baseline.lookup("opus|new_user_ask|s|notools")
assert n == 1
assert mean == 420.0
def test_tool_result_makes_session_have_tools(self, tmp_path):
p = _write_session(
tmp_path,
"s",
[
_user("task", ts="2026-01-01T00:00:00Z"),
_assistant("reading", ts="2026-01-01T00:00:01Z", out_tokens=50),
_tool_result(ts="2026-01-01T00:00:02Z"),
_assistant("done", ts="2026-01-01T00:00:03Z", out_tokens=200),
],
)
_, baseline = extract_signals([p])
# Every response in a tool-using session is stratified as has_tools.
assert any("|tools" in k for k in baseline.strata)
assert not any("|notools" in k for k in baseline.strata)
def test_tool_result_reply_not_counted_as_human(self, tmp_path):
p = _write_session(
tmp_path,
"s",
[
_user("task", ts="2026-01-01T00:00:00Z"),
_assistant("reading", ts="2026-01-01T00:00:01Z"),
_tool_result(ts="2026-01-01T00:00:02Z"),
],
)
sig, _ = extract_signals([p])
assert sig.human_msgs == 1 # only the real ask, not the tool_result
class TestRecommendLevel:
def _sig(self, *, human, interrupts, skip_eligible, fast_skips) -> VerbositySignals:
s = VerbositySignals()
s.human_msgs = human
s.interrupts = interrupts
s.skip_eligible = skip_eligible
s.fast_skips = fast_skips
return s
def test_too_few_turns_defaults_l2_low(self):
level, conf, _ = recommend_level(
self._sig(human=3, interrupts=0, skip_eligible=0, fast_skips=0)
)
assert level == 2
assert conf == "low"
def test_low_pressure_user_gets_l1(self):
# 100 turns, almost no interrupts/skips.
s = self._sig(human=100, interrupts=1, skip_eligible=100, fast_skips=2)
level, conf, _ = recommend_level(s)
assert level == 1
assert conf == "high"
def test_moderate_pressure_gets_l2(self):
s = self._sig(human=80, interrupts=8, skip_eligible=80, fast_skips=12)
level, _, _ = recommend_level(s)
assert level == 2
def test_high_pressure_gets_l3(self):
# Mirrors the real measured user: ~11% interrupt, ~26% skip.
s = self._sig(human=200, interrupts=29, skip_eligible=119, fast_skips=31)
level, conf, _ = recommend_level(s)
assert level == 3
assert conf == "high"
class TestAnalyze:
def test_llm_judge_overrides_heuristic(self, tmp_path):
p = _write_session(
tmp_path,
"s",
[_user("x", ts="2026-01-01T00:00:00Z"), _assistant("y", ts="2026-01-01T00:00:01Z")]
* 20,
)
def judge(signals_dict):
return 4, "LLM says this user wants caveman mode"
profile, _ = analyze([p], "/proj", llm_judge=judge)
assert profile.level == 4
assert profile.source == "llm"
assert "caveman" in profile.rationale
def test_llm_judge_failure_falls_back_to_heuristic(self, tmp_path):
p = _write_session(
tmp_path,
"s",
[_user("x", ts="2026-01-01T00:00:00Z"), _assistant("y", ts="2026-01-01T00:00:01Z")]
* 20,
)
def bad_judge(signals_dict):
raise RuntimeError("no api key")
profile, _ = analyze([p], "/proj", llm_judge=bad_judge)
assert profile.source == "heuristic"
def test_profile_roundtrip(self, tmp_path):
from headroom.learn.verbosity import VerbosityProfile
prof = VerbosityProfile(
project_path="/proj",
level=3,
confidence="high",
source="heuristic",
rationale="because",
signals={"interrupt_rate": 0.11},
)
path = tmp_path / "verbosity.json"
prof.save(path)
loaded = VerbosityProfile.load(path)
assert loaded is not None
assert loaded.level == 3
assert loaded.confidence == "high"
def test_load_missing_returns_none(self, tmp_path):
from headroom.learn.verbosity import VerbosityProfile
assert VerbosityProfile.load(tmp_path / "nope.json") is None