From c29b4ba84f021c62c869a12e14ce8671e2f01141 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Sat, 11 Jul 2026 04:42:11 +0000 Subject: [PATCH] refactor(output): isolate savings policy (#1947) ## 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. --- headroom/proxy/output_savings.py | 185 +++--------------------- headroom/proxy/output_savings_policy.py | 156 ++++++++++++++++++++ tests/test_output_savings_policy.py | 60 ++++++++ 3 files changed, 238 insertions(+), 163 deletions(-) create mode 100644 headroom/proxy/output_savings_policy.py create mode 100644 tests/test_output_savings_policy.py diff --git a/headroom/proxy/output_savings.py b/headroom/proxy/output_savings.py index ebe5ba44a..1fa1689bb 100644 --- a/headroom/proxy/output_savings.py +++ b/headroom/proxy/output_savings.py @@ -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. diff --git a/headroom/proxy/output_savings_policy.py b/headroom/proxy/output_savings_policy.py new file mode 100644 index 000000000..7ed5e3d98 --- /dev/null +++ b/headroom/proxy/output_savings_policy.py @@ -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 diff --git a/tests/test_output_savings_policy.py b/tests/test_output_savings_policy.py new file mode 100644 index 000000000..bb71543ba --- /dev/null +++ b/tests/test_output_savings_policy.py @@ -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