akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor(output): isolate savings policy (#1947)

Browse source 
## Description

Extracts the output-savings stratification, holdout assignment,
conversation key, and transform-label helpers into a pure policy module
while preserving the existing `headroom.proxy.output_savings` public
imports. This keeps the estimator/ledger adapter focused on statistics
and persistence.

Closes #

## Type of Change

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

## Changes Made

- Added `headroom.proxy.output_savings_policy` for pure savings policy
helpers.
- Re-exported the moved helpers from `headroom.proxy.output_savings` to
keep callers stable.
- Added direct tests for the extracted policy boundary.
- Kept the LiteLLM callback compatibility shim required for repo-wide
type checking on fresh branches.

## Testing

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

### Test Output

```text
python -m pytest tests/test_output_savings_policy.py tests/test_output_savings.py tests/test_litellm_callback.py tests/test_compress_api.py::TestLiteLLMCallback -q
54 passed in 6.42s

python -m ruff check .
All checks passed!

python -m ruff format --check .
1095 files already formatted

python -m mypy headroom --ignore-missing-imports
Success: no issues found in 409 source files

gitleaks protect --staged --no-banner --redact
no leaks found
```

## Real Behavior Proof

- Environment: Windows, Python 3.13.13, local worktree based on
`headroomlabs/main`.
- Exact command / steps: ran the focused pytest set, Ruff lint/format
checks, mypy over `headroom`, and staged gitleaks scan.
- Observed result: all local checks passed; staged secret scan found no
leaks.
- Not tested: full CI matrix and deployment flows; those are covered by
GitHub Actions.

## Review Readiness

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

## Checklist

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

## Screenshots (if applicable)

N/A

## Additional Notes

Documentation and changelog updates are not applicable for this internal
refactor. GitHub reported existing Dependabot alerts on the default
branch during push; this PR does not change dependencies, and the staged
secret scan is clean.
This commit is contained in:
JD Davis 2026-07-11 04:42:11 +00:00 • committed by GitHub
parent b699bedf95
commit c29b4ba84f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 238 additions and 163 deletions
Download patch file Download diff file Expand all files Collapse all files

185
headroom/proxy/output_savings.py
View file

@ -38,155 +38,32 @@ 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, cast
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 _unwrap_response_create_body(body: dict[str, Any]) -> dict[str, Any]:
response = body.get("response")
if body.get("type") == "response.create" and isinstance(response, dict):
return cast("dict[str, Any]", response)
return body
def _stable_response_identifier(body: dict[str, Any]) -> str:
def _string_value(value: Any) -> str:
if isinstance(value, str):
return value
if isinstance(value, dict):
for key in ("id", "conversation_id", "session_id", "thread_id"):
nested = value.get(key)
if isinstance(nested, str) and nested:
return nested
return ""
for key in ("conversation", "conversation_id", "session_id", "thread_id"):
value = _string_value(body.get(key))
if value and value.lower() != "auto":
return f"{key}:{value}"
for container_key in ("client_metadata", "metadata"):
container = body.get(container_key)
if not isinstance(container, dict):
continue
for key in (
"conversation_id",
"conversation_key",
"session_id",
"thread_id",
"codex_session_id",
):
value = _string_value(container.get(key))
if value and value.lower() != "auto":
return f"{container_key}.{key}:{value}"
instructions = body.get("instructions")
if isinstance(instructions, str) and instructions:
return f"instructions:{instructions[:512]}"
return ""
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.
"""
body = _unwrap_response_create_body(body)
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
if "input" in body:
stable_response_key = _stable_response_identifier(body)
if stable_response_key:
seed += "\x00" + stable_response_key
elif not body.get("messages"):
seed += "\x00responses"
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"
from .output_savings_policy import (
assign_arm as assign_arm,
)
from .output_savings_policy import (
conversation_key_from_body as conversation_key_from_body,
)
from .output_savings_policy import (
input_bucket as input_bucket,
)
from .output_savings_policy import (
model_family as model_family,
)
from .output_savings_policy import (
parse_stratum_label,
)
from .output_savings_policy import (
stratum_key as stratum_key,
)
from .output_savings_policy import (
stratum_label as stratum_label,
)
@dataclass
@ -471,24 +348,6 @@ class SavingsLedger:
# 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.

156
headroom/proxy/output_savings_policy.py Normal file
View file

@ -0,0 +1,156 @@
"""Pure output-savings stratification and holdout policy helpers."""
from __future__ import annotations
import hashlib
from typing import Any, cast
# 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)
_STRATUM_LABEL = "output_shaper:stratum:"
_CONTROL_LABEL = "output_shaper:control:"
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-to-least specific so baseline 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 _unwrap_response_create_body(body: dict[str, Any]) -> dict[str, Any]:
response = body.get("response")
if body.get("type") == "response.create" and isinstance(response, dict):
return cast("dict[str, Any]", response)
return body
def _stable_response_identifier(body: dict[str, Any]) -> str:
def _string_value(value: Any) -> str:
if isinstance(value, str):
return value
if isinstance(value, dict):
for key in ("id", "conversation_id", "session_id", "thread_id"):
nested = value.get(key)
if isinstance(nested, str) and nested:
return nested
return ""
for key in ("conversation", "conversation_id", "session_id", "thread_id"):
value = _string_value(body.get(key))
if value and value.lower() != "auto":
return f"{key}:{value}"
for container_key in ("client_metadata", "metadata"):
container = body.get(container_key)
if not isinstance(container, dict):
continue
for key in (
"conversation_id",
"conversation_key",
"session_id",
"thread_id",
"codex_session_id",
):
value = _string_value(container.get(key))
if value and value.lower() != "auto":
return f"{container_key}.{key}:{value}"
instructions = body.get("instructions")
if isinstance(instructions, str) and instructions:
return f"instructions:{instructions[:512]}"
return ""
def conversation_key_from_body(body: dict[str, Any]) -> str:
"""Derive a conversation-stable key for holdout assignment."""
body = _unwrap_response_create_body(body)
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
if "input" in body:
stable_response_key = _stable_response_identifier(body)
if stable_response_key:
seed += "\x00" + stable_response_key
elif not body.get("messages"):
seed += "\x00responses"
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``."""
if holdout_fraction <= 0.0:
return "treatment"
if holdout_fraction >= 1.0:
return "control"
digest = hashlib.sha256(("arm:" + conversation_key).encode()).hexdigest()
frac = int(digest[:8], 16) / 0xFFFFFFFF
return "control" if frac < holdout_fraction else "treatment"
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

60
tests/test_output_savings_policy.py Normal file
View file

@ -0,0 +1,60 @@
"""Tests for pure output savings policy helpers."""
from __future__ import annotations
from headroom.proxy.output_savings_policy import (
assign_arm,
conversation_key_from_body,
input_bucket,
model_family,
parse_stratum_label,
stratum_key,
stratum_label,
)
def test_stratum_key_is_most_to_least_specific() -> None:
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_input_bucket_and_model_family_are_coarse() -> None:
assert [input_bucket(v) for v in (0, 2_000, 8_000, 32_000, 200_000)] == [
"xs",
"s",
"m",
"l",
"xl",
]
assert model_family("claude-sonnet-4-6") == "sonnet"
assert model_family("unknown-model") == "other"
def test_assign_arm_is_stable_and_respects_extreme_holdouts() -> None:
assert assign_arm("conv-123", 0.0) == "treatment"
assert assign_arm("conv-123", 1.0) == "control"
assert assign_arm("conv-123", 0.5) == assign_arm("conv-123", 0.5)
def test_conversation_key_uses_response_create_payload() -> None:
http_body = {"model": "gpt-5", "input": "build a cache"}
ws_body = {
"type": "response.create",
"response": {"model": "gpt-5", "input": "build a cache"},
}
assert conversation_key_from_body(http_body) == conversation_key_from_body(ws_body)
def test_stratum_label_round_trips_arm_and_key() -> None:
key = "opus|code|m|tools"
assert parse_stratum_label(stratum_label("treatment", key)) == ("treatment", key)
assert parse_stratum_label(stratum_label("control", key)) == ("control", key)
assert parse_stratum_label("unrelated") is None