mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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`).
This commit is contained in:
Tejas Chopra
GitHub
parent
b7be3814f1
commit
a99dc61424
21 changed files with 3678 additions and 3 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -21,6 +21,7 @@ scripts/*
|
|||
!scripts/tests/
|
||||
!scripts/README.md
|
||||
!scripts/repro_codex_replay.py
|
||||
!scripts/eval_output_shaper.py
|
||||
!scripts/fixtures/
|
||||
!scripts/fixtures/*.json
|
||||
!scripts/record_fixtures.py
|
||||
|
|
|
|||
49
README.md
49
README.md
|
|
@ -52,6 +52,7 @@ Headroom compresses everything your AI agent reads — tool outputs, logs, RAG c
|
|||
- **MCP server** — `headroom_compress`, `headroom_retrieve`, `headroom_stats` for any MCP client
|
||||
- **Cross-agent memory** — shared store across Claude, Codex, Gemini, auto-dedup
|
||||
- **`headroom learn`** — mines failed sessions, writes corrections to `CLAUDE.md` / `AGENTS.md`
|
||||
- **Output token reduction** — trims what the model *writes back* (not just what you send): drops ceremony/restated code and skips deep "thinking" on routine steps. See [Output token reduction](#output-token-reduction-cut-what-the-model-writes-back).
|
||||
- **Reversible (CCR)** — originals are cached for retrieval on demand
|
||||
|
||||
## How it works (30 seconds)
|
||||
|
|
@ -123,6 +124,54 @@ Granular extras: `[proxy]`, `[mcp]`, `[ml]`, `[code]`, `[memory]`, `[relevance]`
|
|||
|
||||
Reproduce: `python -m headroom.evals suite --tier 1` · [Full benchmarks & methodology](https://headroom-docs.vercel.app/docs/benchmarks)
|
||||
|
||||
## Output token reduction (cut what the model writes back)
|
||||
|
||||
Everything above shrinks the prompt you **send**. But you also pay for every
|
||||
token the model **writes back** — and on Opus-class models output costs 5× input.
|
||||
A lot of that output is waste: "Great, let me…" preambles, re-printing code you
|
||||
just showed it, and deep "thinking" on routine steps like reading a file.
|
||||
|
||||
Headroom can trim that too, from the proxy, without you changing any code:
|
||||
|
||||
- **Verbosity steering** — appends a short "be terse, don't restate context"
|
||||
note to the end of the system prompt (so your prompt cache still hits).
|
||||
- **Effort routing** — when a turn is just the model resuming after a tool result
|
||||
(a file read, a passing test), it dials the model's thinking effort down. New
|
||||
questions and errors keep full effort.
|
||||
|
||||
Turn it on:
|
||||
|
||||
```bash
|
||||
export HEADROOM_OUTPUT_SHAPER=1 # off by default
|
||||
headroom proxy --port 8787
|
||||
```
|
||||
|
||||
**Learn the right terseness for you.** People don't *say* how terse they want
|
||||
answers — they *show* it (they interrupt long replies, or move on before they
|
||||
could have read them). `headroom learn --verbosity` reads your past sessions and
|
||||
picks the level automatically:
|
||||
|
||||
```bash
|
||||
headroom learn --verbosity # preview what it found (dry run)
|
||||
headroom learn --verbosity --apply # save it; the proxy uses it from now on
|
||||
```
|
||||
|
||||
**See how many output tokens you saved.** Output savings are *counterfactual* —
|
||||
we never see what the model *would* have written — so Headroom reports an honest
|
||||
**estimate with a confidence range**, never a made-up number:
|
||||
|
||||
```bash
|
||||
headroom output-savings
|
||||
# Reduction: 31.7% (95% CI 27.7% … 35.7%) [estimated]
|
||||
```
|
||||
|
||||
Want a *measured* number instead of an estimate? Leave 10% of conversations
|
||||
unshaped as a control group: `export HEADROOM_OUTPUT_HOLDOUT=0.1`. The dashboard
|
||||
shows an **Output Tokens Saved** card next to input compression, labelled
|
||||
`measured` or `estimated` with the confidence band.
|
||||
|
||||
→ Full write-up incl. the measurement methodology: [`docs/proposals/output-token-reduction.md`](docs/proposals/output-token-reduction.md)
|
||||
|
||||
<a href="https://www.star-history.com/?repos=chopratejas%2Fheadroom&type=date&legend=top-left">
|
||||
<picture>
|
||||
<img alt="Star History Chart" src="https://api.star-history.com/chart?repos=chopratejas/headroom&type=date&legend=top-left" />
|
||||
|
|
|
|||
136
docs/output-token-reduction-guide.md
Normal file
136
docs/output-token-reduction-guide.md
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
# Output Token Reduction — User Guide
|
||||
|
||||
A plain-English guide to cutting the tokens the model **writes back**.
|
||||
|
||||
## Why this exists
|
||||
|
||||
Headroom normally shrinks the prompt you **send**. This feature shrinks what the
|
||||
model **returns**. That matters because:
|
||||
|
||||
- Output tokens cost **5× more** than input on Opus-class models.
|
||||
- A lot of model output is waste: "Great, let me help with that…" intros,
|
||||
re-printing code you already showed it, restating tool results, and long
|
||||
internal "thinking" even on trivial steps.
|
||||
|
||||
You don't change any code. It runs in the Headroom proxy.
|
||||
|
||||
## Turn it on
|
||||
|
||||
```bash
|
||||
export HEADROOM_OUTPUT_SHAPER=1 # off by default
|
||||
headroom proxy --port 8787
|
||||
```
|
||||
|
||||
That's it. Two things now happen on every request:
|
||||
|
||||
1. **Verbosity steering** — a short "be terse, don't restate context" instruction
|
||||
is added to the **end** of the system prompt. (The end, so your prompt cache
|
||||
still works.)
|
||||
2. **Effort routing** — if a turn is just the model continuing after a tool ran
|
||||
(e.g. it read a file and there were no errors), Headroom turns the model's
|
||||
"thinking effort" down for that one turn. Real questions and error-handling
|
||||
turns keep full effort.
|
||||
|
||||
## The verbosity dial (levels 0–4)
|
||||
|
||||
| Level | What the model is told | Good for |
|
||||
|------:|------------------------|----------|
|
||||
| 0 | (off) | disable steering |
|
||||
| 1 | Skip the intro/outro chit-chat | people who read everything |
|
||||
| 2 | Also: don't restate code/output already on screen | **default** — safe |
|
||||
| 3 | Also: conclusions only, skip the reasoning | people who skim |
|
||||
| 4 | Bare minimum, fragments OK | maximum savings, terse |
|
||||
|
||||
Set it by hand if you want:
|
||||
|
||||
```bash
|
||||
export HEADROOM_VERBOSITY_LEVEL=3
|
||||
```
|
||||
|
||||
Or — better — let Headroom learn it from your habits (next section).
|
||||
|
||||
## Let Headroom pick the level for you
|
||||
|
||||
People rarely *say* "be brief." They *show* it: they interrupt long answers, or
|
||||
reply so fast they couldn't have read the whole thing. `headroom learn
|
||||
--verbosity` reads your past sessions and picks a level from those signals.
|
||||
|
||||
```bash
|
||||
# Preview what it found (doesn't change anything)
|
||||
headroom learn --verbosity
|
||||
|
||||
# Save it — the proxy uses this level from now on
|
||||
headroom learn --verbosity --apply
|
||||
```
|
||||
|
||||
Example output:
|
||||
|
||||
```
|
||||
Verbosity — headroom
|
||||
Interrupts: 29 (11% of turns) ← push-back signal
|
||||
Fast-skips: 31 / 119 long answers (26% unread) ← strongest signal
|
||||
>> Recommended verbosity level: 3 (confidence: high)
|
||||
```
|
||||
|
||||
Add `--llm-judge` to have an LLM double-check the level (needs an API key).
|
||||
|
||||
## See how much you saved
|
||||
|
||||
Here's the honest part. We **can't directly measure** output savings — we never
|
||||
see what the model *would* have written without our nudge. So Headroom reports an
|
||||
**estimate with a confidence range**, never a fake exact number:
|
||||
|
||||
```bash
|
||||
headroom output-savings
|
||||
```
|
||||
|
||||
```
|
||||
Output-token reduction
|
||||
Method: ESTIMATED (synthetic control)
|
||||
Requests: 1,240 shaped
|
||||
Saved: ~410,000 output tokens
|
||||
Reduction: 28.0% (95% CI 24.1% … 31.9%)
|
||||
```
|
||||
|
||||
- **ESTIMATED** = compared against a baseline of your past (unshaped) sessions.
|
||||
- **MEASURED** = the gold standard, if you opt into a holdout (below).
|
||||
|
||||
### Want a *measured* number?
|
||||
|
||||
Leave a slice of traffic unshaped as a control group:
|
||||
|
||||
```bash
|
||||
export HEADROOM_OUTPUT_HOLDOUT=0.1 # 10% of conversations stay unshaped
|
||||
```
|
||||
|
||||
Now `headroom output-savings` compares shaped vs unshaped directly and reports a
|
||||
**measured** reduction. The trade-off: you give up the savings on that 10%.
|
||||
|
||||
## On the dashboard
|
||||
|
||||
Open `http://localhost:8787/dashboard`. Next to the input-compression card
|
||||
you'll see an **Output Tokens Saved** card showing the token count, the percent,
|
||||
a `measured`/`estimated` badge, and the confidence range.
|
||||
|
||||
## FAQ
|
||||
|
||||
**Will this make answers worse?**
|
||||
At level 2 (default), no — in our tests the model finds the same bugs and writes
|
||||
the same fixes; it just stops re-printing code and skipping the "let me…" intro.
|
||||
Levels 3–4 are terser by design; that's why learning the level per user matters.
|
||||
|
||||
**Does it break prompt caching?**
|
||||
No. The steering text is added at the *end* of the system prompt and is
|
||||
byte-stable, so your cached prefix is untouched.
|
||||
|
||||
**Is it safe with extended thinking / tool loops?**
|
||||
Yes. It never disables thinking outright (that can error), it only lowers effort
|
||||
on routine turns, and it never adds settings the model doesn't support.
|
||||
|
||||
**How do I turn it off?**
|
||||
Unset `HEADROOM_OUTPUT_SHAPER` (or set it to `0`) and restart the proxy. You can
|
||||
also send `x-headroom-bypass: true` on a request to skip it for that call.
|
||||
|
||||
---
|
||||
|
||||
Deep dive (design + the counterfactual math): [`proposals/output-token-reduction.md`](proposals/output-token-reduction.md)
|
||||
461
docs/proposals/output-token-reduction.md
Normal file
461
docs/proposals/output-token-reduction.md
Normal file
|
|
@ -0,0 +1,461 @@
|
|||
# Output Token Reduction
|
||||
|
||||
**Branches:** `feat/output-token-reduction` (Phase 1) → `feat/verbosity-learning-and-counterfactual` (Phase 2).
|
||||
**Status:** Phase 1 (output shaper) and Phase 2 (`learn --verbosity`, AIMD controller, counterfactual estimator, dashboard) both implemented + tested. Runtime AIMD signal-capture is staged (controller built/tested, live emission off by default).
|
||||
|
||||
See **§7** for the counterfactual measurement methodology (how we honestly report a number we can't directly observe).
|
||||
|
||||
---
|
||||
|
||||
## 1. The problem in one line
|
||||
|
||||
Headroom's entire transform pipeline compresses what goes **into** the model.
|
||||
Nothing today touches what comes **out**. But output tokens are billed at
|
||||
5× input on Opus-class models ($25 vs $5 per MTok on `claude-opus-4-8`), and in
|
||||
agentic coding loops a large fraction of the bill is output: thinking tokens,
|
||||
restated code, ceremony ("Great, let me…"), and full-file rewrites where a
|
||||
10-line edit would do.
|
||||
|
||||
**The constraint that shapes everything:** the proxy never generates output
|
||||
tokens — the model does. Once a token is streamed it is already billed. So
|
||||
every output-token lever is **request-side**: change what we ask for, cap what
|
||||
we allow, or avoid the generation entirely. There is no post-hoc lever.
|
||||
|
||||
That gives three lever families plus a learning loop:
|
||||
|
||||
| Lever | Mechanism | Status |
|
||||
|---|---|---|
|
||||
| **Verbosity steering** | Append a terse-style instruction to the system-prompt tail | ✅ built |
|
||||
| **Effort routing** | Lower `output_config.effort` on mechanical turns | ✅ built |
|
||||
| **Thinking-budget clamp** | Clamp legacy `thinking.budget_tokens` on mechanical turns | ✅ built |
|
||||
| **Per-user learned level** | Mine past sessions for the right verbosity per user (`learn --verbosity`) | ✅ built |
|
||||
| **Counterfactual estimator** | Honestly estimate output tokens saved + dashboard surfacing | ✅ built |
|
||||
| **Runtime AIMD auto-tune** | Adjust level live from interrupt / skip signals | 🟡 controller built/tested; live signal emission off by default |
|
||||
|
||||
---
|
||||
|
||||
## 2. Phase 1 — the output shaper (built)
|
||||
|
||||
### 2.1 What it is
|
||||
|
||||
`headroom/proxy/output_shaper.py` — a request-body rewriter invoked in
|
||||
`handle_anthropic_messages` after every other body mutation (so the turn
|
||||
classifier sees the final message list) and gated behind the same
|
||||
`x-headroom-bypass` header as compression. Opt-in via `HEADROOM_OUTPUT_SHAPER=1`.
|
||||
|
||||
### 2.2 Lever A — verbosity steering
|
||||
|
||||
A deterministic instruction block is appended to the **tail** of the system
|
||||
prompt. Five levels, cumulative:
|
||||
|
||||
- **L0** — off (touch nothing).
|
||||
- **L1** — no ceremony: skip preamble/postamble, don't announce what you're about to do.
|
||||
- **L2** — L1 + no echo: never restate code/diffs/tool output already in context; reference by path:line; don't narrate tool results. **(default)**
|
||||
- **L3** — L2 + conclusions only, omit rationale unless asked, prefer smallest edit.
|
||||
- **L4** — caveman: fragments, minimum tokens, nothing but the answer.
|
||||
|
||||
**Why the tail, not the head.** Prompt caching is a prefix match — any byte
|
||||
change ahead of a `cache_control` breakpoint invalidates everything after it.
|
||||
Prepending steering text would bust the provider prefix cache and cost more
|
||||
than it saves. Appending after the last system block leaves the cached prefix
|
||||
byte-identical; only the small, byte-stable steering block is reprocessed. The
|
||||
steering text is frozen per level and applied idempotently (sentinel-tagged),
|
||||
so repeated requests keep an identical prefix and a level change replaces the
|
||||
block in place rather than stacking.
|
||||
|
||||
### 2.3 Lever B — effort routing
|
||||
|
||||
In an agentic loop, most API calls are **mechanical continuations**: the last
|
||||
message is a clean `tool_result` (a file read, a passing test) and the model is
|
||||
just resuming. Harnesses like Claude Code pin `output_config.effort` at `xhigh`
|
||||
for *every* turn, including these — and effort drives thinking depth, which
|
||||
bills as output. The router lowers effort to `low` on mechanical turns only.
|
||||
|
||||
Turn classification is **purely structural** — no content regexes, no keyword
|
||||
lists (per the project's no-hardcoded-patterns rule):
|
||||
|
||||
| Last user message contains… | Classification | Action |
|
||||
|---|---|---|
|
||||
| Any text / image / document block | `NEW_USER_ASK` | leave effort alone |
|
||||
| Only `tool_result`, none `is_error` | `MECHANICAL_CONTINUATION` | lower effort → `low` |
|
||||
| Any `tool_result` with `is_error: true` | `ERROR_CONTINUATION` | leave effort alone (model must reason about the failure) |
|
||||
| Anything else | `UNKNOWN` | leave alone |
|
||||
|
||||
### 2.4 Lever C — legacy thinking-budget clamp
|
||||
|
||||
On older models still sending `thinking: {type: "enabled", budget_tokens: N}`,
|
||||
the router clamps `N` to the API floor (1024) on mechanical turns. The `type`
|
||||
field is **never** toggled.
|
||||
|
||||
### 2.5 Safety rules (each prevents a concrete failure)
|
||||
|
||||
1. **Never inject `output_config.effort` where the client didn't send it.**
|
||||
Models without effort support return 400 on it. Lowering an
|
||||
already-present value is always valid — its presence proves the target
|
||||
model accepts the param.
|
||||
2. **Never toggle `thinking.type`.** Disabling thinking while history carries
|
||||
thinking blocks 400s on some models, and the toggle busts the messages
|
||||
cache tier (per the caching invalidation hierarchy).
|
||||
3. **Byte-stable, idempotent steering** — repeated requests keep an identical
|
||||
prefix; cache stays warm.
|
||||
4. **Respect `x-headroom-bypass`** — sub-agent calls that opt out of
|
||||
compression also opt out of shaping.
|
||||
|
||||
### 2.6 Configuration
|
||||
|
||||
| Env var | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `HEADROOM_OUTPUT_SHAPER` | off | master switch (`1`/`true`/`yes`) |
|
||||
| `HEADROOM_VERBOSITY_LEVEL` | `2` | 0–4 (clamped) |
|
||||
| `HEADROOM_EFFORT_ROUTER` | on | set `0` to disable effort routing |
|
||||
| `HEADROOM_MECHANICAL_EFFORT` | `low` | floor effort for mechanical turns |
|
||||
|
||||
### 2.7 Tests
|
||||
|
||||
`tests/test_output_shaper.py` — 34 tests, all passing. Covers turn
|
||||
classification (every block-type path), cache-safe steering (string→block
|
||||
conversion, append-after-`cache_control`, idempotency, level change),
|
||||
effort routing (lower / never-inject / untouched-on-non-mechanical /
|
||||
configurable target), legacy budget clamp, and the env gate. Ruff + mypy clean.
|
||||
|
||||
---
|
||||
|
||||
## 3. Live before/after results
|
||||
|
||||
Measured against `claude-opus-4-8` via `scripts/eval_output_shaper.py`,
|
||||
comparing the exact body a client sends vs. the body the proxy forwards.
|
||||
`usage.output_tokens` includes thinking.
|
||||
|
||||
### 3.1 Verbosity steering — complex code-review ask
|
||||
|
||||
Prompt: "review this TTLCache, find every bug, show fixes."
|
||||
|
||||
| Condition | Mean output tokens | Reduction |
|
||||
|---|---|---|
|
||||
| Baseline | ~1,800–1,930 | — |
|
||||
| L2 (no ceremony, no echo) | ~1,180–1,470 | **−22% to −39%** |
|
||||
| L3 (conclusions only) | ~670 | **−63%** |
|
||||
|
||||
**Quality check — same bugs found, only redundancy removed.** The baseline
|
||||
opens with a title + "Let me go through them," finds the bugs, **re-prints the
|
||||
entire fixed class**, then adds a **summary table restating all six bugs**,
|
||||
then a trailing notes section — the same information appears ~2.5×. L2 opens
|
||||
directly with the findings, gives the same fixed code block once, and stops.
|
||||
Both correctly identify the no-locking race, `popitem(last=True)` evicting the
|
||||
wrong end, and the mutate-during-iteration `RuntimeError`. Nothing of substance
|
||||
is lost at L2; L3 additionally drops rationale prose (only for users who don't
|
||||
read explanations).
|
||||
|
||||
### 3.2 Effort routing — agentic mechanical continuation
|
||||
|
||||
Transcript ending in a clean `tool_result`, `effort: xhigh` the way Claude Code
|
||||
sends it.
|
||||
|
||||
| Condition | Mean output tokens | Reduction |
|
||||
|---|---|---|
|
||||
| Baseline (xhigh) | ~1,120 | — |
|
||||
| Shaped (routed to low) | 793 | **−29%** |
|
||||
|
||||
### 3.3 Honest framing
|
||||
|
||||
Caveman (L4) would get ~60–70% but degrades the experience. The taxonomy-driven
|
||||
default (L2 + effort routing on mechanical turns) gets a realistic **25–40%**
|
||||
with the user barely noticing. The learning loop (Phase 2) finds where each
|
||||
user sits between those poles.
|
||||
|
||||
---
|
||||
|
||||
## 4. Phase 2 — learning the right level per user
|
||||
|
||||
### 4.1 Why this is necessary
|
||||
|
||||
The right verbosity is **per-user**, not global. A fixed `HEADROOM_VERBOSITY_LEVEL`
|
||||
is a guess. We can do better: mine the user's own past sessions to infer what
|
||||
they actually tolerate — exactly the philosophy of `headroom learn`, which
|
||||
already reads `~/.claude/projects/*.jsonl` and turns history into learned
|
||||
context.
|
||||
|
||||
### 4.2 The key insight (validated on real data)
|
||||
|
||||
I prototyped the signal extraction and ran it over **24 real sessions** for
|
||||
this project. The finding that shapes the whole design:
|
||||
|
||||
| Signal | This user's data | Carries signal? |
|
||||
|---|---|---|
|
||||
| Explicit "be brief" keywords | **1** of 210 human msgs | ❌ almost none |
|
||||
| Explicit "explain more" keywords | **1** of 210 | ❌ almost none |
|
||||
| **Interruptions** (user cuts Claude off) | **29** (~1 per 7 turns) | ✅ strong |
|
||||
| **Fast-skips** (reply <30s after a >250-word answer) | **15 of 100 long outputs** | ✅ strong |
|
||||
| Long-output frequency | 100 outputs >250 words | ✅ context |
|
||||
|
||||
**Users almost never *say* what verbosity they want — they show it
|
||||
behaviorally.** Keyword matching (the obvious approach) is nearly empty. The
|
||||
behavioral signals are rich. The strongest is **fast-skip**: a reply arriving
|
||||
faster than the answer could have been read is a direct measurement of
|
||||
"generated tokens nobody consumed," computable from timestamps + lengths
|
||||
already in the JSONL.
|
||||
|
||||
For this user — 29 interrupts, 15% unread-long-output rate, zero "explain more"
|
||||
requests — the data reads as a clear **L2, arguably L3** user.
|
||||
|
||||
### 4.3 Design — `headroom learn --verbosity`
|
||||
|
||||
Slots into the existing `learn` architecture. `headroom/learn/plugins/claude.py`
|
||||
already parses the JSONL; `analyzer.py` already does cheap-extraction → digest →
|
||||
LLM-returns-JSON. We add a verbosity analysis path:
|
||||
|
||||
**Step 1 — structural pass (pure Python, no patterns).** Per session, compute:
|
||||
- **interrupt rate** — `[Request interrupted by user…]` markers per human turn
|
||||
- **fast-skip rate** — human reply latency vs. preceding assistant output length
|
||||
- **long-output frequency** — share of assistant texts over N words
|
||||
- **echo ratio** — n-gram overlap between assistant output and prior context (restated code)
|
||||
- reply-latency distribution vs. output length
|
||||
|
||||
All mechanical, all from data already on disk. (Prototype:
|
||||
`scripts/verbosity_scan.py`, validated above.)
|
||||
|
||||
**Step 2 — LLM judgment pass (mirrors today's `learn`).** Feed a digest of
|
||||
human messages + the structural stats to the analyzer LLM with one question:
|
||||
*"Does this user read long explanations, and when they push back, is it for
|
||||
more detail or less?"* This replaces brittle keyword lists with judgment —
|
||||
consistent with the no-hardcoded-patterns rule. Returns a level + confidence +
|
||||
rationale as JSON.
|
||||
|
||||
**Step 3 — output (this is what `--verbosity` produces).** See §4.4.
|
||||
|
||||
**Step 4 — runtime AIMD auto-tune.** The offline pass sets the *starting* level;
|
||||
a live loop tracks drift. An interrupt or fast-skip nudges the level up one; an
|
||||
"explain"/"why" follow-up drops it back. Hysteresis (require 2–3 consistent
|
||||
signals before moving) prevents oscillation. Like TCP congestion control: probe
|
||||
toward terser, back off on a "too terse" signal.
|
||||
|
||||
### 4.4 What `--verbosity` outputs — and how it helps
|
||||
|
||||
The command produces three concrete artifacts:
|
||||
|
||||
**(a) A human-readable report (stdout):**
|
||||
```
|
||||
Verbosity analysis — /Users/tcms/demo/headroom (24 sessions, 210 turns)
|
||||
|
||||
Interrupts: 29 (1 per 7.2 turns) ← strong "too much" signal
|
||||
Fast-skips: 15 / 100 long outputs ← 15% of long answers unread
|
||||
Explicit brevity: 1 explicit verbose: 1 ← behavioral, not stated
|
||||
|
||||
LLM read: "User interrupts frequently and rarely reads long
|
||||
explanations; pushes back for less, never more detail."
|
||||
|
||||
Recommended verbosity level: 2 (confidence: high)
|
||||
Estimated output-token reduction at L2: ~25–35%
|
||||
```
|
||||
|
||||
**(b) A persisted setting** written to `~/.headroom/` (alongside the existing
|
||||
savings tracker) — per-project verbosity level + confidence:
|
||||
```json
|
||||
{"project": "/Users/tcms/demo/headroom",
|
||||
"verbosity_level": 2, "confidence": "high",
|
||||
"signals": {"interrupt_rate": 0.138, "fast_skip_rate": 0.15},
|
||||
"learned_at": "2026-06-12T…"}
|
||||
```
|
||||
|
||||
**(c) The shaper reads it as its default.** `OutputShaperSettings.from_env()`
|
||||
gains a fallback: if `HEADROOM_VERBOSITY_LEVEL` is unset, load the learned
|
||||
per-project level instead of the hardcoded `2`. So **the output of `--verbosity`
|
||||
directly becomes the live verbosity the proxy applies** — no manual tuning.
|
||||
|
||||
**How it helps, concretely:**
|
||||
1. **Removes the guess.** Today you set `HEADROOM_VERBOSITY_LEVEL=2` by hand.
|
||||
After `learn --verbosity`, the level is derived from *your* behavior — a
|
||||
heavy-interrupter gets L3, a "read everything" user gets L1.
|
||||
2. **Per-project, not global.** Your exploratory side-project and your
|
||||
production repo can carry different levels.
|
||||
3. **Justified, not magic.** The report shows the signals and the LLM's read,
|
||||
so the recommendation is auditable (matches the dashboard philosophy of
|
||||
showing directional data, not opaque scores).
|
||||
4. **Seeds the runtime loop.** The learned level is the AIMD starting point;
|
||||
live signals refine it without re-running the offline pass.
|
||||
|
||||
### 4.5 Mapping signals → level (initial heuristic, LLM-refined)
|
||||
|
||||
| Interrupt rate | Fast-skip rate | "explain more" present | → Level |
|
||||
|---|---|---|---|
|
||||
| low | low | yes | 1 |
|
||||
| low–med | low–med | no | 2 |
|
||||
| high | high | no | 3 |
|
||||
| very high | very high | no | 4 (offer, don't auto-apply) |
|
||||
|
||||
The LLM judgment pass adjusts this — the table is the prior, not the verdict.
|
||||
|
||||
---
|
||||
|
||||
## 5. Files
|
||||
|
||||
| File | Status | Purpose |
|
||||
|---|---|---|
|
||||
| `headroom/proxy/output_shaper.py` | ✅ built | the shaper (steering + effort routing) |
|
||||
| `headroom/proxy/handlers/anthropic.py` | ✅ wired | invoke shaper after body mutations |
|
||||
| `tests/test_output_shaper.py` | ✅ 34 passing | unit coverage |
|
||||
| `scripts/eval_output_shaper.py` | ✅ built | live before/after eval |
|
||||
| `scripts/verbosity_scan.py` | 🔬 prototype | session-mining signal extraction |
|
||||
| `headroom/learn/plugins/claude.py` (+ analyzer) | 🔜 extend | `--verbosity` analysis path |
|
||||
| `~/.headroom/verbosity.json` | 🔜 | persisted per-project learned level |
|
||||
|
||||
---
|
||||
|
||||
## 6. Roadmap
|
||||
|
||||
1. ✅ **Measure** — live eval establishes the realistic ceiling and baseline.
|
||||
2. ✅ **Effort router** — biggest win per unit risk, fully mechanical, invisible.
|
||||
3. ✅ **Verbosity ladder at fixed L2** — safe default, cache-safe tail injection.
|
||||
4. 🔜 **`learn --verbosity`** — derive the per-user starting level from sessions.
|
||||
5. 🔜 **Runtime AIMD auto-tune** — refine the level live from interrupt/skip signals.
|
||||
6. 🔭 **Waste taxonomy on the dashboard** — echo ratio, ceremony ratio,
|
||||
full-file-rewrite detection, as token counts (no dollar estimates).
|
||||
7. 🔭 **Budget whispering** — tell the model its token budget per turn-type,
|
||||
sized from the historical output distribution in SQLite.
|
||||
|
||||
---
|
||||
|
||||
## 7. Counterfactual measurement — how we show a % we can't directly observe
|
||||
|
||||
This is the hard part, and it deserves its own section.
|
||||
|
||||
### 7.1 Why output savings are not directly measurable
|
||||
|
||||
Input compression is a **pure function**: Headroom takes a request, shrinks it,
|
||||
and can count `tokens_before` and `tokens_after` — both are observable on the
|
||||
same request. Output is different. When the shaper makes a request terser, the
|
||||
model emits N output tokens. We **never observe** what it *would* have emitted
|
||||
without the steering. Only one side of the counterfactual ever happens. So a
|
||||
flat "we saved 30%" is a guess dressed as a fact.
|
||||
|
||||
The design rule that follows: **never report a single number as if it were
|
||||
measured.** Separate what is genuinely measured from what is estimated, label
|
||||
each, and always attach uncertainty.
|
||||
|
||||
### 7.2 Three tiers of honesty
|
||||
|
||||
**Tier 1 — Estimated (synthetic control).** Build a per-stratum baseline of
|
||||
*unshaped* output tokens from session history that predates the shaper
|
||||
(`learn --verbosity` does this in the same pass that picks the level). For each
|
||||
shaped request, the expected unshaped output is the baseline mean for that
|
||||
request's stratum. Aggregate estimate:
|
||||
|
||||
```
|
||||
tokens_saved = Σ over shaped requests ( baseline_mean[stratum] − observed_output )
|
||||
```
|
||||
|
||||
Summed as **signed** deltas — never clamped per-request. Clamping each delta at
|
||||
zero would throw away the cases where a shaped response happened to be *longer*
|
||||
by chance, biasing the total upward. Over many requests the noise averages out;
|
||||
the systematic effect remains. Reported with a propagated 95% CI (see §7.5) and
|
||||
always labelled "estimated."
|
||||
|
||||
**Tier 2 — Measured (A/B holdout).** Set `HEADROOM_OUTPUT_HOLDOUT=0.1` and 10%
|
||||
of conversations are deliberately left **unshaped** as a control arm. Within
|
||||
each stratum, `mean(control) − mean(treatment)` is an **unbiased causal
|
||||
estimate** of the per-request saving. This is the only number we call
|
||||
"measured." It self-corrects: if steering doesn't actually help on some
|
||||
workload, the holdout reveals it. The cost is real (you forgo savings on 10% of
|
||||
traffic), so it's opt-in; default holdout is 0 (estimate-only).
|
||||
|
||||
**Tier 3 — Direct waste (no counterfactual at all).** Echo ratio — the n-gram
|
||||
overlap between a response and the context it was given — is a property of a
|
||||
single response. "32% of this output restated context already on screen" needs
|
||||
no counterfactual; it's a measured fact about output we *did* see, and it's
|
||||
exactly what the shaper targets. Surfaced in `learn --verbosity` as
|
||||
`mean_echo_ratio`.
|
||||
|
||||
### 7.3 Stratification — comparing like with like
|
||||
|
||||
You can't compare a "fix this typo" response to "design a caching layer." The
|
||||
estimator buckets every request by features observable **before** the response:
|
||||
|
||||
```
|
||||
stratum = model_family | turn_kind | input_token_bucket | has_tools
|
||||
= e.g. "opus | mechanical_continuation | xl | tools"
|
||||
```
|
||||
|
||||
Coarse on purpose (~25–50 strata) so per-stratum baselines stay dense. The live
|
||||
proxy computes the stratum the exact same way the offline baseline does, so
|
||||
treatment requests line up with their baseline. Unseen strata fall back
|
||||
hierarchically (drop `has_tools`, then the bucket, …, then the global mean).
|
||||
|
||||
### 7.4 The two constraints that happen to align
|
||||
|
||||
Holdout assignment is **conversation-stable** — a whole conversation is either
|
||||
treatment or control, decided by hashing a conversation-stable key (model +
|
||||
first user message). This matters for two independent reasons that point the
|
||||
same way:
|
||||
|
||||
1. **Measurement validity** — mixing shaped and unshaped turns within one
|
||||
conversation would contaminate the comparison (the history itself differs).
|
||||
2. **Cache safety** — flipping a conversation's verbosity mid-stream changes the
|
||||
system-prompt tail, which busts the provider prefix cache.
|
||||
|
||||
So the same rule (assign per conversation, never per turn) is forced by both
|
||||
the statistics and the caching. Nice when constraints agree.
|
||||
|
||||
### 7.5 The confidence interval
|
||||
|
||||
For the estimated tier, uncertainty comes from two sources, both propagated:
|
||||
|
||||
```
|
||||
Var(tokens_saved) ≈ Σ_s [ n_s · σ²_observed,s + n_s² · σ²_baseline,s / m_s ]
|
||||
└─ spread of shaped outputs ─┘ └─ finite-baseline error ─┘
|
||||
```
|
||||
|
||||
where `n_s` is treatment count in stratum `s` and `m_s` the baseline sample
|
||||
count. For the measured tier it's the standard difference-of-means variance
|
||||
`σ²_c/n_c + σ²_t/n_t` per stratum. The 95% band is `point ± 1.96·√Var`, surfaced
|
||||
everywhere the number is (CLI and dashboard), so the reader sees the precision,
|
||||
not just a point estimate.
|
||||
|
||||
Output-token counts are right-skewed (a few huge responses). Means are still the
|
||||
right statistic for *totals* (you're billed on the sum), but the CI widens
|
||||
honestly when a stratum is dominated by a few large responses — which is the
|
||||
correct signal that the estimate is soft there.
|
||||
|
||||
### 7.6 How it flows end to end
|
||||
|
||||
```
|
||||
learn --verbosity --apply
|
||||
├─ writes verbosity.json (the level the shaper applies)
|
||||
└─ seeds output_savings.json (the per-stratum baseline = synthetic control)
|
||||
|
||||
proxy request (HEADROOM_OUTPUT_SHAPER=1)
|
||||
├─ assign_arm(conversation) → treatment | control (holdout)
|
||||
├─ stratum_key(request features)
|
||||
├─ treatment: shape body; control: leave unshaped
|
||||
└─ tag (arm, stratum) onto transforms_applied ← rides existing plumbing
|
||||
|
||||
response completes → emit_request_outcome (one funnel, all paths)
|
||||
└─ recorder.record(arm, stratum, output_tokens) → output_savings.json
|
||||
|
||||
headroom output-savings / dashboard "Output Tokens Saved" card
|
||||
└─ best_estimate(): measured if a holdout exists, else estimated; with CI
|
||||
```
|
||||
|
||||
The recording rides the existing `transforms_applied` label channel, so it
|
||||
works for streaming, non-streaming, and backend paths with no change to
|
||||
`RequestOutcome` or its construction sites.
|
||||
|
||||
### 7.7 What the user sees
|
||||
|
||||
- **CLI:** `headroom output-savings` →
|
||||
`Reduction: 31.7% (95% CI 27.7% … 35.7%) [MEASURED, 400 shaped requests]`
|
||||
- **Dashboard:** an "Output Tokens Saved" hero card next to input compression —
|
||||
token count, percent, a `measured`/`estimated` badge, and the CI band.
|
||||
- **No dollar estimates** on the output card (per project convention) — token
|
||||
counts and directional percentages only.
|
||||
|
||||
### 7.8 Honest limitations
|
||||
|
||||
- Estimated-tier accuracy depends on the baseline matching current workload; if
|
||||
your tasks drift, re-run `learn --verbosity` or turn on a small holdout.
|
||||
- The baseline must come from *unshaped* history. If you learn from sessions
|
||||
where the shaper was already active, the baseline is contaminated — the live
|
||||
holdout is the clean path forward.
|
||||
- Runtime AIMD upward-ratcheting is gated off by default: we can reliably detect
|
||||
"too much output" (fast-skip timing, stream cancellation) but not yet "too
|
||||
little" at runtime without content heuristics, so auto-escalation stays
|
||||
behind `HEADROOM_VERBOSITY_AUTOTUNE` until both directions are trustworthy.
|
||||
|
|
@ -3,7 +3,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import click
|
||||
|
||||
|
|
@ -105,6 +105,21 @@ Use 'auto' (default) to scan all detected agents."""
|
|||
help="Only scan top-level main sessions, skipping nested subagent/workflow "
|
||||
"transcripts (Claude Code). Default scans everything.",
|
||||
)
|
||||
@click.option(
|
||||
"--verbosity",
|
||||
"verbosity_mode",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Learn the user's preferred OUTPUT verbosity from behavioral signals "
|
||||
"(interrupts, fast-skips) instead of analyzing failures. Writes the level "
|
||||
"the output shaper applies, and seeds the savings baseline. --apply persists.",
|
||||
)
|
||||
@click.option(
|
||||
"--llm-judge",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="With --verbosity: let an LLM override the heuristic level (needs an API key).",
|
||||
)
|
||||
def learn(
|
||||
project: Path | None,
|
||||
analyze_all: bool,
|
||||
|
|
@ -113,6 +128,8 @@ def learn(
|
|||
model: str | None,
|
||||
workers: int | None,
|
||||
main_only: bool,
|
||||
verbosity_mode: bool,
|
||||
llm_judge: bool,
|
||||
) -> None:
|
||||
"""Learn from past tool call failures to prevent future ones.
|
||||
|
||||
|
|
@ -139,6 +156,19 @@ def learn(
|
|||
|
||||
max_workers = workers if workers is not None else min(os.cpu_count() or 4, 8)
|
||||
|
||||
# Verbosity learning is a distinct flow: it mines behavioral signals (no
|
||||
# failure analysis) and needs no LLM unless --llm-judge is set.
|
||||
if verbosity_mode:
|
||||
_run_verbosity(
|
||||
project=project,
|
||||
analyze_all=analyze_all,
|
||||
apply=apply,
|
||||
agent=agent,
|
||||
llm_judge=llm_judge,
|
||||
model=model,
|
||||
)
|
||||
return
|
||||
|
||||
# Resolve model early to fail fast with a clear message
|
||||
try:
|
||||
resolved_model = model or _detect_default_model()
|
||||
|
|
@ -271,3 +301,152 @@ def learn(
|
|||
f"Total: {total_projects} projects, {total_failures} failures, "
|
||||
f"{total_recommendations} recommendations"
|
||||
)
|
||||
|
||||
|
||||
def _make_llm_judge(model: str) -> Any:
|
||||
"""Build an LLM judge callable for verbosity, or None if unavailable.
|
||||
|
||||
The judge gets the behavioral signals and returns (level, rationale). Kept
|
||||
best-effort: any failure (no key, parse error) returns None so the caller
|
||||
falls back to the heuristic.
|
||||
"""
|
||||
|
||||
def judge(signals: dict) -> tuple[int, str] | None:
|
||||
try:
|
||||
import json
|
||||
|
||||
import litellm
|
||||
except ImportError:
|
||||
return None
|
||||
prompt = (
|
||||
"You tune how terse an AI coding assistant should be for one user, "
|
||||
"from their behavioral signals. Levels: 1=light (skip ceremony), "
|
||||
"2=no ceremony+no echo, 3=conclusions only, 4=caveman/fragments. "
|
||||
"Users who interrupt often and reply faster than an answer could be "
|
||||
"read (fast-skip) want LESS output.\n\n"
|
||||
f"Signals: {json.dumps(signals)}\n\n"
|
||||
'Return ONLY JSON: {"level": <1-4>, "rationale": "<one sentence>"}'
|
||||
)
|
||||
try:
|
||||
resp = litellm.completion(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
max_tokens=200,
|
||||
)
|
||||
text = resp["choices"][0]["message"]["content"]
|
||||
start, end = text.find("{"), text.rfind("}")
|
||||
data = json.loads(text[start : end + 1])
|
||||
return int(data["level"]), str(data.get("rationale", "LLM judgment"))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
return judge
|
||||
|
||||
|
||||
def _run_verbosity(
|
||||
*,
|
||||
project: Path | None,
|
||||
analyze_all: bool,
|
||||
apply: bool,
|
||||
agent: str,
|
||||
llm_judge: bool,
|
||||
model: str | None,
|
||||
) -> None:
|
||||
"""Learn preferred output verbosity from session transcripts."""
|
||||
from ..learn.registry import auto_detect_plugins, get_plugin
|
||||
from ..learn.verbosity import analyze
|
||||
from ..paths import ensure_workspace_dir
|
||||
from ..proxy.output_savings import SavingsLedger
|
||||
|
||||
# Verbosity mining reads Claude Code transcripts; restrict to that plugin.
|
||||
if agent == "auto":
|
||||
plugins = [p for p in auto_detect_plugins() if p.name == "claude"]
|
||||
if not plugins:
|
||||
click.echo("Verbosity learning currently supports Claude Code transcripts only.")
|
||||
return
|
||||
plugin = plugins[0]
|
||||
else:
|
||||
plugin = get_plugin(agent)
|
||||
if plugin.name != "claude":
|
||||
click.echo("Verbosity learning currently supports Claude Code transcripts only.")
|
||||
return
|
||||
|
||||
all_projects = plugin.discover_projects()
|
||||
if not all_projects:
|
||||
click.echo("No Claude Code project data found.")
|
||||
return
|
||||
|
||||
if analyze_all:
|
||||
targets = all_projects
|
||||
elif project:
|
||||
resolved = project.resolve()
|
||||
targets = [p for p in all_projects if p.project_path == resolved]
|
||||
else:
|
||||
cwd = Path.cwd().resolve()
|
||||
targets = [p for p in all_projects if p.project_path == cwd]
|
||||
if not targets:
|
||||
for parent in cwd.parents:
|
||||
targets = [p for p in all_projects if p.project_path == parent]
|
||||
if targets:
|
||||
break
|
||||
if not targets:
|
||||
click.echo("No matching project. Try --all or --project <path>.")
|
||||
return
|
||||
|
||||
judge = _make_llm_judge(model or "claude-sonnet-4-6") if llm_judge else None
|
||||
|
||||
for proj in targets:
|
||||
session_paths = sorted(proj.data_path.glob("*.jsonl"))
|
||||
if not session_paths:
|
||||
continue
|
||||
profile, baseline = analyze(session_paths, str(proj.project_path), llm_judge=judge)
|
||||
sig = profile.signals
|
||||
|
||||
click.echo(f"\n{'=' * 60}")
|
||||
click.echo(f"Verbosity — {proj.name}")
|
||||
click.echo(f"Path: {proj.project_path}")
|
||||
click.echo(f"{'=' * 60}")
|
||||
click.echo(
|
||||
f" Sessions: {sig.get('sessions')} human turns: {sig.get('human_msgs')} "
|
||||
f"responses: {sig.get('asst_responses')}"
|
||||
)
|
||||
click.echo(
|
||||
f" Interrupts: {sig.get('interrupts')} "
|
||||
f"({sig.get('interrupt_rate', 0):.0%} of turns) "
|
||||
"← push-back signal"
|
||||
)
|
||||
click.echo(
|
||||
f" Fast-skips: {sig.get('fast_skips')} / {sig.get('skip_eligible')} long "
|
||||
f"answers ({sig.get('fast_skip_rate', 0):.0%} unread) ← strongest signal"
|
||||
)
|
||||
click.echo(f" Echo ratio: {sig.get('mean_echo_ratio', 0):.1%} of output restated context")
|
||||
click.echo(f"\n Source: {profile.source}")
|
||||
click.echo(f" {profile.rationale}")
|
||||
click.echo(
|
||||
f"\n >> Recommended verbosity level: {profile.level} "
|
||||
f"(confidence: {profile.confidence})"
|
||||
)
|
||||
|
||||
if apply:
|
||||
ws = ensure_workspace_dir()
|
||||
from datetime import datetime, timezone
|
||||
|
||||
profile.learned_at = datetime.now(timezone.utc).isoformat()
|
||||
profile.save(ws / "verbosity.json")
|
||||
# Seed the savings baseline: replace baseline, preserve any live
|
||||
# treatment/control already accumulated.
|
||||
ledger_path = ws / "output_savings.json"
|
||||
ledger = SavingsLedger.load(ledger_path)
|
||||
ledger.baseline = baseline
|
||||
ledger.save(ledger_path)
|
||||
click.echo(f"\n [WROTE] {ws / 'verbosity.json'} (level {profile.level})")
|
||||
click.echo(
|
||||
f" [WROTE] {ledger_path} (baseline: {baseline.total_samples} samples, "
|
||||
f"{len(baseline.strata)} strata)"
|
||||
)
|
||||
click.echo(
|
||||
"\n The output shaper now uses this level when "
|
||||
"HEADROOM_OUTPUT_SHAPER=1 and HEADROOM_VERBOSITY_LEVEL is unset."
|
||||
)
|
||||
else:
|
||||
click.echo("\n Dry run — use --apply to persist the level and baseline.")
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ def _register_commands() -> None:
|
|||
install, # noqa: F401
|
||||
learn, # noqa: F401
|
||||
mcp, # noqa: F401
|
||||
output_savings, # noqa: F401
|
||||
perf, # noqa: F401
|
||||
proxy, # noqa: F401
|
||||
tools, # noqa: F401
|
||||
|
|
|
|||
61
headroom/cli/output_savings.py
Normal file
61
headroom/cli/output_savings.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
"""CLI: show counterfactual output-token reduction."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import click
|
||||
|
||||
from .main import main
|
||||
|
||||
|
||||
@main.command(name="output-savings")
|
||||
def output_savings() -> None:
|
||||
"""Show estimated/measured output-token reduction from the shaper.
|
||||
|
||||
Output tokens are counterfactual — we never see what the model *would* have
|
||||
emitted unshaped. This reports the honest estimate:
|
||||
|
||||
* "measured" — from an A/B holdout (set HEADROOM_OUTPUT_HOLDOUT>0), the
|
||||
unbiased difference between unshaped and shaped arms.
|
||||
* "estimated" — synthetic control: shaped output vs. the per-stratum
|
||||
baseline learned by `headroom learn --verbosity`.
|
||||
|
||||
Both are shown with a 95% confidence band so the uncertainty is explicit.
|
||||
"""
|
||||
from ..paths import workspace_dir
|
||||
from ..proxy.output_savings import SavingsLedger
|
||||
|
||||
path = workspace_dir() / "output_savings.json"
|
||||
if not path.exists():
|
||||
click.echo("No output-savings data yet.")
|
||||
click.echo("Run `headroom learn --verbosity --apply` to seed the baseline,")
|
||||
click.echo("then enable the shaper (HEADROOM_OUTPUT_SHAPER=1) and send traffic.")
|
||||
return
|
||||
|
||||
ledger = SavingsLedger.load(path)
|
||||
est = ledger.best_estimate()
|
||||
|
||||
click.echo(f"\n{'=' * 56}")
|
||||
click.echo("Output-token reduction")
|
||||
click.echo(f"{'=' * 56}")
|
||||
if est.n_requests == 0:
|
||||
click.echo(" No shaped requests recorded yet.")
|
||||
click.echo(
|
||||
f" Baseline: {ledger.baseline.total_samples} samples, "
|
||||
f"{len(ledger.baseline.strata)} strata."
|
||||
)
|
||||
return
|
||||
|
||||
label = "MEASURED (A/B holdout)" if est.kind == "measured" else "ESTIMATED (synthetic control)"
|
||||
click.echo(f" Method: {label}")
|
||||
click.echo(f" Requests: {est.n_requests:,} shaped")
|
||||
click.echo(f" Baseline: {est.baseline_tokens:,.0f} output tokens expected")
|
||||
click.echo(f" Saved: {est.tokens_saved:,.0f} output tokens")
|
||||
click.echo(
|
||||
f" Reduction: {est.pct:.1f}% (95% CI {est.ci_low_pct:.1f}% … {est.ci_high_pct:.1f}%)"
|
||||
)
|
||||
if est.kind == "estimated":
|
||||
click.echo(
|
||||
"\n Note: estimated vs the learned baseline. For a measured number,"
|
||||
"\n set HEADROOM_OUTPUT_HOLDOUT=0.1 to leave 10% of conversations"
|
||||
"\n unshaped as a control arm."
|
||||
)
|
||||
|
|
@ -178,8 +178,8 @@
|
|||
<main class="p-6 max-w-7xl mx-auto">
|
||||
<template x-if="viewMode === 'session'">
|
||||
<div>
|
||||
<!-- Hero Metrics: Savings $ -> Compression % -> Overhead -->
|
||||
<div class="grid grid-cols-1 md:grid-cols-3 gap-4 mb-6">
|
||||
<!-- Hero Metrics: Savings $ -> Input Compression -> Output Reduction -> Overhead -->
|
||||
<div class="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-4 mb-6">
|
||||
<!-- Savings ($) - proxy compression only, priced at model list rate -->
|
||||
<div class="bg-surface rounded-lg p-4 border border-border">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Proxy $ Saved</div>
|
||||
|
|
@ -229,6 +229,36 @@
|
|||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Output Tokens Saved (counterfactual) -->
|
||||
<div class="bg-surface rounded-lg p-4 border border-border">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Output Tokens Saved</div>
|
||||
<template x-if="stats.tokens?.output_reduction?.available">
|
||||
<div>
|
||||
<div class="flex items-baseline gap-2">
|
||||
<span class="text-3xl font-light tabular-nums text-accent" x-text="formatNumber(stats.tokens?.output_saved || 0)"></span>
|
||||
<span class="text-sm text-accent" x-text="(stats.tokens?.output_reduction_percent || 0).toFixed(1) + '%'"></span>
|
||||
</div>
|
||||
<div class="mt-1 text-xs text-gray-500 leading-relaxed">
|
||||
<!-- "measured" = A/B holdout (unbiased); "estimated" = vs learned baseline -->
|
||||
<span class="uppercase tracking-wide"
|
||||
:class="stats.tokens?.output_reduction?.method === 'measured' ? 'text-emerald-400' : 'text-gray-400'"
|
||||
x-text="stats.tokens?.output_reduction?.method || ''"></span>
|
||||
<span x-text="'· 95% CI ' + (stats.tokens?.output_reduction?.ci_low_percent || 0).toFixed(1) + '–' + (stats.tokens?.output_reduction?.ci_high_percent || 0).toFixed(1) + '%'"></span>
|
||||
</div>
|
||||
<div class="mt-1 text-xs text-gray-600 leading-relaxed"
|
||||
x-text="formatNumber(stats.tokens?.output_reduction?.requests || 0) + ' shaped responses · counterfactual'">
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
<template x-if="!stats.tokens?.output_reduction?.available">
|
||||
<div class="mt-1 text-xs text-gray-500 leading-relaxed">
|
||||
<span class="text-2xl font-light tabular-nums text-gray-600">—</span>
|
||||
<div class="mt-1">Enable the output shaper (HEADROOM_OUTPUT_SHAPER=1) and run
|
||||
<code class="text-gray-400">headroom learn --verbosity --apply</code> to start measuring.</div>
|
||||
</div>
|
||||
</template>
|
||||
</div>
|
||||
|
||||
<!-- Headroom Overhead -->
|
||||
<div class="bg-surface rounded-lg p-4 border border-border">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Overhead</div>
|
||||
|
|
|
|||
473
headroom/learn/verbosity.py
Normal file
473
headroom/learn/verbosity.py
Normal file
|
|
@ -0,0 +1,473 @@
|
|||
"""Learn a user's preferred output verbosity from past sessions.
|
||||
|
||||
The premise (validated on real transcripts): users almost never *say* how terse
|
||||
they want answers — explicit "be brief" feedback is near-zero — but they *show*
|
||||
it behaviorally. They interrupt long answers, and they reply faster than a long
|
||||
answer could possibly have been read. Those signals are mechanical to extract
|
||||
from Claude Code's JSONL transcripts.
|
||||
|
||||
This module:
|
||||
|
||||
1. Parses transcripts into per-response records (output tokens, word count,
|
||||
timestamps, the preceding turn's structural kind).
|
||||
2. Extracts behavioral signals — interrupt rate, fast-skip rate, long-output
|
||||
frequency, echo ratio — using length-adaptive thresholds (a "fast skip" is
|
||||
defined relative to how long the answer would take to *read*, not a fixed
|
||||
number of seconds; "long" is relative to the user's own median).
|
||||
3. Recommends a verbosity level (heuristic prior; an optional LLM judgment pass
|
||||
can override it — see ``analyze``).
|
||||
4. Builds the per-stratum output-token baseline that
|
||||
:mod:`headroom.proxy.output_savings` uses as its synthetic control — so the
|
||||
same pass that picks the level also establishes how to measure its effect.
|
||||
|
||||
The structural signals are *inputs* to the decision, not the decision itself —
|
||||
which is why thresholds here are interpretable, length-adaptive, and (when an
|
||||
LLM is available) advisory rather than final.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from ..proxy.output_savings import BaselineModel, echo_ratio, stratum_key
|
||||
from ..proxy.output_shaper import classify_turn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Average adult reading speed (words/min) for technical prose. Used to turn a
|
||||
# response's length into "how long it would take to read", so a fast-skip is
|
||||
# defined relative to answer length rather than a fixed wall-clock cutoff.
|
||||
_READING_WPM = 250.0
|
||||
# A reply arriving in less than this fraction of the read-time means the answer
|
||||
# was almost certainly not read.
|
||||
_SKIP_READ_FRACTION = 0.5
|
||||
# Don't score skips on trivially short answers — there's nothing to skip.
|
||||
_MIN_WORDS_FOR_SKIP = 150
|
||||
# Floor for the per-user adaptive "long output" threshold (words).
|
||||
_LONG_OUTPUT_FLOOR = 200
|
||||
# Only sample the few preceding messages for echo context (bounded cost).
|
||||
_ECHO_CONTEXT_LOOKBACK = 4
|
||||
|
||||
_INTERRUPT_MARKER = "[Request interrupted by user"
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Response:
|
||||
"""One assistant response, with the request features that produced it."""
|
||||
|
||||
words: int
|
||||
output_tokens: int
|
||||
input_tokens: int
|
||||
model: str
|
||||
turn_kind: str
|
||||
has_tools: bool
|
||||
ts: float | None
|
||||
echo: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class _HumanMsg:
|
||||
ts: float | None
|
||||
is_interrupt: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class VerbositySignals:
|
||||
"""Behavioral signals aggregated across a project's sessions."""
|
||||
|
||||
sessions: int = 0
|
||||
human_msgs: int = 0
|
||||
interrupts: int = 0
|
||||
asst_responses: int = 0
|
||||
asst_words: int = 0
|
||||
long_outputs: int = 0
|
||||
fast_skips: int = 0
|
||||
skip_eligible: int = 0
|
||||
mean_echo_ratio: float = 0.0
|
||||
|
||||
@property
|
||||
def interrupt_rate(self) -> float:
|
||||
denom = self.human_msgs + self.interrupts
|
||||
return self.interrupts / denom if denom else 0.0
|
||||
|
||||
@property
|
||||
def fast_skip_rate(self) -> float:
|
||||
return self.fast_skips / self.skip_eligible if self.skip_eligible else 0.0
|
||||
|
||||
@property
|
||||
def long_output_rate(self) -> float:
|
||||
return self.long_outputs / self.asst_responses if self.asst_responses else 0.0
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"sessions": self.sessions,
|
||||
"human_msgs": self.human_msgs,
|
||||
"interrupts": self.interrupts,
|
||||
"interrupt_rate": round(self.interrupt_rate, 4),
|
||||
"asst_responses": self.asst_responses,
|
||||
"long_outputs": self.long_outputs,
|
||||
"long_output_rate": round(self.long_output_rate, 4),
|
||||
"fast_skips": self.fast_skips,
|
||||
"skip_eligible": self.skip_eligible,
|
||||
"fast_skip_rate": round(self.fast_skip_rate, 4),
|
||||
"mean_echo_ratio": round(self.mean_echo_ratio, 4),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class VerbosityProfile:
|
||||
"""The learned recommendation for a project."""
|
||||
|
||||
project_path: str
|
||||
level: int
|
||||
confidence: str # "low" | "medium" | "high"
|
||||
source: str # "heuristic" | "llm"
|
||||
rationale: str
|
||||
signals: dict[str, Any] = field(default_factory=dict)
|
||||
learned_at: str | None = None # caller stamps (Date.now unavailable here)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"project_path": self.project_path,
|
||||
"verbosity_level": self.level,
|
||||
"confidence": self.confidence,
|
||||
"source": self.source,
|
||||
"rationale": self.rationale,
|
||||
"signals": self.signals,
|
||||
"learned_at": self.learned_at,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: Path) -> VerbosityProfile | None:
|
||||
try:
|
||||
d = json.loads(Path(path).read_text())
|
||||
except (OSError, json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
return cls(
|
||||
project_path=d.get("project_path", ""),
|
||||
level=int(d.get("verbosity_level", 2)),
|
||||
confidence=d.get("confidence", "low"),
|
||||
source=d.get("source", "heuristic"),
|
||||
rationale=d.get("rationale", ""),
|
||||
signals=d.get("signals", {}),
|
||||
learned_at=d.get("learned_at"),
|
||||
)
|
||||
|
||||
def save(self, path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(self.to_dict(), indent=2))
|
||||
|
||||
|
||||
def _parse_ts(s: str | None) -> float | None:
|
||||
if not s:
|
||||
return None
|
||||
try:
|
||||
from datetime import datetime
|
||||
|
||||
return datetime.fromisoformat(s.replace("Z", "+00:00")).timestamp()
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _assistant_words_and_text(content: Any) -> tuple[int, str]:
|
||||
if not isinstance(content, list):
|
||||
return 0, ""
|
||||
text = " ".join(
|
||||
b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"
|
||||
)
|
||||
return len(text.split()), text
|
||||
|
||||
|
||||
def _human_text(content: Any) -> str | None:
|
||||
"""Return human-typed text, or None for tool results / slash commands / meta."""
|
||||
if isinstance(content, str):
|
||||
t = content
|
||||
elif isinstance(content, list):
|
||||
if any(isinstance(b, dict) and b.get("type") == "tool_result" for b in content):
|
||||
return None
|
||||
t = " ".join(
|
||||
b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"
|
||||
)
|
||||
else:
|
||||
return None
|
||||
if "<command-name>" in t or "<local-command-stdout>" in t:
|
||||
return None
|
||||
return t
|
||||
|
||||
|
||||
def _parse_session(path: Path) -> tuple[list[_Response], list[_HumanMsg], bool]:
|
||||
"""Parse one transcript into responses + human messages + has_tools flag."""
|
||||
responses: list[_Response] = []
|
||||
humans: list[_HumanMsg] = []
|
||||
has_tools = False
|
||||
prior_messages: list[dict[str, Any]] = []
|
||||
recent_context: list[str] = []
|
||||
|
||||
try:
|
||||
lines = path.read_text().splitlines()
|
||||
except (OSError, UnicodeDecodeError):
|
||||
return [], [], False
|
||||
|
||||
for line in lines:
|
||||
try:
|
||||
d = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
ltype = d.get("type")
|
||||
msg = d.get("message", {}) if isinstance(d.get("message"), dict) else {}
|
||||
ts = _parse_ts(d.get("timestamp"))
|
||||
|
||||
if ltype == "assistant":
|
||||
content = msg.get("content", [])
|
||||
if isinstance(content, list) and any(
|
||||
isinstance(b, dict) and b.get("type") == "tool_use" for b in content
|
||||
):
|
||||
has_tools = True
|
||||
words, text = _assistant_words_and_text(content)
|
||||
usage = msg.get("usage", {}) if isinstance(msg.get("usage"), dict) else {}
|
||||
in_tok = (
|
||||
usage.get("input_tokens", 0)
|
||||
+ usage.get("cache_read_input_tokens", 0)
|
||||
+ usage.get("cache_creation_input_tokens", 0)
|
||||
)
|
||||
out_tok = usage.get("output_tokens", 0)
|
||||
if words > 0 or out_tok > 0:
|
||||
ctx = " ".join(recent_context[-_ECHO_CONTEXT_LOOKBACK:])
|
||||
responses.append(
|
||||
_Response(
|
||||
words=words,
|
||||
output_tokens=out_tok,
|
||||
input_tokens=in_tok,
|
||||
model=str(msg.get("model", "")),
|
||||
turn_kind=classify_turn(prior_messages).value,
|
||||
has_tools=False, # filled after the session scan
|
||||
ts=ts,
|
||||
echo=echo_ratio(text, ctx) if text and ctx else 0.0,
|
||||
)
|
||||
)
|
||||
prior_messages.append({"role": "assistant", "content": content})
|
||||
|
||||
elif ltype == "user":
|
||||
content = msg.get("content")
|
||||
prior_messages.append({"role": "user", "content": content})
|
||||
# Feed tool results / user text into echo context.
|
||||
if isinstance(content, list):
|
||||
for b in content:
|
||||
if isinstance(b, dict) and b.get("type") == "tool_result":
|
||||
rc = b.get("content", "")
|
||||
recent_context.append(rc if isinstance(rc, str) else str(rc))
|
||||
has_tools = True
|
||||
human = _human_text(content)
|
||||
if human is None:
|
||||
continue
|
||||
if _INTERRUPT_MARKER in human:
|
||||
humans.append(_HumanMsg(ts=ts, is_interrupt=True))
|
||||
else:
|
||||
humans.append(_HumanMsg(ts=ts, is_interrupt=False))
|
||||
recent_context.append(human)
|
||||
|
||||
# Backfill has_tools (a session-level property of the harness).
|
||||
for r in responses:
|
||||
r.has_tools = has_tools
|
||||
return responses, humans, has_tools
|
||||
|
||||
|
||||
def _ordered_events(path: Path) -> list[tuple[float | None, str, _Response | _HumanMsg]]:
|
||||
"""Re-read interleaving order so fast-skip can pair a human reply to the
|
||||
assistant response immediately before it. Kept simple: re-parse with a tag.
|
||||
"""
|
||||
out: list[tuple[float | None, str, Any]] = []
|
||||
try:
|
||||
lines = path.read_text().splitlines()
|
||||
except (OSError, UnicodeDecodeError):
|
||||
return out
|
||||
responses, humans, _ = _parse_session(path)
|
||||
# The two lists are already in file order; interleave by re-walking lines.
|
||||
ri = hi = 0
|
||||
for line in lines:
|
||||
try:
|
||||
d = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
ltype = d.get("type")
|
||||
if ltype == "assistant" and ri < len(responses):
|
||||
out.append((responses[ri].ts, "assistant", responses[ri]))
|
||||
ri += 1
|
||||
elif ltype == "user":
|
||||
msg = d.get("message", {}) if isinstance(d.get("message"), dict) else {}
|
||||
text = _human_text(msg.get("content"))
|
||||
if text is None:
|
||||
continue
|
||||
if hi < len(humans):
|
||||
out.append((humans[hi].ts, "human", humans[hi]))
|
||||
hi += 1
|
||||
return out
|
||||
|
||||
|
||||
def extract_signals(
|
||||
session_paths: list[Path],
|
||||
) -> tuple[VerbositySignals, BaselineModel]:
|
||||
"""Compute behavioral signals and the per-stratum output-token baseline."""
|
||||
sig = VerbositySignals()
|
||||
baseline = BaselineModel()
|
||||
all_response_words: list[int] = []
|
||||
|
||||
# First pass: collect every response word-count to derive the per-user
|
||||
# adaptive "long" threshold (median), so "long" scales to the user.
|
||||
parsed: list[tuple[list[_Response], list[_HumanMsg]]] = []
|
||||
for p in session_paths:
|
||||
responses, humans, _ = _parse_session(p)
|
||||
if not responses and not humans:
|
||||
continue
|
||||
sig.sessions += 1
|
||||
parsed.append((responses, humans))
|
||||
all_response_words.extend(r.words for r in responses if r.words > 0)
|
||||
|
||||
long_threshold = _LONG_OUTPUT_FLOOR
|
||||
if all_response_words:
|
||||
all_response_words.sort()
|
||||
median = all_response_words[len(all_response_words) // 2]
|
||||
long_threshold = max(_LONG_OUTPUT_FLOOR, median)
|
||||
|
||||
echo_sum = 0.0
|
||||
echo_n = 0
|
||||
for responses, humans in parsed:
|
||||
for r in responses:
|
||||
sig.asst_responses += 1
|
||||
sig.asst_words += r.words
|
||||
if r.words >= long_threshold:
|
||||
sig.long_outputs += 1
|
||||
if r.echo > 0:
|
||||
echo_sum += r.echo
|
||||
echo_n += 1
|
||||
if r.output_tokens > 0:
|
||||
key = stratum_key(
|
||||
turn_kind=r.turn_kind,
|
||||
input_tokens=r.input_tokens,
|
||||
model=r.model or "unknown",
|
||||
has_tools=r.has_tools,
|
||||
)
|
||||
baseline.observe(key, r.output_tokens)
|
||||
for h in humans:
|
||||
if h.is_interrupt:
|
||||
sig.interrupts += 1
|
||||
else:
|
||||
sig.human_msgs += 1
|
||||
|
||||
# Second pass: fast-skip pairing via interleaved order.
|
||||
for p in session_paths:
|
||||
events = _ordered_events(p)
|
||||
last_resp: _Response | None = None
|
||||
for ts, kind, obj in events:
|
||||
if kind == "assistant":
|
||||
last_resp = obj # type: ignore[assignment]
|
||||
elif kind == "human":
|
||||
hm: _HumanMsg = obj # type: ignore[assignment]
|
||||
if (
|
||||
last_resp is not None
|
||||
and not hm.is_interrupt
|
||||
and last_resp.words >= _MIN_WORDS_FOR_SKIP
|
||||
and ts is not None
|
||||
and last_resp.ts is not None
|
||||
):
|
||||
sig.skip_eligible += 1
|
||||
read_secs = last_resp.words / _READING_WPM * 60.0
|
||||
if (ts - last_resp.ts) < _SKIP_READ_FRACTION * read_secs:
|
||||
sig.fast_skips += 1
|
||||
last_resp = None
|
||||
|
||||
sig.mean_echo_ratio = echo_sum / echo_n if echo_n else 0.0
|
||||
return sig, baseline
|
||||
|
||||
|
||||
def recommend_level(sig: VerbositySignals) -> tuple[int, str, str]:
|
||||
"""Heuristic prior mapping signals → (level, confidence, rationale).
|
||||
|
||||
This is the prior; an LLM judgment pass (in :func:`analyze`) may override
|
||||
it. Bands are interpretable: the more a user interrupts and fast-skips, the
|
||||
less of the output they consume, so the terser we should make it.
|
||||
"""
|
||||
if sig.human_msgs + sig.interrupts < 10:
|
||||
return 2, "low", "Too few human turns to calibrate; defaulting to L2."
|
||||
|
||||
ir = sig.interrupt_rate
|
||||
fsr = sig.fast_skip_rate
|
||||
pressure = ir + fsr # combined "too much output" pressure
|
||||
|
||||
confidence = "high" if (sig.human_msgs + sig.interrupts) >= 60 else "medium"
|
||||
|
||||
if pressure < 0.10:
|
||||
return (
|
||||
1,
|
||||
confidence,
|
||||
(
|
||||
f"Low push-back (interrupt {ir:.0%}, fast-skip {fsr:.0%}); user reads "
|
||||
"answers — light touch (L1)."
|
||||
),
|
||||
)
|
||||
if pressure < 0.30:
|
||||
return (
|
||||
2,
|
||||
confidence,
|
||||
(
|
||||
f"Moderate push-back (interrupt {ir:.0%}, fast-skip {fsr:.0%}); "
|
||||
"drop ceremony and echo (L2)."
|
||||
),
|
||||
)
|
||||
if pressure < 0.55:
|
||||
return (
|
||||
3,
|
||||
confidence,
|
||||
(
|
||||
f"High push-back (interrupt {ir:.0%}, fast-skip {fsr:.0%}); user "
|
||||
"rarely reads long answers — conclusions only (L3)."
|
||||
),
|
||||
)
|
||||
return (
|
||||
3,
|
||||
confidence,
|
||||
(
|
||||
f"Very high push-back (interrupt {ir:.0%}, fast-skip {fsr:.0%}); capping "
|
||||
"at L3 rather than auto-applying caveman L4."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def analyze(
|
||||
session_paths: list[Path],
|
||||
project_path: str,
|
||||
*,
|
||||
llm_judge: Any | None = None,
|
||||
) -> tuple[VerbosityProfile, BaselineModel]:
|
||||
"""Full analysis: signals → recommendation → profile, plus the baseline.
|
||||
|
||||
``llm_judge``, if given, is a callable ``(signals_dict) -> (level, rationale)``
|
||||
that overrides the heuristic. Kept injectable so the core stays LLM-free and
|
||||
testable; the CLI wires a real LLM call.
|
||||
"""
|
||||
sig, baseline = extract_signals(session_paths)
|
||||
level, confidence, rationale = recommend_level(sig)
|
||||
source = "heuristic"
|
||||
if llm_judge is not None:
|
||||
try:
|
||||
verdict = llm_judge(sig.to_dict())
|
||||
if verdict is not None:
|
||||
level, rationale = verdict
|
||||
level = max(0, min(4, int(level)))
|
||||
source = "llm"
|
||||
except Exception as e: # LLM is advisory — never fail the analysis
|
||||
logger.warning("verbosity LLM judge failed, using heuristic: %s", e)
|
||||
|
||||
profile = VerbosityProfile(
|
||||
project_path=project_path,
|
||||
level=level,
|
||||
confidence=confidence,
|
||||
source=source,
|
||||
rationale=rationale,
|
||||
signals=sig.to_dict(),
|
||||
)
|
||||
return profile, baseline
|
||||
|
|
@ -1697,6 +1697,64 @@ class AnthropicHandlerMixin:
|
|||
optimized_tokens = tokenizer.count_messages(body["messages"])
|
||||
tokens_saved = max(0, original_tokens - optimized_tokens)
|
||||
|
||||
# Output shaping (opt-in via HEADROOM_OUTPUT_SHAPER): verbosity
|
||||
# steering appended to the system-prompt tail + effort routing on
|
||||
# mechanical tool_result continuations. Runs after every other
|
||||
# body mutation so the turn classifier sees the final messages,
|
||||
# and respects the same bypass header as compression.
|
||||
if not _bypass:
|
||||
from headroom.proxy.output_savings import (
|
||||
assign_arm,
|
||||
conversation_key_from_body,
|
||||
stratum_key,
|
||||
stratum_label,
|
||||
)
|
||||
from headroom.proxy.output_shaper import (
|
||||
OutputShaperSettings,
|
||||
classify_turn,
|
||||
resolve_verbosity_level,
|
||||
shape_request,
|
||||
)
|
||||
|
||||
_shaper_settings = OutputShaperSettings.from_env()
|
||||
if _shaper_settings.enabled:
|
||||
# Conversation-stable holdout assignment: a whole
|
||||
# conversation is treatment or control. This keeps the A/B
|
||||
# comparison clean AND keeps the prefix cache stable (we
|
||||
# never flip a conversation's system-prompt tail mid-stream).
|
||||
import os as _os
|
||||
|
||||
_holdout = 0.0
|
||||
try:
|
||||
_holdout = float(_os.environ.get("HEADROOM_OUTPUT_HOLDOUT", "0") or "0")
|
||||
except ValueError:
|
||||
_holdout = 0.0
|
||||
_arm = assign_arm(conversation_key_from_body(body), _holdout)
|
||||
|
||||
# Stratum from request features observable now (mirrors the
|
||||
# offline baseline so live and learned strata line up).
|
||||
_turn_kind = classify_turn(body.get("messages", [])).value
|
||||
_stratum = stratum_key(
|
||||
turn_kind=_turn_kind,
|
||||
input_tokens=original_tokens,
|
||||
model=model,
|
||||
has_tools=bool(body.get("tools")),
|
||||
)
|
||||
# Carry (arm, stratum) on the existing label channel so the
|
||||
# outcome funnel can feed the savings ledger from any path.
|
||||
transforms_applied.append(stratum_label(_arm, _stratum))
|
||||
|
||||
if _arm == "treatment":
|
||||
_level, _src = resolve_verbosity_level(_shaper_settings)
|
||||
shape_result = shape_request(body, _shaper_settings, level_override=_level)
|
||||
if shape_result.changed:
|
||||
body_mutation_tracker.mark_mutated("output_shaper")
|
||||
transforms_applied.extend(shape_result.labels or [])
|
||||
logger.info(
|
||||
f"[{request_id}] OutputShaper(L{_level}/{_src}): "
|
||||
f"{shape_result.labels}"
|
||||
)
|
||||
|
||||
# Unit 2: mark end of pre-upstream phase. Everything after this
|
||||
# point is upstream I/O or post-response bookkeeping.
|
||||
stage_timer.record(
|
||||
|
|
|
|||
|
|
@ -333,6 +333,18 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
|
|||
from headroom.proxy.models import RequestLog
|
||||
from headroom.proxy.project_context import get_current_project
|
||||
|
||||
# Output-shaping savings ledger (counterfactual estimator). The shaper
|
||||
# tags each request's (arm, stratum) onto ``transforms_applied``; feed the
|
||||
# observed output tokens to the recorder so it can produce an honest
|
||||
# reduction estimate. Best-effort: never let bookkeeping break a response.
|
||||
if any(str(t).startswith("output_shaper:") for t in outcome.transforms_applied):
|
||||
try:
|
||||
from headroom.proxy.output_savings import get_recorder
|
||||
|
||||
get_recorder().record_from_labels(outcome.transforms_applied, outcome.output_tokens)
|
||||
except Exception: # pragma: no cover - defensive
|
||||
pass
|
||||
|
||||
# Project attribution: explicit outcome field wins, else the value the
|
||||
# HTTP middleware / WS accept captured from ``X-Headroom-Project``.
|
||||
project = outcome.project or get_current_project()
|
||||
|
|
|
|||
501
headroom/proxy/output_savings.py
Normal file
501
headroom/proxy/output_savings.py
Normal file
|
|
@ -0,0 +1,501 @@
|
|||
"""Counterfactual estimation of output-token reduction.
|
||||
|
||||
The hard problem: output-token savings are **counterfactual**. When the shaper
|
||||
makes a request terser, the model emits N output tokens — but we never observe
|
||||
what it *would* have emitted unshaped. Input compression is a pure function, so
|
||||
``tokens_before``/``tokens_after`` are both observable. Output is not: only one
|
||||
side of the counterfactual happens per request. So a flat "we save 30%" claim
|
||||
is marketing, not measurement.
|
||||
|
||||
This module makes the estimate honest by separating three tiers:
|
||||
|
||||
1. **Estimated (synthetic control).** A per-stratum baseline of unshaped output
|
||||
tokens — built by ``learn --verbosity`` from session history that predates
|
||||
the shaper — gives an expected output for each request's feature stratum.
|
||||
``estimate = Σ (baseline_mean[stratum] − observed_output)`` over shaped
|
||||
requests, summed as **signed** deltas (never clamped per-request — clamping
|
||||
biases upward). Reported with a propagated confidence interval and always
|
||||
labelled an estimate, never "measured".
|
||||
|
||||
2. **Measured (A/B holdout).** When a small holdout fraction of conversations
|
||||
is left unshaped, the difference of per-stratum means between the treatment
|
||||
and control arms is an unbiased causal estimate. This is the only number we
|
||||
call "measured". Assignment is **conversation-stable** (a whole conversation
|
||||
is in one arm) for two reasons that happen to align: mixing shaped and
|
||||
unshaped turns within one conversation would (a) pollute the comparison and
|
||||
(b) bust the prefix cache by changing the system-prompt tail mid-stream.
|
||||
|
||||
3. **Direct waste (no counterfactual).** Echo ratio — n-gram overlap between a
|
||||
response and the context it was given — is a property of a single response,
|
||||
measurable with no counterfactual. "32% of output restated existing context"
|
||||
is an honest standalone fact and the shaper's target. See ``echo_ratio``.
|
||||
|
||||
Stratification uses only features observable at request time (never the output):
|
||||
turn kind, input-token bucket, model family, whether tools are present.
|
||||
|
||||
Pure module: no I/O except explicit ``load``/``save``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from typing import Any
|
||||
|
||||
# Coarse input-token buckets. Coarse on purpose: too many strata make
|
||||
# per-stratum baselines sparse and noisy. Boundaries in tokens.
|
||||
_INPUT_BUCKETS = (2_000, 8_000, 32_000, 128_000)
|
||||
|
||||
|
||||
def input_bucket(input_tokens: int) -> str:
|
||||
"""Map an input-token count to a coarse bucket label."""
|
||||
if input_tokens < _INPUT_BUCKETS[0]:
|
||||
return "xs"
|
||||
if input_tokens < _INPUT_BUCKETS[1]:
|
||||
return "s"
|
||||
if input_tokens < _INPUT_BUCKETS[2]:
|
||||
return "m"
|
||||
if input_tokens < _INPUT_BUCKETS[3]:
|
||||
return "l"
|
||||
return "xl"
|
||||
|
||||
|
||||
def model_family(model: str) -> str:
|
||||
"""Collapse a model id to a coarse family for stratification.
|
||||
|
||||
Token-spend behaviour clusters by family far more than by point release,
|
||||
so we bucket (e.g.) every ``claude-opus-*`` together.
|
||||
"""
|
||||
m = model.lower()
|
||||
for fam in ("opus", "sonnet", "haiku", "fable", "mythos", "gpt", "gemini"):
|
||||
if fam in m:
|
||||
return fam
|
||||
return "other"
|
||||
|
||||
|
||||
def stratum_key(
|
||||
*,
|
||||
turn_kind: str,
|
||||
input_tokens: int,
|
||||
model: str,
|
||||
has_tools: bool,
|
||||
) -> str:
|
||||
"""Build a stratum key from request features observable BEFORE the response.
|
||||
|
||||
Order is most→least specific so :meth:`BaselineModel.lookup` can back off
|
||||
by trimming trailing fields.
|
||||
"""
|
||||
return "|".join(
|
||||
(
|
||||
model_family(model),
|
||||
turn_kind,
|
||||
input_bucket(input_tokens),
|
||||
"tools" if has_tools else "notools",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def conversation_key_from_body(body: dict[str, Any]) -> str:
|
||||
"""Derive a conversation-stable key for holdout assignment.
|
||||
|
||||
Stable across every turn of one conversation (so the whole conversation
|
||||
lands in one arm) and cheap: a hash of the model plus the first user
|
||||
message's text. The first user turn is immutable for a conversation's
|
||||
lifetime, which is exactly the stability we need.
|
||||
"""
|
||||
model = str(body.get("model", ""))
|
||||
seed = model
|
||||
for msg in body.get("messages", []):
|
||||
if isinstance(msg, dict) and msg.get("role") == "user":
|
||||
content = msg.get("content")
|
||||
if isinstance(content, str):
|
||||
seed += "\x00" + content[:512]
|
||||
elif isinstance(content, list):
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
seed += "\x00" + str(block.get("text", ""))[:512]
|
||||
break
|
||||
break
|
||||
return hashlib.sha256(seed.encode("utf-8", "ignore")).hexdigest()
|
||||
|
||||
|
||||
def assign_arm(conversation_key: str, holdout_fraction: float) -> str:
|
||||
"""Deterministically assign a conversation to ``treatment`` or ``control``.
|
||||
|
||||
``holdout_fraction`` in [0, 1] is the share routed to ``control`` (left
|
||||
unshaped for measurement). Hashing the conversation key keeps assignment
|
||||
stable across the conversation's turns and uniform across conversations.
|
||||
"""
|
||||
if holdout_fraction <= 0.0:
|
||||
return "treatment"
|
||||
if holdout_fraction >= 1.0:
|
||||
return "control"
|
||||
digest = hashlib.sha256(("arm:" + conversation_key).encode()).hexdigest()
|
||||
# Map the first 8 hex digits to [0, 1).
|
||||
frac = int(digest[:8], 16) / 0xFFFFFFFF
|
||||
return "control" if frac < holdout_fraction else "treatment"
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Accum:
|
||||
"""Running count / sum / sum-of-squares for online mean & variance."""
|
||||
|
||||
n: int = 0
|
||||
sum: float = 0.0
|
||||
sumsq: float = 0.0
|
||||
|
||||
def add(self, x: float) -> None:
|
||||
self.n += 1
|
||||
self.sum += x
|
||||
self.sumsq += x * x
|
||||
|
||||
@property
|
||||
def mean(self) -> float:
|
||||
return self.sum / self.n if self.n else 0.0
|
||||
|
||||
@property
|
||||
def var(self) -> float:
|
||||
"""Sample variance (unbiased). 0 when fewer than 2 observations."""
|
||||
if self.n < 2:
|
||||
return 0.0
|
||||
return max(0.0, (self.sumsq - self.sum * self.sum / self.n) / (self.n - 1))
|
||||
|
||||
def to_dict(self) -> dict[str, float]:
|
||||
return {"n": self.n, "sum": self.sum, "sumsq": self.sumsq}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict[str, float]) -> _Accum:
|
||||
a = cls()
|
||||
a.n = int(d.get("n", 0))
|
||||
a.sum = float(d.get("sum", 0.0))
|
||||
a.sumsq = float(d.get("sumsq", 0.0))
|
||||
return a
|
||||
|
||||
|
||||
@dataclass
|
||||
class BaselineModel:
|
||||
"""Per-stratum baseline of unshaped output tokens (the synthetic control).
|
||||
|
||||
Built offline by ``learn --verbosity`` from pre-shaper history. ``strata``
|
||||
maps a stratum key to its accumulator; ``glob`` is the all-requests
|
||||
fallback for strata never seen during training.
|
||||
"""
|
||||
|
||||
strata: dict[str, _Accum] = field(default_factory=dict)
|
||||
glob: _Accum = field(default_factory=_Accum)
|
||||
|
||||
def observe(self, key: str, output_tokens: int) -> None:
|
||||
self.strata.setdefault(key, _Accum()).add(output_tokens)
|
||||
self.glob.add(output_tokens)
|
||||
|
||||
def lookup(self, key: str) -> tuple[float, float, int]:
|
||||
"""Return ``(mean, var, n)`` for *key* with hierarchical back-off.
|
||||
|
||||
Falls back by trimming trailing (least-specific) stratum fields, then
|
||||
to the global mean. Back-off keeps the estimate defined for strata the
|
||||
baseline never saw, at the cost of specificity.
|
||||
"""
|
||||
acc = self.strata.get(key)
|
||||
if acc is not None and acc.n > 0:
|
||||
return acc.mean, acc.var, acc.n
|
||||
parts = key.split("|")
|
||||
while len(parts) > 1:
|
||||
parts = parts[:-1]
|
||||
prefix = "|".join(parts)
|
||||
for k, a in self.strata.items():
|
||||
if k.startswith(prefix + "|") and a.n > 0:
|
||||
return a.mean, a.var, a.n
|
||||
return self.glob.mean, self.glob.var, self.glob.n
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"strata": {k: a.to_dict() for k, a in self.strata.items()},
|
||||
"glob": self.glob.to_dict(),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict[str, Any]) -> BaselineModel:
|
||||
m = cls()
|
||||
for k, a in (d.get("strata") or {}).items():
|
||||
m.strata[k] = _Accum.from_dict(a)
|
||||
m.glob = _Accum.from_dict(d.get("glob") or {})
|
||||
return m
|
||||
|
||||
@property
|
||||
def total_samples(self) -> int:
|
||||
return self.glob.n
|
||||
|
||||
|
||||
@dataclass
|
||||
class SavingsEstimate:
|
||||
"""Result of an estimation pass."""
|
||||
|
||||
tokens_saved: float
|
||||
baseline_tokens: float
|
||||
pct: float
|
||||
ci_low_pct: float
|
||||
ci_high_pct: float
|
||||
n_requests: int
|
||||
kind: str # "estimated" (synthetic control) or "measured" (A/B holdout)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SavingsLedger:
|
||||
"""Accumulates shaped (treatment) and unshaped (control) observations and
|
||||
produces honest reduction estimates.
|
||||
|
||||
``baseline`` is the offline synthetic control. ``treatment``/``control``
|
||||
are live per-stratum accumulators of observed output tokens, used both for
|
||||
the A/B "measured" number (when a holdout exists) and to keep the ledger
|
||||
self-describing.
|
||||
"""
|
||||
|
||||
baseline: BaselineModel = field(default_factory=BaselineModel)
|
||||
treatment: dict[str, _Accum] = field(default_factory=dict)
|
||||
control: dict[str, _Accum] = field(default_factory=dict)
|
||||
|
||||
# ---- recording -------------------------------------------------------
|
||||
|
||||
def record(self, arm: str, key: str, output_tokens: int) -> None:
|
||||
target = self.treatment if arm == "treatment" else self.control
|
||||
target.setdefault(key, _Accum()).add(output_tokens)
|
||||
|
||||
# ---- estimation ------------------------------------------------------
|
||||
|
||||
def estimate_from_baseline(self) -> SavingsEstimate:
|
||||
"""Synthetic-control estimate: treatment output vs. offline baseline.
|
||||
|
||||
Aggregate signed delta ``Σ_s n_s·(μ_s − ȳ_s)`` where μ_s is the
|
||||
baseline mean and ȳ_s the observed treatment mean. Variance propagates
|
||||
both the observed-output spread and the finite-baseline-sample error:
|
||||
|
||||
Var ≈ Σ_s [ n_s·σ²_y,s + n_s²·σ²_μ,s / m_s ]
|
||||
"""
|
||||
total_saved = 0.0
|
||||
total_baseline = 0.0
|
||||
var = 0.0
|
||||
n_requests = 0
|
||||
for key, acc in self.treatment.items():
|
||||
if acc.n == 0:
|
||||
continue
|
||||
mu, mu_var, m = self.baseline.lookup(key)
|
||||
if m == 0:
|
||||
continue
|
||||
n = acc.n
|
||||
n_requests += n
|
||||
total_saved += n * (mu - acc.mean)
|
||||
total_baseline += n * mu
|
||||
var += n * acc.var
|
||||
if m > 0:
|
||||
var += (n * n) * (mu_var / m)
|
||||
return self._finalize(total_saved, total_baseline, var, n_requests, "estimated")
|
||||
|
||||
def estimate_from_holdout(self) -> SavingsEstimate | None:
|
||||
"""A/B measurement: per-stratum control mean minus treatment mean.
|
||||
|
||||
Only strata with data in BOTH arms contribute. Returns ``None`` if no
|
||||
such stratum exists (no holdout traffic yet). Weighted by treatment
|
||||
volume; this is the unbiased causal number.
|
||||
"""
|
||||
total_saved = 0.0
|
||||
total_baseline = 0.0
|
||||
var = 0.0
|
||||
n_requests = 0
|
||||
contributing = 0
|
||||
for key, t in self.treatment.items():
|
||||
c = self.control.get(key)
|
||||
if c is None or c.n == 0 or t.n == 0:
|
||||
continue
|
||||
contributing += 1
|
||||
n = t.n
|
||||
n_requests += n
|
||||
delta = c.mean - t.mean # tokens saved per request in this stratum
|
||||
total_saved += n * delta
|
||||
total_baseline += n * c.mean
|
||||
# Var of (c.mean - t.mean) = σ²_c/n_c + σ²_t/n_t, scaled by n².
|
||||
var += (n * n) * (c.var / c.n + t.var / t.n)
|
||||
if contributing == 0:
|
||||
return None
|
||||
return self._finalize(total_saved, total_baseline, var, n_requests, "measured")
|
||||
|
||||
@staticmethod
|
||||
def _finalize(
|
||||
total_saved: float,
|
||||
total_baseline: float,
|
||||
var: float,
|
||||
n_requests: int,
|
||||
kind: str,
|
||||
) -> SavingsEstimate:
|
||||
pct = (total_saved / total_baseline * 100.0) if total_baseline > 0 else 0.0
|
||||
se = math.sqrt(var)
|
||||
# 95% normal-approx band on the token total, converted to percent.
|
||||
lo = total_saved - 1.96 * se
|
||||
hi = total_saved + 1.96 * se
|
||||
ci_low = (lo / total_baseline * 100.0) if total_baseline > 0 else 0.0
|
||||
ci_high = (hi / total_baseline * 100.0) if total_baseline > 0 else 0.0
|
||||
return SavingsEstimate(
|
||||
tokens_saved=total_saved,
|
||||
baseline_tokens=total_baseline,
|
||||
pct=pct,
|
||||
ci_low_pct=ci_low,
|
||||
ci_high_pct=ci_high,
|
||||
n_requests=n_requests,
|
||||
kind=kind,
|
||||
)
|
||||
|
||||
def best_estimate(self) -> SavingsEstimate:
|
||||
"""Prefer the measured A/B number; fall back to the baseline estimate."""
|
||||
measured = self.estimate_from_holdout()
|
||||
return measured if measured is not None else self.estimate_from_baseline()
|
||||
|
||||
# ---- persistence -----------------------------------------------------
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"baseline": self.baseline.to_dict(),
|
||||
"treatment": {k: a.to_dict() for k, a in self.treatment.items()},
|
||||
"control": {k: a.to_dict() for k, a in self.control.items()},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict[str, Any]) -> SavingsLedger:
|
||||
ledger = cls(baseline=BaselineModel.from_dict(d.get("baseline") or {}))
|
||||
for k, a in (d.get("treatment") or {}).items():
|
||||
ledger.treatment[k] = _Accum.from_dict(a)
|
||||
for k, a in (d.get("control") or {}).items():
|
||||
ledger.control[k] = _Accum.from_dict(a)
|
||||
return ledger
|
||||
|
||||
def save(self, path: Any) -> None:
|
||||
from pathlib import Path
|
||||
|
||||
p = Path(path)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text(json.dumps(self.to_dict(), separators=(",", ":")))
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: Any) -> SavingsLedger:
|
||||
from pathlib import Path
|
||||
|
||||
p = Path(path)
|
||||
if not p.exists():
|
||||
return cls()
|
||||
try:
|
||||
return cls.from_dict(json.loads(p.read_text()))
|
||||
except (json.JSONDecodeError, ValueError, OSError):
|
||||
return cls()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Live recording — rides the existing ``transforms_applied`` label channel so
|
||||
# every response path (streaming, non-streaming, backend) feeds the ledger with
|
||||
# no changes to RequestOutcome or its construction sites.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
_STRATUM_LABEL = "output_shaper:stratum:"
|
||||
_CONTROL_LABEL = "output_shaper:control:"
|
||||
|
||||
|
||||
def stratum_label(arm: str, key: str) -> str:
|
||||
"""Encode (arm, stratum) as a transforms_applied label."""
|
||||
prefix = _STRATUM_LABEL if arm == "treatment" else _CONTROL_LABEL
|
||||
return prefix + key
|
||||
|
||||
|
||||
def parse_stratum_label(label: str) -> tuple[str, str] | None:
|
||||
"""Decode a label into ``(arm, stratum)``, or None if not one of ours."""
|
||||
if label.startswith(_STRATUM_LABEL):
|
||||
return "treatment", label[len(_STRATUM_LABEL) :]
|
||||
if label.startswith(_CONTROL_LABEL):
|
||||
return "control", label[len(_CONTROL_LABEL) :]
|
||||
return None
|
||||
|
||||
|
||||
class SavingsRecorder:
|
||||
"""In-memory ledger with periodic flush, safe for concurrent requests.
|
||||
|
||||
Loads the baseline (written by ``learn --verbosity``) from disk, accumulates
|
||||
live treatment/control observations in memory, and flushes every
|
||||
``flush_every`` records so a busy proxy doesn't do a read-modify-write of the
|
||||
JSON file on every request.
|
||||
"""
|
||||
|
||||
def __init__(self, path: Any, flush_every: int = 25) -> None:
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
self._path = Path(path)
|
||||
self._lock = threading.Lock()
|
||||
self._ledger = SavingsLedger.load(self._path)
|
||||
self._flush_every = flush_every
|
||||
self._since_flush = 0
|
||||
|
||||
def record_from_labels(self, labels: Any, output_tokens: int) -> bool:
|
||||
"""Record one outcome given its transforms_applied labels. Returns True
|
||||
if a shaping label was found and recorded."""
|
||||
for label in labels or ():
|
||||
parsed = parse_stratum_label(str(label))
|
||||
if parsed is None:
|
||||
continue
|
||||
arm, key = parsed
|
||||
with self._lock:
|
||||
self._ledger.record(arm, key, output_tokens)
|
||||
self._since_flush += 1
|
||||
if self._since_flush >= self._flush_every:
|
||||
self._flush_locked()
|
||||
return True
|
||||
return False
|
||||
|
||||
def _flush_locked(self) -> None:
|
||||
try:
|
||||
self._ledger.save(self._path)
|
||||
self._since_flush = 0
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def flush(self) -> None:
|
||||
with self._lock:
|
||||
self._flush_locked()
|
||||
|
||||
def estimate(self) -> SavingsEstimate:
|
||||
with self._lock:
|
||||
return self._ledger.best_estimate()
|
||||
|
||||
|
||||
_RECORDER: SavingsRecorder | None = None
|
||||
|
||||
|
||||
def get_recorder() -> SavingsRecorder:
|
||||
"""Process-wide recorder singleton, rooted at the workspace dir."""
|
||||
global _RECORDER
|
||||
if _RECORDER is None:
|
||||
from ..paths import workspace_dir
|
||||
|
||||
_RECORDER = SavingsRecorder(workspace_dir() / "output_savings.json")
|
||||
return _RECORDER
|
||||
|
||||
|
||||
def echo_ratio(output_text: str, context_text: str, n: int = 8) -> float:
|
||||
"""Fraction of the response's n-grams that already appear in the context.
|
||||
|
||||
A measured (non-counterfactual) waste signal: high overlap means the model
|
||||
re-emitted code/text it was already shown. Token-ish word n-grams; cheap
|
||||
and language-agnostic. Returns 0.0 when the output is shorter than *n*.
|
||||
"""
|
||||
out_words = output_text.split()
|
||||
if len(out_words) < n:
|
||||
return 0.0
|
||||
ctx_words = context_text.split()
|
||||
ctx_grams = {" ".join(ctx_words[i : i + n]) for i in range(max(0, len(ctx_words) - n + 1))}
|
||||
if not ctx_grams:
|
||||
return 0.0
|
||||
out_grams = [" ".join(out_words[i : i + n]) for i in range(len(out_words) - n + 1)]
|
||||
if not out_grams:
|
||||
return 0.0
|
||||
hits = sum(1 for g in out_grams if g in ctx_grams)
|
||||
return hits / len(out_grams)
|
||||
361
headroom/proxy/output_shaper.py
Normal file
361
headroom/proxy/output_shaper.py
Normal file
|
|
@ -0,0 +1,361 @@
|
|||
"""Output token shaping for proxied Anthropic requests.
|
||||
|
||||
Headroom's transforms compress what goes INTO the model. This module is the
|
||||
first request-side lever on what comes OUT of it. The proxy never generates
|
||||
output tokens, so every lever here works by reshaping the request:
|
||||
|
||||
1. Verbosity steering — a deterministic instruction block appended to the
|
||||
TAIL of the system prompt (after any ``cache_control`` breakpoint, so the
|
||||
provider prefix cache is preserved). Five levels, from "no ceremony" to
|
||||
full caveman.
|
||||
|
||||
2. Effort routing — agentic loops are mostly mechanical continuations (the
|
||||
last message is a clean tool_result: a file read, a passing test). Thinking
|
||||
bills as output tokens, and harnesses like Claude Code pin
|
||||
``output_config.effort`` at ``xhigh`` for every turn. On turns classified
|
||||
as mechanical we lower an explicitly-present effort; on errors or new user
|
||||
asks we leave it alone. For legacy models still sending
|
||||
``thinking.budget_tokens`` we clamp the budget to the API floor instead.
|
||||
|
||||
Safety rules (each prevents a concrete failure mode):
|
||||
- Never INJECT ``output_config.effort`` where the client didn't send it —
|
||||
models without effort support 400 on it. Lowering an existing value is
|
||||
always valid.
|
||||
- Never toggle ``thinking.type`` — disabling thinking while history carries
|
||||
thinking blocks 400s on some models, and the toggle busts the messages
|
||||
cache tier.
|
||||
- Steering text is byte-stable per level and applied idempotently, so
|
||||
repeated requests keep an identical prefix.
|
||||
|
||||
Turn classification is purely structural (block types, roles, ``is_error``
|
||||
flags) — no content regexes or keyword patterns.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Documented Anthropic API minimum for thinking.budget_tokens on models
|
||||
# that still accept the legacy enabled/budget_tokens form.
|
||||
LEGACY_THINKING_FLOOR = 1024
|
||||
|
||||
# Ordering for output_config.effort values. Unknown values are left alone.
|
||||
_EFFORT_RANK = {"low": 0, "medium": 1, "high": 2, "xhigh": 3, "max": 4}
|
||||
|
||||
# Sentinel prefix marks the steering block so application is idempotent and
|
||||
# the block is recognizable in logs/diffs.
|
||||
_STEERING_SENTINEL = "<headroom_output_shaping>"
|
||||
_STEERING_SUFFIX = "</headroom_output_shaping>"
|
||||
|
||||
# Levels are cumulative: each includes everything above it. Text must stay
|
||||
# byte-stable across releases for prefix-cache friendliness — treat edits to
|
||||
# these strings as cache-busting changes.
|
||||
_VERBOSITY_LEVELS = {
|
||||
1: (
|
||||
"Skip preamble and postamble. Do not announce what you are about to "
|
||||
"do or recap what you just did; start with the substance."
|
||||
),
|
||||
2: (
|
||||
"Skip preamble and postamble; start with the substance. Never restate "
|
||||
"code, file contents, diffs, or tool output that already appear in "
|
||||
"this conversation — reference them by path and line instead. After a "
|
||||
"tool call succeeds, continue without narrating the result."
|
||||
),
|
||||
3: (
|
||||
"Skip preamble and postamble. Never restate code, file contents, "
|
||||
"diffs, or tool output already in this conversation — reference by "
|
||||
"path and line. Give conclusions only; omit rationale unless the user "
|
||||
"asks why. Prefer the smallest edit over rewriting whole files. Keep "
|
||||
"prose to the minimum needed to be unambiguous."
|
||||
),
|
||||
4: (
|
||||
"Minimum tokens. Fragments fine. No preamble, no postamble, no "
|
||||
"restating context, no rationale. Answer, smallest-possible edits, "
|
||||
"nothing else."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class TurnKind(Enum):
|
||||
"""Structural classification of the latest conversation turn."""
|
||||
|
||||
NEW_USER_ASK = "new_user_ask"
|
||||
MECHANICAL_CONTINUATION = "mechanical_continuation"
|
||||
ERROR_CONTINUATION = "error_continuation"
|
||||
UNKNOWN = "unknown"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OutputShaperSettings:
|
||||
"""Runtime settings, resolved once per request from the environment.
|
||||
|
||||
Env-driven (like HEADROOM_INTERCEPT_ENABLED) so the proxy picks it up
|
||||
without config plumbing through the server. Off by default.
|
||||
"""
|
||||
|
||||
enabled: bool = False
|
||||
verbosity_level: int = 2
|
||||
effort_router_enabled: bool = True
|
||||
mechanical_effort: str = "low"
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> OutputShaperSettings:
|
||||
enabled = os.environ.get("HEADROOM_OUTPUT_SHAPER", "").lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
)
|
||||
try:
|
||||
level = int(os.environ.get("HEADROOM_VERBOSITY_LEVEL", "2"))
|
||||
except ValueError:
|
||||
level = 2
|
||||
level = max(0, min(4, level))
|
||||
router = os.environ.get("HEADROOM_EFFORT_ROUTER", "1").lower() not in (
|
||||
"0",
|
||||
"false",
|
||||
"no",
|
||||
)
|
||||
mech = os.environ.get("HEADROOM_MECHANICAL_EFFORT", "low")
|
||||
if mech not in _EFFORT_RANK:
|
||||
mech = "low"
|
||||
return cls(
|
||||
enabled=enabled,
|
||||
verbosity_level=level,
|
||||
effort_router_enabled=router,
|
||||
mechanical_effort=mech,
|
||||
)
|
||||
|
||||
|
||||
def resolve_verbosity_level(settings: OutputShaperSettings) -> tuple[int, str]:
|
||||
"""Resolve the live verbosity level and its source.
|
||||
|
||||
Precedence:
|
||||
1. ``HEADROOM_VERBOSITY_LEVEL`` set explicitly → manual override.
|
||||
2. AIMD controller state (when ``HEADROOM_VERBOSITY_AUTOTUNE`` is on).
|
||||
3. Learned ``verbosity.json`` from ``learn --verbosity``.
|
||||
4. The settings default.
|
||||
|
||||
Returns ``(level, source)``. Kept separate from :func:`shape_request` so the
|
||||
body-mutating core stays a pure function of an explicit level.
|
||||
"""
|
||||
import os
|
||||
|
||||
if os.environ.get("HEADROOM_VERBOSITY_LEVEL"):
|
||||
return settings.verbosity_level, "env"
|
||||
|
||||
try:
|
||||
from ..paths import workspace_dir
|
||||
|
||||
ws = workspace_dir()
|
||||
except Exception:
|
||||
return settings.verbosity_level, "default"
|
||||
|
||||
autotune = os.environ.get("HEADROOM_VERBOSITY_AUTOTUNE", "").lower() in ("1", "true", "yes")
|
||||
if autotune:
|
||||
ctrl_path = ws / "verbosity_controller.json"
|
||||
if ctrl_path.exists():
|
||||
try:
|
||||
import json as _json
|
||||
|
||||
level = int(
|
||||
_json.loads(ctrl_path.read_text()).get("level", settings.verbosity_level)
|
||||
)
|
||||
return max(0, min(4, level)), "controller"
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
|
||||
prof_path = ws / "verbosity.json"
|
||||
if prof_path.exists():
|
||||
try:
|
||||
import json as _json
|
||||
|
||||
level = int(_json.loads(prof_path.read_text()).get("verbosity_level", -1))
|
||||
if 0 <= level <= 4:
|
||||
return level, "learned"
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
|
||||
return settings.verbosity_level, "default"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ShapeResult:
|
||||
"""What the shaper did to a request body."""
|
||||
|
||||
changed: bool = False
|
||||
labels: list[str] | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.labels is None:
|
||||
self.labels = []
|
||||
|
||||
|
||||
def classify_turn(messages: list[dict[str, Any]]) -> TurnKind:
|
||||
"""Classify the latest turn from message structure alone.
|
||||
|
||||
- Any text block in the last user message → the user is asking something
|
||||
new: full effort.
|
||||
- Only tool_result blocks, none flagged ``is_error`` → mechanical
|
||||
continuation: the model is resuming after a routine tool call.
|
||||
- Any tool_result with ``is_error: true`` → error continuation: the model
|
||||
must reason about a failure, keep full effort.
|
||||
"""
|
||||
if not messages:
|
||||
return TurnKind.UNKNOWN
|
||||
last = messages[-1]
|
||||
if not isinstance(last, dict) or last.get("role") != "user":
|
||||
return TurnKind.UNKNOWN
|
||||
|
||||
content = last.get("content")
|
||||
if isinstance(content, str):
|
||||
return TurnKind.NEW_USER_ASK if content.strip() else TurnKind.UNKNOWN
|
||||
if not isinstance(content, list) or not content:
|
||||
return TurnKind.UNKNOWN
|
||||
|
||||
saw_tool_result = False
|
||||
saw_error = False
|
||||
for block in content:
|
||||
if not isinstance(block, dict):
|
||||
return TurnKind.UNKNOWN
|
||||
btype = block.get("type")
|
||||
if btype == "tool_result":
|
||||
saw_tool_result = True
|
||||
if block.get("is_error") is True:
|
||||
saw_error = True
|
||||
elif btype == "text":
|
||||
# Fresh user text alongside (or instead of) tool results means
|
||||
# the user interjected — treat as a new ask.
|
||||
return TurnKind.NEW_USER_ASK
|
||||
elif btype in ("image", "document"):
|
||||
return TurnKind.NEW_USER_ASK
|
||||
# Unknown block types are ignored rather than guessed at.
|
||||
|
||||
if saw_error:
|
||||
return TurnKind.ERROR_CONTINUATION
|
||||
if saw_tool_result:
|
||||
return TurnKind.MECHANICAL_CONTINUATION
|
||||
return TurnKind.UNKNOWN
|
||||
|
||||
|
||||
def steering_text(level: int) -> str | None:
|
||||
"""The full steering block for a verbosity level, or None for level 0."""
|
||||
text = _VERBOSITY_LEVELS.get(level)
|
||||
if text is None:
|
||||
return None
|
||||
return f"{_STEERING_SENTINEL}\n{text}\n{_STEERING_SUFFIX}"
|
||||
|
||||
|
||||
def apply_verbosity_steering(body: dict[str, Any], level: int) -> bool:
|
||||
"""Append the steering block to the tail of the system prompt.
|
||||
|
||||
Appending AFTER the last system block keeps any ``cache_control``
|
||||
breakpoint on an earlier block intact — the cached prefix is unchanged
|
||||
and only the (small, byte-stable) steering block is reprocessed.
|
||||
|
||||
A string system prompt is converted to block form so the original text
|
||||
keeps its exact bytes as the first block.
|
||||
"""
|
||||
text = steering_text(level)
|
||||
if text is None:
|
||||
return False
|
||||
|
||||
system = body.get("system")
|
||||
if system is None:
|
||||
body["system"] = [{"type": "text", "text": text}]
|
||||
return True
|
||||
if isinstance(system, str):
|
||||
body["system"] = [
|
||||
{"type": "text", "text": system},
|
||||
{"type": "text", "text": text},
|
||||
]
|
||||
return True
|
||||
if isinstance(system, list):
|
||||
for block in system:
|
||||
if isinstance(block, dict) and block.get("text", "").startswith(_STEERING_SENTINEL):
|
||||
if block["text"] == text:
|
||||
return False # already applied at this level
|
||||
block["text"] = text # level changed mid-session
|
||||
return True
|
||||
system.append({"type": "text", "text": text})
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def route_effort(
|
||||
body: dict[str, Any],
|
||||
kind: TurnKind,
|
||||
settings: OutputShaperSettings,
|
||||
) -> list[str]:
|
||||
"""Lower thinking/effort spend on mechanical continuations.
|
||||
|
||||
Returns labels for each mutation made (empty list = untouched).
|
||||
"""
|
||||
if kind is not TurnKind.MECHANICAL_CONTINUATION:
|
||||
return []
|
||||
|
||||
labels: list[str] = []
|
||||
|
||||
# Modern lever: output_config.effort. Only lower a value the client
|
||||
# explicitly sent — presence proves the target model accepts the param.
|
||||
output_config = body.get("output_config")
|
||||
if isinstance(output_config, dict):
|
||||
effort = output_config.get("effort")
|
||||
if (
|
||||
isinstance(effort, str)
|
||||
and effort in _EFFORT_RANK
|
||||
and _EFFORT_RANK[effort] > _EFFORT_RANK[settings.mechanical_effort]
|
||||
):
|
||||
output_config["effort"] = settings.mechanical_effort
|
||||
labels.append(f"output_shaper:effort:{effort}->{settings.mechanical_effort}")
|
||||
|
||||
# Legacy lever: clamp thinking.budget_tokens on models still using the
|
||||
# enabled/budget_tokens form. The type field itself is never touched.
|
||||
thinking = body.get("thinking")
|
||||
if isinstance(thinking, dict) and thinking.get("type") == "enabled":
|
||||
budget = thinking.get("budget_tokens")
|
||||
if isinstance(budget, int) and budget > LEGACY_THINKING_FLOOR:
|
||||
thinking["budget_tokens"] = LEGACY_THINKING_FLOOR
|
||||
labels.append(f"output_shaper:thinking_budget:{budget}->{LEGACY_THINKING_FLOOR}")
|
||||
|
||||
return labels
|
||||
|
||||
|
||||
def shape_request(
|
||||
body: dict[str, Any],
|
||||
settings: OutputShaperSettings | None = None,
|
||||
level_override: int | None = None,
|
||||
) -> ShapeResult:
|
||||
"""Apply all output-shaping levers to an Anthropic request body in place.
|
||||
|
||||
``level_override`` supersedes ``settings.verbosity_level`` when given — the
|
||||
handler passes the level resolved by :func:`resolve_verbosity_level` (learned
|
||||
profile / controller / env) so the body-mutating core stays level-agnostic.
|
||||
"""
|
||||
if settings is None:
|
||||
settings = OutputShaperSettings.from_env()
|
||||
result = ShapeResult()
|
||||
if not settings.enabled:
|
||||
return result
|
||||
|
||||
assert result.labels is not None # __post_init__ guarantees this
|
||||
|
||||
level = settings.verbosity_level if level_override is None else level_override
|
||||
if level > 0 and apply_verbosity_steering(body, level):
|
||||
result.changed = True
|
||||
result.labels.append(f"output_shaper:verbosity:L{level}")
|
||||
|
||||
if settings.effort_router_enabled:
|
||||
kind = classify_turn(body.get("messages", []))
|
||||
labels = route_effort(body, kind, settings)
|
||||
if labels:
|
||||
result.changed = True
|
||||
result.labels.extend(labels)
|
||||
logger.debug("OutputShaper: turn=%s mutations=%s", kind.value, labels)
|
||||
|
||||
return result
|
||||
|
|
@ -2464,6 +2464,30 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
global_output_tokens=m.tokens_output_total,
|
||||
)
|
||||
|
||||
# Output-side reduction (counterfactual estimate from the shaper's
|
||||
# ledger). Distinct from input compression above: these are OUTPUT
|
||||
# tokens the model didn't emit because we steered verbosity / routed
|
||||
# effort down. Always labelled estimated-vs-measured + a CI so it's
|
||||
# never mistaken for an exact count. Best-effort — never break /stats.
|
||||
output_reduction: dict[str, Any] = {"available": False}
|
||||
try:
|
||||
from headroom.proxy.output_savings import get_recorder
|
||||
|
||||
_oest = get_recorder().estimate()
|
||||
if _oest.n_requests > 0:
|
||||
output_reduction = {
|
||||
"available": True,
|
||||
"method": _oest.kind, # "measured" | "estimated"
|
||||
"tokens_saved": round(_oest.tokens_saved),
|
||||
"baseline_tokens": round(_oest.baseline_tokens),
|
||||
"reduction_percent": round(_oest.pct, 1),
|
||||
"ci_low_percent": round(_oest.ci_low_pct, 1),
|
||||
"ci_high_percent": round(_oest.ci_high_pct, 1),
|
||||
"requests": _oest.n_requests,
|
||||
}
|
||||
except Exception: # pragma: no cover - defensive
|
||||
pass
|
||||
|
||||
return {
|
||||
"summary": summary,
|
||||
"agent_usage": agent_usage,
|
||||
|
|
@ -2520,6 +2544,15 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"baseline caching is provider-native."
|
||||
),
|
||||
},
|
||||
"output_shaping": {
|
||||
**output_reduction,
|
||||
"description": (
|
||||
"OUTPUT tokens the model didn't emit because the shaper "
|
||||
"steered verbosity / routed effort down. Counterfactual — "
|
||||
"shown as an estimate (vs a learned baseline) or measured "
|
||||
"(A/B holdout), always with a confidence band."
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
"requests": {
|
||||
|
|
@ -2534,6 +2567,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"tokens": {
|
||||
"input": m.tokens_input_total,
|
||||
"output": m.tokens_output_total,
|
||||
"output_saved": output_reduction.get("tokens_saved", 0),
|
||||
"output_reduction_percent": output_reduction.get("reduction_percent", 0),
|
||||
"output_reduction": output_reduction,
|
||||
"saved": all_layers_tokens_saved,
|
||||
"proxy_compression_saved": proxy_compression_tokens,
|
||||
"cli_filtering_saved": cli_tokens_avoided,
|
||||
|
|
|
|||
108
headroom/proxy/verbosity_controller.py
Normal file
108
headroom/proxy/verbosity_controller.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
"""AIMD controller for live verbosity adjustment.
|
||||
|
||||
The offline ``learn --verbosity`` pass sets the *starting* level. This controller
|
||||
tracks drift during a session and nudges the level from runtime signals, using
|
||||
the congestion-control intuition:
|
||||
|
||||
- **Additive increase** toward terser output: only after *sustained* "the user
|
||||
isn't reading this" pressure (a streak of TOO_MUCH signals) do we step the
|
||||
level up by one. Probing up is cheap to get wrong only slowly.
|
||||
- **Multiplicative-style decrease** on a TOO_LITTLE signal (the user asked for
|
||||
more): back off immediately by a level and enter a cooldown that suppresses
|
||||
re-escalation. Annoying the user is the expensive event — like congestion —
|
||||
so we react fast and then hold off.
|
||||
|
||||
The controller is a pure state machine over an abstract signal. Detecting the
|
||||
signals at the proxy (fast-skip from reply timing, interrupt from a cancelled
|
||||
stream) is the caller's job; keeping detection out of here makes the control
|
||||
logic deterministic and testable, and lets the live path enable only the
|
||||
signals it can measure reliably.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import asdict, dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class Signal(Enum):
|
||||
"""An abstract feedback signal about the last response's verbosity."""
|
||||
|
||||
TOO_MUCH = "too_much" # interrupted / fast-skipped → output went unread
|
||||
TOO_LITTLE = "too_little" # user asked to explain / expand
|
||||
NEUTRAL = "neutral" # engaged normally
|
||||
|
||||
|
||||
@dataclass
|
||||
class ControllerState:
|
||||
"""Per-conversation (or per-project) controller state."""
|
||||
|
||||
level: int
|
||||
up_streak: int = 0
|
||||
cooldown: int = 0
|
||||
|
||||
def to_dict(self) -> dict[str, int]:
|
||||
return asdict(self)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict[str, int]) -> ControllerState:
|
||||
return cls(
|
||||
level=int(d.get("level", 2)),
|
||||
up_streak=int(d.get("up_streak", 0)),
|
||||
cooldown=int(d.get("cooldown", 0)),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VerbosityController:
|
||||
"""Pure AIMD controller. ``observe`` maps (state, signal) → new state."""
|
||||
|
||||
floor: int = 1
|
||||
ceil: int = 4
|
||||
probe_threshold: int = 3 # consecutive TOO_MUCH before stepping up
|
||||
cooldown_turns: int = 5 # turns after a back-off during which we don't re-probe
|
||||
|
||||
def observe(self, state: ControllerState, signal: Signal) -> ControllerState:
|
||||
level = state.level
|
||||
up_streak = state.up_streak
|
||||
cooldown = max(0, state.cooldown - 1) # one turn elapses per observation
|
||||
|
||||
if signal is Signal.TOO_LITTLE:
|
||||
# Fast back-off: drop a level immediately and suppress re-escalation.
|
||||
return ControllerState(
|
||||
level=max(self.floor, level - 1),
|
||||
up_streak=0,
|
||||
cooldown=self.cooldown_turns,
|
||||
)
|
||||
|
||||
if signal is Signal.TOO_MUCH:
|
||||
if cooldown > 0:
|
||||
# Recently backed off — don't re-terse yet; keep cooling down.
|
||||
return ControllerState(level=level, up_streak=0, cooldown=cooldown)
|
||||
up_streak += 1
|
||||
if up_streak >= self.probe_threshold and level < self.ceil:
|
||||
return ControllerState(level=level + 1, up_streak=0, cooldown=0)
|
||||
return ControllerState(level=level, up_streak=up_streak, cooldown=cooldown)
|
||||
|
||||
# NEUTRAL: engagement resets the upward streak (we require *consecutive*
|
||||
# pressure) and lets any cooldown tick down.
|
||||
return ControllerState(level=level, up_streak=0, cooldown=cooldown)
|
||||
|
||||
|
||||
def load_state(path: Path, default_level: int, floor: int, ceil: int) -> ControllerState:
|
||||
"""Load controller state, clamped to [floor, ceil]; seed from default."""
|
||||
try:
|
||||
d = json.loads(Path(path).read_text())
|
||||
state = ControllerState.from_dict(d)
|
||||
except (OSError, json.JSONDecodeError, ValueError):
|
||||
state = ControllerState(level=default_level)
|
||||
state.level = max(floor, min(ceil, state.level))
|
||||
return state
|
||||
|
||||
|
||||
def save_state(path: Path, state: ControllerState) -> None:
|
||||
p = Path(path)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text(json.dumps(state.to_dict(), separators=(",", ":")))
|
||||
235
scripts/eval_output_shaper.py
Normal file
235
scripts/eval_output_shaper.py
Normal file
|
|
@ -0,0 +1,235 @@
|
|||
"""Live before/after eval for the output shaper.
|
||||
|
||||
Sends the SAME request to the Anthropic API twice — once as a client would
|
||||
send it (baseline) and once after `shape_request` rewrites it (exactly what
|
||||
the proxy forwards upstream) — and compares `usage.output_tokens`, which
|
||||
includes thinking tokens.
|
||||
|
||||
Scenario A (verbosity steering): a complex code-review ask. Baseline vs
|
||||
verbosity levels 2 and 3.
|
||||
|
||||
Scenario B (effort routing): an agentic transcript whose last message is a
|
||||
clean tool_result (mechanical continuation) with `output_config.effort` set
|
||||
to "xhigh" the way Claude Code pins it. The shaper lowers effort to "low"
|
||||
for this turn only.
|
||||
|
||||
Usage:
|
||||
source .venv/bin/activate && python scripts/eval_output_shaper.py
|
||||
Requires ANTHROPIC_API_KEY in the environment or in ./.env.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
import anthropic # noqa: E402
|
||||
|
||||
from headroom.proxy.output_shaper import OutputShaperSettings, shape_request # noqa: E402
|
||||
|
||||
MODEL = "claude-opus-4-8"
|
||||
TRIALS = 2
|
||||
|
||||
BUGGY_CODE = '''\
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
|
||||
class TTLCache:
|
||||
"""LRU cache with per-entry TTL."""
|
||||
|
||||
def __init__(self, max_size=128, ttl=300):
|
||||
self.max_size = max_size
|
||||
self.ttl = ttl
|
||||
self._store = OrderedDict()
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def get(self, key, now):
|
||||
entry = self._store.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
value, expires_at = entry
|
||||
if now > expires_at:
|
||||
del self._store[key]
|
||||
return None
|
||||
self._store.move_to_end(key)
|
||||
return value
|
||||
|
||||
def put(self, key, value, now):
|
||||
with self._lock:
|
||||
if key in self._store:
|
||||
self._store.move_to_end(key)
|
||||
self._store[key] = (value, now + self.ttl)
|
||||
if len(self._store) > self.max_size:
|
||||
self._store.popitem(last=True)
|
||||
|
||||
def cleanup(self, now):
|
||||
for key, (_, expires_at) in self._store.items():
|
||||
if now > expires_at:
|
||||
del self._store[key]
|
||||
'''
|
||||
|
||||
|
||||
def load_env() -> None:
|
||||
env_path = Path(__file__).resolve().parent.parent / ".env"
|
||||
if not env_path.exists() or os.environ.get("ANTHROPIC_API_KEY"):
|
||||
return
|
||||
for line in env_path.read_text().splitlines():
|
||||
line = line.strip()
|
||||
if line and not line.startswith("#") and "=" in line:
|
||||
key, _, value = line.partition("=")
|
||||
value = value.strip().strip("'\"")
|
||||
os.environ.setdefault(key.strip(), value)
|
||||
|
||||
|
||||
def scenario_a_body() -> dict[str, Any]:
|
||||
"""Complex single-turn ask — exercises verbosity steering."""
|
||||
return {
|
||||
"model": MODEL,
|
||||
"max_tokens": 8000,
|
||||
"system": "You are a senior Python engineer doing code review.",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"Review this cache implementation. Identify every bug and "
|
||||
"thread-safety issue, then show how to fix each one:\n\n"
|
||||
f"```python\n{BUGGY_CODE}```"
|
||||
),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def scenario_b_body() -> dict[str, Any]:
|
||||
"""Agentic mechanical continuation — exercises effort routing."""
|
||||
return {
|
||||
"model": MODEL,
|
||||
"max_tokens": 8000,
|
||||
"thinking": {"type": "adaptive"},
|
||||
"output_config": {"effort": "xhigh"},
|
||||
"system": (
|
||||
"You are a coding agent. Use the Read tool to inspect files, then "
|
||||
"report findings concisely."
|
||||
),
|
||||
"tools": [
|
||||
{
|
||||
"name": "Read",
|
||||
"description": "Read a file from the repository.",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"path": {"type": "string"}},
|
||||
"required": ["path"],
|
||||
},
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Check whether cache.py has thread-safety issues.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Reading cache.py first."},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_eval_01",
|
||||
"name": "Read",
|
||||
"input": {"path": "cache.py"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_eval_01",
|
||||
"content": BUGGY_CODE,
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def run(client: anthropic.Anthropic, body: dict[str, Any]) -> dict[str, int]:
|
||||
# The installed SDK may predate output_config as a typed kwarg; the API
|
||||
# accepts it either way, so pass it through extra_body.
|
||||
body = dict(body)
|
||||
extra_body = None
|
||||
if "output_config" in body:
|
||||
extra_body = {"output_config": body.pop("output_config")}
|
||||
response = client.messages.create(**body, extra_body=extra_body)
|
||||
if response.stop_reason == "refusal":
|
||||
raise RuntimeError("request was refused by safety classifiers")
|
||||
return {
|
||||
"input_tokens": response.usage.input_tokens,
|
||||
"output_tokens": response.usage.output_tokens,
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
load_env()
|
||||
if not os.environ.get("ANTHROPIC_API_KEY"):
|
||||
print("ANTHROPIC_API_KEY not found (env or .env)", file=sys.stderr)
|
||||
return 1
|
||||
client = anthropic.Anthropic()
|
||||
which = sys.argv[1].upper() if len(sys.argv) > 1 else "ALL"
|
||||
|
||||
conditions: list[tuple[str, str, dict[str, Any]]] = []
|
||||
|
||||
if which in ("A", "ALL"):
|
||||
# Scenario A: baseline vs steered.
|
||||
conditions.append(("A:verbosity", "baseline", scenario_a_body()))
|
||||
for level in (2, 3):
|
||||
body = scenario_a_body()
|
||||
shape_request(body, OutputShaperSettings(enabled=True, verbosity_level=level))
|
||||
conditions.append(("A:verbosity", f"shaped L{level}", body))
|
||||
|
||||
if which in ("B", "ALL"):
|
||||
# Scenario B: baseline (effort=xhigh) vs shaped (effort routed to low).
|
||||
conditions.append(("B:effort-routing", "baseline xhigh", scenario_b_body()))
|
||||
body = scenario_b_body()
|
||||
result = shape_request(body, OutputShaperSettings(enabled=True, verbosity_level=0))
|
||||
assert body["output_config"]["effort"] == "low", result.labels
|
||||
conditions.append(("B:effort-routing", "shaped low", body))
|
||||
|
||||
print(f"model={MODEL} trials={TRIALS}\n")
|
||||
print(f"{'scenario':<18} {'condition':<16} {'trial':<6} {'in_tok':>7} {'out_tok':>8}")
|
||||
print("-" * 60)
|
||||
|
||||
results: dict[tuple[str, str], list[int]] = {}
|
||||
for scenario, condition, body in conditions:
|
||||
for trial in range(1, TRIALS + 1):
|
||||
usage = run(client, copy.deepcopy(body))
|
||||
results.setdefault((scenario, condition), []).append(usage["output_tokens"])
|
||||
print(
|
||||
f"{scenario:<18} {condition:<16} {trial:<6} "
|
||||
f"{usage['input_tokens']:>7} {usage['output_tokens']:>8}"
|
||||
)
|
||||
|
||||
print("\n=== Summary (mean output tokens, reduction vs baseline) ===")
|
||||
baselines: dict[str, float] = {}
|
||||
for (scenario, condition), outs in results.items():
|
||||
if condition.startswith("baseline"):
|
||||
baselines[scenario] = statistics.mean(outs)
|
||||
for (scenario, condition), outs in results.items():
|
||||
mean = statistics.mean(outs)
|
||||
base = baselines.get(scenario, 0)
|
||||
if condition.startswith("baseline") or not base:
|
||||
print(f"{scenario:<18} {condition:<16} {mean:>8.0f} (baseline)")
|
||||
else:
|
||||
pct = (base - mean) / base * 100
|
||||
print(f"{scenario:<18} {condition:<16} {mean:>8.0f} ({pct:+.1f}% vs baseline)")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
287
tests/test_output_savings.py
Normal file
287
tests/test_output_savings.py
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
"""Tests for headroom.proxy.output_savings — the counterfactual estimator."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from headroom.proxy.output_savings import (
|
||||
BaselineModel,
|
||||
SavingsLedger,
|
||||
assign_arm,
|
||||
conversation_key_from_body,
|
||||
echo_ratio,
|
||||
input_bucket,
|
||||
model_family,
|
||||
stratum_key,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stratification primitives
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStratification:
|
||||
def test_input_buckets_monotone(self):
|
||||
assert input_bucket(0) == "xs"
|
||||
assert input_bucket(1_999) == "xs"
|
||||
assert input_bucket(2_000) == "s"
|
||||
assert input_bucket(8_000) == "m"
|
||||
assert input_bucket(32_000) == "l"
|
||||
assert input_bucket(200_000) == "xl"
|
||||
|
||||
def test_model_family_collapses_point_releases(self):
|
||||
assert model_family("claude-opus-4-8") == "opus"
|
||||
assert model_family("claude-opus-4-7") == "opus"
|
||||
assert model_family("claude-sonnet-4-6") == "sonnet"
|
||||
assert model_family("gpt-4o") == "gpt"
|
||||
assert model_family("something-weird") == "other"
|
||||
|
||||
def test_stratum_key_is_most_to_least_specific(self):
|
||||
key = stratum_key(
|
||||
turn_kind="new_user_ask", input_tokens=5000, model="claude-opus-4-8", has_tools=True
|
||||
)
|
||||
assert key == "opus|new_user_ask|s|tools"
|
||||
|
||||
def test_stratum_key_distinguishes_tools(self):
|
||||
a = stratum_key(turn_kind="x", input_tokens=100, model="m", has_tools=True)
|
||||
b = stratum_key(turn_kind="x", input_tokens=100, model="m", has_tools=False)
|
||||
assert a != b
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# holdout arm assignment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestArmAssignment:
|
||||
def test_zero_holdout_always_treatment(self):
|
||||
assert assign_arm("anything", 0.0) == "treatment"
|
||||
|
||||
def test_full_holdout_always_control(self):
|
||||
assert assign_arm("anything", 1.0) == "control"
|
||||
|
||||
def test_assignment_is_stable_for_same_key(self):
|
||||
assert assign_arm("conv-123", 0.5) == assign_arm("conv-123", 0.5)
|
||||
|
||||
def test_roughly_matches_fraction(self):
|
||||
keys = [f"conv-{i}" for i in range(4000)]
|
||||
control = sum(1 for k in keys if assign_arm(k, 0.1) == "control")
|
||||
# 10% holdout over 4000 keys — allow generous slack for hash noise.
|
||||
assert 250 < control < 550
|
||||
|
||||
def test_conversation_key_stable_across_turns(self):
|
||||
first = {
|
||||
"model": "claude-opus-4-8",
|
||||
"messages": [{"role": "user", "content": "build a cache"}],
|
||||
}
|
||||
later = {
|
||||
"model": "claude-opus-4-8",
|
||||
"messages": [
|
||||
{"role": "user", "content": "build a cache"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
{"role": "user", "content": [{"type": "tool_result", "content": "x"}]},
|
||||
],
|
||||
}
|
||||
assert conversation_key_from_body(first) == conversation_key_from_body(later)
|
||||
|
||||
def test_conversation_key_differs_by_first_message(self):
|
||||
a = {"model": "m", "messages": [{"role": "user", "content": "task A"}]}
|
||||
b = {"model": "m", "messages": [{"role": "user", "content": "task B"}]}
|
||||
assert conversation_key_from_body(a) != conversation_key_from_body(b)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# baseline model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBaselineModel:
|
||||
def test_observe_and_lookup_exact(self):
|
||||
m = BaselineModel()
|
||||
for v in (100, 200, 300):
|
||||
m.observe("opus|new_user_ask|s|tools", v)
|
||||
mean, var, n = m.lookup("opus|new_user_ask|s|tools")
|
||||
assert mean == 200.0
|
||||
assert n == 3
|
||||
assert var > 0
|
||||
|
||||
def test_lookup_backs_off_to_prefix(self):
|
||||
m = BaselineModel()
|
||||
m.observe("opus|new_user_ask|s|tools", 500)
|
||||
# Query a sibling stratum (different tools flag) — backs off on prefix.
|
||||
mean, _, n = m.lookup("opus|new_user_ask|s|notools")
|
||||
assert mean == 500.0
|
||||
assert n == 1
|
||||
|
||||
def test_lookup_falls_back_to_global(self):
|
||||
m = BaselineModel()
|
||||
m.observe("opus|a|s|tools", 100)
|
||||
m.observe("sonnet|b|m|notools", 300)
|
||||
mean, _, n = m.lookup("gpt|totally|xl|tools")
|
||||
assert mean == 200.0 # global mean of 100 and 300
|
||||
assert n == 2
|
||||
|
||||
def test_roundtrip_serialization(self):
|
||||
m = BaselineModel()
|
||||
for v in (10, 20, 30):
|
||||
m.observe("k|a|s|tools", v)
|
||||
m2 = BaselineModel.from_dict(m.to_dict())
|
||||
assert m2.lookup("k|a|s|tools") == m.lookup("k|a|s|tools")
|
||||
assert m2.total_samples == 3
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# synthetic-control estimate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEstimateFromBaseline:
|
||||
def _ledger_with_baseline(self, baseline_val: float, n: int = 50) -> SavingsLedger:
|
||||
ledger = SavingsLedger()
|
||||
for _ in range(n):
|
||||
ledger.baseline.observe("opus|new_user_ask|s|tools", baseline_val)
|
||||
return ledger
|
||||
|
||||
def test_positive_savings_when_treatment_below_baseline(self):
|
||||
ledger = self._ledger_with_baseline(1000.0)
|
||||
for _ in range(20):
|
||||
ledger.record("treatment", "opus|new_user_ask|s|tools", 700)
|
||||
est = ledger.estimate_from_baseline()
|
||||
assert est.kind == "estimated"
|
||||
assert est.n_requests == 20
|
||||
# 20 requests * (1000 - 700) = 6000 tokens saved.
|
||||
assert abs(est.tokens_saved - 6000) < 1e-6
|
||||
assert abs(est.pct - 30.0) < 1e-6
|
||||
|
||||
def test_signed_delta_not_clamped(self):
|
||||
# A treatment request LARGER than baseline must reduce the total, not
|
||||
# be clamped to zero (clamping would bias the estimate upward).
|
||||
ledger = self._ledger_with_baseline(1000.0)
|
||||
ledger.record("treatment", "opus|new_user_ask|s|tools", 700)
|
||||
ledger.record("treatment", "opus|new_user_ask|s|tools", 1400)
|
||||
est = ledger.estimate_from_baseline()
|
||||
# (1000-700) + (1000-1400) = 300 - 400 = -100
|
||||
assert abs(est.tokens_saved - (-100)) < 1e-6
|
||||
|
||||
def test_zero_baseline_samples_yields_zero(self):
|
||||
ledger = SavingsLedger()
|
||||
ledger.record("treatment", "opus|x|s|tools", 500)
|
||||
est = ledger.estimate_from_baseline()
|
||||
# No baseline at all -> global is empty -> nothing contributes.
|
||||
assert est.n_requests == 0
|
||||
assert est.tokens_saved == 0.0
|
||||
|
||||
def test_ci_band_brackets_point_estimate(self):
|
||||
ledger = SavingsLedger()
|
||||
for v in (900, 1000, 1100):
|
||||
for _ in range(20):
|
||||
ledger.baseline.observe("opus|new_user_ask|s|tools", v)
|
||||
for v in (600, 700, 800):
|
||||
for _ in range(20):
|
||||
ledger.record("treatment", "opus|new_user_ask|s|tools", v)
|
||||
est = ledger.estimate_from_baseline()
|
||||
assert est.ci_low_pct <= est.pct <= est.ci_high_pct
|
||||
assert est.ci_low_pct < est.ci_high_pct # nonzero band given spread
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# A/B measured estimate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEstimateFromHoldout:
|
||||
def test_none_without_control_data(self):
|
||||
ledger = SavingsLedger()
|
||||
ledger.record("treatment", "opus|x|s|tools", 500)
|
||||
assert ledger.estimate_from_holdout() is None
|
||||
|
||||
def test_measured_difference_of_means(self):
|
||||
ledger = SavingsLedger()
|
||||
for _ in range(30):
|
||||
ledger.record("control", "opus|new_user_ask|s|tools", 1000)
|
||||
ledger.record("treatment", "opus|new_user_ask|s|tools", 750)
|
||||
est = ledger.estimate_from_holdout()
|
||||
assert est is not None
|
||||
assert est.kind == "measured"
|
||||
# 30 * (1000 - 750) = 7500 saved; 25% of the 1000 baseline.
|
||||
assert abs(est.tokens_saved - 7500) < 1e-6
|
||||
assert abs(est.pct - 25.0) < 1e-6
|
||||
|
||||
def test_only_strata_present_in_both_arms_contribute(self):
|
||||
ledger = SavingsLedger()
|
||||
for _ in range(10):
|
||||
ledger.record("control", "opus|a|s|tools", 1000)
|
||||
ledger.record("treatment", "opus|a|s|tools", 800)
|
||||
# Treatment-only stratum must not contribute (no control to compare).
|
||||
ledger.record("treatment", "opus|b|m|notools", 50)
|
||||
est = ledger.estimate_from_holdout()
|
||||
assert est is not None
|
||||
assert est.n_requests == 10
|
||||
|
||||
def test_best_estimate_prefers_measured(self):
|
||||
ledger = SavingsLedger()
|
||||
for _ in range(10):
|
||||
ledger.baseline.observe("opus|a|s|tools", 1000)
|
||||
ledger.record("control", "opus|a|s|tools", 1000)
|
||||
ledger.record("treatment", "opus|a|s|tools", 900)
|
||||
assert ledger.best_estimate().kind == "measured"
|
||||
|
||||
def test_best_estimate_falls_back_to_estimated(self):
|
||||
ledger = SavingsLedger()
|
||||
for _ in range(10):
|
||||
ledger.baseline.observe("opus|a|s|tools", 1000)
|
||||
ledger.record("treatment", "opus|a|s|tools", 900)
|
||||
assert ledger.best_estimate().kind == "estimated"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# persistence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLedgerPersistence:
|
||||
def test_roundtrip(self, tmp_path):
|
||||
ledger = SavingsLedger()
|
||||
ledger.baseline.observe("opus|a|s|tools", 1000)
|
||||
ledger.record("treatment", "opus|a|s|tools", 800)
|
||||
ledger.record("control", "opus|a|s|tools", 1000)
|
||||
path = tmp_path / "savings.json"
|
||||
ledger.save(path)
|
||||
loaded = SavingsLedger.load(path)
|
||||
assert loaded.estimate_from_baseline().tokens_saved == (
|
||||
ledger.estimate_from_baseline().tokens_saved
|
||||
)
|
||||
assert loaded.estimate_from_holdout() is not None
|
||||
|
||||
def test_load_missing_returns_empty(self, tmp_path):
|
||||
ledger = SavingsLedger.load(tmp_path / "nope.json")
|
||||
assert ledger.baseline.total_samples == 0
|
||||
|
||||
def test_load_corrupt_returns_empty(self, tmp_path):
|
||||
p = tmp_path / "bad.json"
|
||||
p.write_text("{not json")
|
||||
ledger = SavingsLedger.load(p)
|
||||
assert ledger.baseline.total_samples == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# echo ratio (direct waste signal)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEchoRatio:
|
||||
def test_full_echo(self):
|
||||
ctx = "the quick brown fox jumps over the lazy dog every single time"
|
||||
assert echo_ratio(ctx, ctx, n=4) == 1.0
|
||||
|
||||
def test_no_echo(self):
|
||||
out = "completely unrelated words appearing nowhere within the given source context here"
|
||||
ctx = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda"
|
||||
assert echo_ratio(out, ctx, n=4) == 0.0
|
||||
|
||||
def test_partial_echo_between_zero_and_one(self):
|
||||
ctx = "alpha beta gamma delta epsilon zeta eta theta"
|
||||
out = "alpha beta gamma delta brand new tokens here now"
|
||||
r = echo_ratio(out, ctx, n=4)
|
||||
assert 0.0 < r < 1.0
|
||||
|
||||
def test_short_output_returns_zero(self):
|
||||
assert echo_ratio("a b", "a b c d e f g h", n=8) == 0.0
|
||||
53
tests/test_output_savings_cli.py
Normal file
53
tests/test_output_savings_cli.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
"""Smoke tests for the output-savings CLI and the outcome→ledger wiring."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.main import main
|
||||
from headroom.proxy.output_savings import SavingsRecorder, stratum_label
|
||||
|
||||
|
||||
def test_output_savings_empty(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("HEADROOM_WORKSPACE_DIR", str(tmp_path))
|
||||
result = CliRunner().invoke(main, ["output-savings"])
|
||||
assert result.exit_code == 0
|
||||
assert "No output-savings data" in result.output
|
||||
|
||||
|
||||
def test_output_savings_reports_estimate(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("HEADROOM_WORKSPACE_DIR", str(tmp_path))
|
||||
# Seed a baseline + treatment observations directly via the ledger.
|
||||
from headroom.proxy.output_savings import SavingsLedger
|
||||
|
||||
ledger = SavingsLedger()
|
||||
for _ in range(50):
|
||||
ledger.baseline.observe("opus|new_user_ask|s|tools", 1000)
|
||||
for _ in range(30):
|
||||
ledger.record("treatment", "opus|new_user_ask|s|tools", 700)
|
||||
ledger.save(tmp_path / "output_savings.json")
|
||||
|
||||
result = CliRunner().invoke(main, ["output-savings"])
|
||||
assert result.exit_code == 0
|
||||
assert "ESTIMATED" in result.output
|
||||
assert "Reduction:" in result.output
|
||||
assert "30.0%" in result.output
|
||||
|
||||
|
||||
def test_recorder_round_trips_via_labels(tmp_path):
|
||||
path = tmp_path / "savings.json"
|
||||
rec = SavingsRecorder(path, flush_every=1)
|
||||
# Baseline so the estimate has something to compare against.
|
||||
rec._ledger.baseline.observe("opus|new_user_ask|s|tools", 1000)
|
||||
labels = ["compress:smartcrush", stratum_label("treatment", "opus|new_user_ask|s|tools")]
|
||||
assert rec.record_from_labels(labels, output_tokens=600) is True
|
||||
assert rec.record_from_labels(["no-shaper-label"], output_tokens=999) is False
|
||||
est = rec.estimate()
|
||||
assert est.n_requests == 1
|
||||
assert est.tokens_saved == 400 # 1000 - 600
|
||||
|
||||
# Persisted to disk.
|
||||
data = json.loads(path.read_text())
|
||||
assert "treatment" in data
|
||||
284
tests/test_output_shaper.py
Normal file
284
tests/test_output_shaper.py
Normal file
|
|
@ -0,0 +1,284 @@
|
|||
"""Tests for headroom.proxy.output_shaper.
|
||||
|
||||
Covers turn classification (structural only), cache-safe verbosity steering,
|
||||
effort routing on mechanical continuations, and the env-driven gate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from typing import Any
|
||||
|
||||
from headroom.proxy.output_shaper import (
|
||||
LEGACY_THINKING_FLOOR,
|
||||
OutputShaperSettings,
|
||||
TurnKind,
|
||||
apply_verbosity_steering,
|
||||
classify_turn,
|
||||
route_effort,
|
||||
shape_request,
|
||||
steering_text,
|
||||
)
|
||||
|
||||
ENABLED = OutputShaperSettings(enabled=True)
|
||||
|
||||
|
||||
def _tool_result(is_error: bool = False) -> dict[str, Any]:
|
||||
block: dict[str, Any] = {
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_01",
|
||||
"content": "ok",
|
||||
}
|
||||
if is_error:
|
||||
block["is_error"] = True
|
||||
return block
|
||||
|
||||
|
||||
def _mechanical_messages() -> list[dict[str, Any]]:
|
||||
return [
|
||||
{"role": "user", "content": "fix the bug in foo.py"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Reading the file."},
|
||||
{"type": "tool_use", "id": "toolu_01", "name": "Read", "input": {}},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": [_tool_result()]},
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# classify_turn
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClassifyTurn:
|
||||
def test_string_user_message_is_new_ask(self):
|
||||
assert classify_turn([{"role": "user", "content": "explain this"}]) == TurnKind.NEW_USER_ASK
|
||||
|
||||
def test_clean_tool_result_is_mechanical(self):
|
||||
assert classify_turn(_mechanical_messages()) == TurnKind.MECHANICAL_CONTINUATION
|
||||
|
||||
def test_multiple_clean_tool_results_are_mechanical(self):
|
||||
msgs = _mechanical_messages()
|
||||
msgs[-1]["content"].append(_tool_result())
|
||||
assert classify_turn(msgs) == TurnKind.MECHANICAL_CONTINUATION
|
||||
|
||||
def test_error_tool_result_is_error_continuation(self):
|
||||
msgs = _mechanical_messages()
|
||||
msgs[-1]["content"] = [_tool_result(), _tool_result(is_error=True)]
|
||||
assert classify_turn(msgs) == TurnKind.ERROR_CONTINUATION
|
||||
|
||||
def test_text_block_alongside_tool_result_is_new_ask(self):
|
||||
msgs = _mechanical_messages()
|
||||
msgs[-1]["content"].append({"type": "text", "text": "also check bar.py"})
|
||||
assert classify_turn(msgs) == TurnKind.NEW_USER_ASK
|
||||
|
||||
def test_image_block_is_new_ask(self):
|
||||
msgs = [{"role": "user", "content": [{"type": "image", "source": {}}]}]
|
||||
assert classify_turn(msgs) == TurnKind.NEW_USER_ASK
|
||||
|
||||
def test_assistant_last_is_unknown(self):
|
||||
msgs = [{"role": "assistant", "content": "hello"}]
|
||||
assert classify_turn(msgs) == TurnKind.UNKNOWN
|
||||
|
||||
def test_empty_messages_is_unknown(self):
|
||||
assert classify_turn([]) == TurnKind.UNKNOWN
|
||||
|
||||
def test_empty_content_list_is_unknown(self):
|
||||
assert classify_turn([{"role": "user", "content": []}]) == TurnKind.UNKNOWN
|
||||
|
||||
def test_whitespace_string_content_is_unknown(self):
|
||||
assert classify_turn([{"role": "user", "content": " "}]) == TurnKind.UNKNOWN
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apply_verbosity_steering
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestVerbositySteering:
|
||||
def test_level_zero_is_noop(self):
|
||||
body = {"system": "You are helpful."}
|
||||
assert apply_verbosity_steering(body, 0) is False
|
||||
assert body["system"] == "You are helpful."
|
||||
|
||||
def test_string_system_converted_to_blocks_with_original_bytes_first(self):
|
||||
body = {"system": "You are helpful."}
|
||||
assert apply_verbosity_steering(body, 2) is True
|
||||
assert body["system"][0] == {"type": "text", "text": "You are helpful."}
|
||||
assert body["system"][1]["text"] == steering_text(2)
|
||||
|
||||
def test_missing_system_creates_steering_only_block(self):
|
||||
body: dict[str, Any] = {}
|
||||
assert apply_verbosity_steering(body, 2) is True
|
||||
assert body["system"] == [{"type": "text", "text": steering_text(2)}]
|
||||
|
||||
def test_block_system_appends_after_cache_control(self):
|
||||
cached = {
|
||||
"type": "text",
|
||||
"text": "Big system prompt.",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
body = {"system": [copy.deepcopy(cached)]}
|
||||
assert apply_verbosity_steering(body, 2) is True
|
||||
# The cached block is byte-identical and still first — prefix intact.
|
||||
assert body["system"][0] == cached
|
||||
assert body["system"][1] == {"type": "text", "text": steering_text(2)}
|
||||
# Our block carries no cache_control (breakpoints are a scarce resource).
|
||||
assert "cache_control" not in body["system"][1]
|
||||
|
||||
def test_idempotent_at_same_level(self):
|
||||
body = {"system": [{"type": "text", "text": "Sys."}]}
|
||||
assert apply_verbosity_steering(body, 2) is True
|
||||
snapshot = copy.deepcopy(body)
|
||||
assert apply_verbosity_steering(body, 2) is False
|
||||
assert body == snapshot
|
||||
|
||||
def test_level_change_replaces_block_in_place(self):
|
||||
body = {"system": [{"type": "text", "text": "Sys."}]}
|
||||
apply_verbosity_steering(body, 2)
|
||||
assert apply_verbosity_steering(body, 4) is True
|
||||
steering_blocks = [
|
||||
b for b in body["system"] if b["text"].startswith("<headroom_output_shaping>")
|
||||
]
|
||||
assert len(steering_blocks) == 1
|
||||
assert steering_blocks[0]["text"] == steering_text(4)
|
||||
|
||||
def test_steering_text_is_deterministic(self):
|
||||
for level in (1, 2, 3, 4):
|
||||
assert steering_text(level) == steering_text(level)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# route_effort
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRouteEffort:
|
||||
def test_lowers_explicit_effort_on_mechanical_turn(self):
|
||||
body = {"output_config": {"effort": "xhigh"}}
|
||||
labels = route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED)
|
||||
assert body["output_config"]["effort"] == "low"
|
||||
assert labels == ["output_shaper:effort:xhigh->low"]
|
||||
|
||||
def test_never_injects_effort_when_absent(self):
|
||||
body: dict[str, Any] = {"messages": []}
|
||||
labels = route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED)
|
||||
assert "output_config" not in body
|
||||
assert labels == []
|
||||
|
||||
def test_effort_untouched_on_new_ask(self):
|
||||
body = {"output_config": {"effort": "xhigh"}}
|
||||
assert route_effort(body, TurnKind.NEW_USER_ASK, ENABLED) == []
|
||||
assert body["output_config"]["effort"] == "xhigh"
|
||||
|
||||
def test_effort_untouched_on_error_continuation(self):
|
||||
body = {"output_config": {"effort": "xhigh"}}
|
||||
assert route_effort(body, TurnKind.ERROR_CONTINUATION, ENABLED) == []
|
||||
assert body["output_config"]["effort"] == "xhigh"
|
||||
|
||||
def test_effort_already_at_target_untouched(self):
|
||||
body = {"output_config": {"effort": "low"}}
|
||||
assert route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED) == []
|
||||
|
||||
def test_unknown_effort_value_untouched(self):
|
||||
body = {"output_config": {"effort": "turbo"}}
|
||||
assert route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED) == []
|
||||
assert body["output_config"]["effort"] == "turbo"
|
||||
|
||||
def test_configurable_mechanical_effort(self):
|
||||
settings = OutputShaperSettings(enabled=True, mechanical_effort="medium")
|
||||
body = {"output_config": {"effort": "xhigh"}}
|
||||
route_effort(body, TurnKind.MECHANICAL_CONTINUATION, settings)
|
||||
assert body["output_config"]["effort"] == "medium"
|
||||
|
||||
def test_legacy_thinking_budget_clamped(self):
|
||||
body = {"thinking": {"type": "enabled", "budget_tokens": 32000}}
|
||||
labels = route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED)
|
||||
assert body["thinking"]["budget_tokens"] == LEGACY_THINKING_FLOOR
|
||||
assert body["thinking"]["type"] == "enabled" # never toggled
|
||||
assert labels == [f"output_shaper:thinking_budget:32000->{LEGACY_THINKING_FLOOR}"]
|
||||
|
||||
def test_legacy_budget_at_floor_untouched(self):
|
||||
body = {"thinking": {"type": "enabled", "budget_tokens": LEGACY_THINKING_FLOOR}}
|
||||
assert route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED) == []
|
||||
|
||||
def test_adaptive_thinking_untouched(self):
|
||||
body = {"thinking": {"type": "adaptive"}}
|
||||
assert route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED) == []
|
||||
assert body["thinking"] == {"type": "adaptive"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# shape_request (end to end)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShapeRequest:
|
||||
def test_disabled_is_noop(self):
|
||||
body = {
|
||||
"system": "Sys.",
|
||||
"messages": _mechanical_messages(),
|
||||
"output_config": {"effort": "xhigh"},
|
||||
}
|
||||
snapshot = copy.deepcopy(body)
|
||||
result = shape_request(body, OutputShaperSettings(enabled=False))
|
||||
assert result.changed is False
|
||||
assert body == snapshot
|
||||
|
||||
def test_enabled_applies_steering_and_effort_routing(self):
|
||||
body = {
|
||||
"system": "Sys.",
|
||||
"messages": _mechanical_messages(),
|
||||
"output_config": {"effort": "xhigh"},
|
||||
"thinking": {"type": "adaptive"},
|
||||
}
|
||||
result = shape_request(body, ENABLED)
|
||||
assert result.changed is True
|
||||
assert result.labels == [
|
||||
"output_shaper:verbosity:L2",
|
||||
"output_shaper:effort:xhigh->low",
|
||||
]
|
||||
assert body["output_config"]["effort"] == "low"
|
||||
assert body["system"][1]["text"] == steering_text(2)
|
||||
|
||||
def test_new_ask_gets_steering_but_keeps_effort(self):
|
||||
body = {
|
||||
"system": "Sys.",
|
||||
"messages": [{"role": "user", "content": "design a cache layer"}],
|
||||
"output_config": {"effort": "xhigh"},
|
||||
}
|
||||
result = shape_request(body, ENABLED)
|
||||
assert result.labels == ["output_shaper:verbosity:L2"]
|
||||
assert body["output_config"]["effort"] == "xhigh"
|
||||
|
||||
def test_second_pass_is_stable(self):
|
||||
body = {"system": "Sys.", "messages": _mechanical_messages()}
|
||||
shape_request(body, ENABLED)
|
||||
snapshot = copy.deepcopy(body)
|
||||
result = shape_request(body, ENABLED)
|
||||
assert result.changed is False
|
||||
assert body == snapshot
|
||||
|
||||
def test_from_env_defaults_off(self, monkeypatch):
|
||||
monkeypatch.delenv("HEADROOM_OUTPUT_SHAPER", raising=False)
|
||||
assert OutputShaperSettings.from_env().enabled is False
|
||||
|
||||
def test_from_env_enabled_with_overrides(self, monkeypatch):
|
||||
monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1")
|
||||
monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "3")
|
||||
monkeypatch.setenv("HEADROOM_MECHANICAL_EFFORT", "medium")
|
||||
settings = OutputShaperSettings.from_env()
|
||||
assert settings.enabled is True
|
||||
assert settings.verbosity_level == 3
|
||||
assert settings.mechanical_effort == "medium"
|
||||
|
||||
def test_from_env_clamps_bad_values(self, monkeypatch):
|
||||
monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "true")
|
||||
monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "99")
|
||||
monkeypatch.setenv("HEADROOM_MECHANICAL_EFFORT", "bogus")
|
||||
settings = OutputShaperSettings.from_env()
|
||||
assert settings.verbosity_level == 4
|
||||
assert settings.mechanical_effort == "low"
|
||||
98
tests/test_verbosity_controller.py
Normal file
98
tests/test_verbosity_controller.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
"""Tests for headroom.proxy.verbosity_controller — the AIMD state machine."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from headroom.proxy.verbosity_controller import (
|
||||
ControllerState,
|
||||
Signal,
|
||||
VerbosityController,
|
||||
load_state,
|
||||
save_state,
|
||||
)
|
||||
|
||||
CTRL = VerbosityController(floor=1, ceil=4, probe_threshold=3, cooldown_turns=5)
|
||||
|
||||
|
||||
def _run(signals, start=2):
|
||||
state = ControllerState(level=start)
|
||||
for s in signals:
|
||||
state = CTRL.observe(state, s)
|
||||
return state
|
||||
|
||||
|
||||
class TestAdditiveIncrease:
|
||||
def test_steps_up_only_after_threshold(self):
|
||||
state = _run([Signal.TOO_MUCH, Signal.TOO_MUCH])
|
||||
assert state.level == 2 # not yet at threshold
|
||||
assert state.up_streak == 2
|
||||
|
||||
def test_steps_up_at_threshold(self):
|
||||
state = _run([Signal.TOO_MUCH] * 3)
|
||||
assert state.level == 3
|
||||
assert state.up_streak == 0 # reset after stepping
|
||||
|
||||
def test_neutral_breaks_the_streak(self):
|
||||
state = _run([Signal.TOO_MUCH, Signal.TOO_MUCH, Signal.NEUTRAL, Signal.TOO_MUCH])
|
||||
assert state.level == 2 # streak was broken; only 1 consecutive at end
|
||||
assert state.up_streak == 1
|
||||
|
||||
def test_does_not_exceed_ceiling(self):
|
||||
state = _run([Signal.TOO_MUCH] * 30, start=4)
|
||||
assert state.level == 4
|
||||
|
||||
|
||||
class TestMultiplicativeDecrease:
|
||||
def test_too_little_backs_off_immediately(self):
|
||||
state = _run([Signal.TOO_LITTLE], start=3)
|
||||
assert state.level == 2
|
||||
assert state.cooldown == 5
|
||||
|
||||
def test_does_not_go_below_floor(self):
|
||||
state = _run([Signal.TOO_LITTLE] * 10, start=2)
|
||||
assert state.level == 1
|
||||
|
||||
def test_cooldown_suppresses_reescalation(self):
|
||||
# Back off, then immediately get TOO_MUCH pressure — must not re-terse
|
||||
# until the cooldown elapses.
|
||||
state = ControllerState(level=3)
|
||||
state = CTRL.observe(state, Signal.TOO_LITTLE) # → level 2, cooldown 5
|
||||
assert state.level == 2
|
||||
for _ in range(3): # would normally step up at 3, but we're cooling down
|
||||
state = CTRL.observe(state, Signal.TOO_MUCH)
|
||||
assert state.level == 2 # held
|
||||
|
||||
def test_reescalates_after_cooldown_expires(self):
|
||||
state = ControllerState(level=2, cooldown=5)
|
||||
# 5 neutral turns drain the cooldown...
|
||||
for _ in range(5):
|
||||
state = CTRL.observe(state, Signal.NEUTRAL)
|
||||
assert state.cooldown == 0
|
||||
# ...then sustained pressure can step up again.
|
||||
for _ in range(3):
|
||||
state = CTRL.observe(state, Signal.TOO_MUCH)
|
||||
assert state.level == 3
|
||||
|
||||
|
||||
class TestPersistence:
|
||||
def test_roundtrip(self, tmp_path):
|
||||
path = tmp_path / "ctrl.json"
|
||||
save_state(path, ControllerState(level=3, up_streak=2, cooldown=1))
|
||||
state = load_state(path, default_level=2, floor=1, ceil=4)
|
||||
assert state.level == 3
|
||||
assert state.up_streak == 2
|
||||
|
||||
def test_missing_uses_default(self, tmp_path):
|
||||
state = load_state(tmp_path / "nope.json", default_level=2, floor=1, ceil=4)
|
||||
assert state.level == 2
|
||||
|
||||
def test_corrupt_uses_default(self, tmp_path):
|
||||
p = tmp_path / "bad.json"
|
||||
p.write_text("{broken")
|
||||
state = load_state(p, default_level=3, floor=1, ceil=4)
|
||||
assert state.level == 3
|
||||
|
||||
def test_loaded_level_clamped(self, tmp_path):
|
||||
p = tmp_path / "ctrl.json"
|
||||
save_state(p, ControllerState(level=9))
|
||||
state = load_state(p, default_level=2, floor=1, ceil=4)
|
||||
assert state.level == 4
|
||||
251
tests/test_verbosity_learn.py
Normal file
251
tests/test_verbosity_learn.py
Normal file
|
|
@ -0,0 +1,251 @@
|
|||
"""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
|
||||
Loading…
Add table
Add a link
Reference in a new issue