mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
refactor(proxy): isolate output effort policy (#1961)
## Description Extracts provider-neutral output effort decisions into a pure `output_effort_policy` module. `output_shaper` still owns request mutation and labels, while the rank comparisons, legacy thinking clamp, and OpenAI text verbosity eligibility now live behind small deterministic functions. 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_effort_policy` for effort lowering, legacy thinking budget clamping, and OpenAI text verbosity decisions. - Updated `output_shaper` to delegate those pure decisions while preserving existing labels and request mutation behavior. - Added focused policy tests for effort rank transitions, thinking clamp boundaries, and verbosity creation/lowering. - Included the current LiteLLM callback signature compatibility shim required for repo-wide mypy on main-based slices. ## 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_effort_policy.py tests/test_output_shaper.py tests/test_litellm_callback.py -q 56 passed in 6.34s 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, clean worktree based on `headroomlabs/main`. - Exact command / steps: targeted pytest, ruff, format check, repo-wide mypy, staged gitleaks scan. - Observed result: output effort policy/shaper/callback tests pass; static checks pass; no staged secrets detected. - Not tested: live provider calls; this slice preserves existing request mutation behavior and only moves pure policy decisions. ## 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 architecture slice. PR-specific GHAS checks will be monitored after opening.
This commit is contained in:
parent
b1e871d51c
commit
094a53c047
3 changed files with 145 additions and 37 deletions
55
headroom/proxy/output_effort_policy.py
Normal file
55
headroom/proxy/output_effort_policy.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
"""Pure output-effort policy decisions.
|
||||
|
||||
The output shaper mutates provider request bodies. This module owns the
|
||||
provider-neutral decisions behind those mutations so rank comparisons and
|
||||
legacy budget clamping stay testable without request dictionaries.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
EFFORT_RANK = {"low": 0, "medium": 1, "high": 2, "xhigh": 3, "max": 4}
|
||||
TEXT_VERBOSITY_RANK = {"low": 0, "medium": 1, "high": 2}
|
||||
LEGACY_THINKING_FLOOR = 1024
|
||||
|
||||
|
||||
def lower_effort_value(current: object, target: str) -> str | None:
|
||||
"""Return ``target`` when an existing effort should be lowered."""
|
||||
if not isinstance(current, str):
|
||||
return None
|
||||
if current not in EFFORT_RANK or target not in EFFORT_RANK:
|
||||
return None
|
||||
if EFFORT_RANK[current] <= EFFORT_RANK[target]:
|
||||
return None
|
||||
return target
|
||||
|
||||
|
||||
def clamp_legacy_thinking_budget(
|
||||
*,
|
||||
thinking_type: object,
|
||||
budget_tokens: object,
|
||||
floor: int = LEGACY_THINKING_FLOOR,
|
||||
) -> int | None:
|
||||
"""Return the clamped budget for legacy enabled thinking, else ``None``."""
|
||||
if thinking_type != "enabled":
|
||||
return None
|
||||
if not isinstance(budget_tokens, int):
|
||||
return None
|
||||
if budget_tokens <= floor:
|
||||
return None
|
||||
return floor
|
||||
|
||||
|
||||
def can_create_openai_text_verbosity(model: object) -> bool:
|
||||
"""Whether it is safe to create a new OpenAI ``text.verbosity`` block."""
|
||||
return str(model or "").lower().startswith("gpt-5")
|
||||
|
||||
|
||||
def lower_text_verbosity_value(current: object) -> str | None:
|
||||
"""Return ``low`` when an existing OpenAI text verbosity should be lowered."""
|
||||
if not isinstance(current, str):
|
||||
return None
|
||||
if current not in TEXT_VERBOSITY_RANK:
|
||||
return None
|
||||
if TEXT_VERBOSITY_RANK[current] <= TEXT_VERBOSITY_RANK["low"]:
|
||||
return None
|
||||
return "low"
|
||||
|
|
@ -39,6 +39,16 @@ from enum import Enum
|
|||
from typing import Any
|
||||
|
||||
from headroom.proxy import runtime_env
|
||||
from headroom.proxy.output_effort_policy import (
|
||||
EFFORT_RANK as _EFFORT_RANK,
|
||||
)
|
||||
from headroom.proxy.output_effort_policy import (
|
||||
LEGACY_THINKING_FLOOR,
|
||||
can_create_openai_text_verbosity,
|
||||
clamp_legacy_thinking_budget,
|
||||
lower_effort_value,
|
||||
lower_text_verbosity_value,
|
||||
)
|
||||
from headroom.proxy.output_steering import (
|
||||
apply_openai_responses_verbosity_steering,
|
||||
apply_verbosity_steering,
|
||||
|
|
@ -66,15 +76,6 @@ __all__ = [
|
|||
"steering_text",
|
||||
]
|
||||
|
||||
# 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}
|
||||
|
||||
_TEXT_VERBOSITY_RANK = {"low": 0, "medium": 1, "high": 2}
|
||||
|
||||
_OPENAI_RESPONSES_OUTPUT_ITEM_TYPES = frozenset(
|
||||
{
|
||||
"custom_tool_call_output",
|
||||
|
|
@ -265,22 +266,24 @@ def route_effort(
|
|||
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}")
|
||||
lowered = lower_effort_value(effort, settings.mechanical_effort)
|
||||
if lowered is not None:
|
||||
output_config["effort"] = lowered
|
||||
labels.append(f"output_shaper:effort:{effort}->{lowered}")
|
||||
|
||||
# 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":
|
||||
if isinstance(thinking, dict):
|
||||
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}")
|
||||
clamped = clamp_legacy_thinking_budget(
|
||||
thinking_type=thinking.get("type"),
|
||||
budget_tokens=budget,
|
||||
floor=LEGACY_THINKING_FLOOR,
|
||||
)
|
||||
if clamped is not None:
|
||||
thinking["budget_tokens"] = clamped
|
||||
labels.append(f"output_shaper:thinking_budget:{budget}->{clamped}")
|
||||
|
||||
return labels
|
||||
|
||||
|
|
@ -363,22 +366,17 @@ def route_openai_reasoning_effort(
|
|||
return []
|
||||
effort = reasoning.get("effort")
|
||||
target = settings.mechanical_effort
|
||||
if (
|
||||
isinstance(effort, str)
|
||||
and effort in _EFFORT_RANK
|
||||
and target in _EFFORT_RANK
|
||||
and _EFFORT_RANK[effort] > _EFFORT_RANK[target]
|
||||
):
|
||||
reasoning["effort"] = target
|
||||
return [f"output_shaper:reasoning_effort:{effort}->{target}"]
|
||||
lowered = lower_effort_value(effort, target)
|
||||
if lowered is not None:
|
||||
reasoning["effort"] = lowered
|
||||
return [f"output_shaper:reasoning_effort:{effort}->{lowered}"]
|
||||
return []
|
||||
|
||||
|
||||
def route_openai_text_verbosity(body: dict[str, Any]) -> list[str]:
|
||||
"""Set or lower OpenAI ``text.verbosity`` conservatively."""
|
||||
model = str(body.get("model") or "").lower()
|
||||
text_config = body.get("text")
|
||||
can_create = model.startswith("gpt-5")
|
||||
can_create = can_create_openai_text_verbosity(body.get("model"))
|
||||
if text_config is None:
|
||||
if not can_create:
|
||||
return []
|
||||
|
|
@ -393,13 +391,10 @@ def route_openai_text_verbosity(body: dict[str, Any]) -> list[str]:
|
|||
return []
|
||||
text_config["verbosity"] = "low"
|
||||
return ["output_shaper:text_verbosity:unset->low"]
|
||||
if (
|
||||
isinstance(verbosity, str)
|
||||
and verbosity in _TEXT_VERBOSITY_RANK
|
||||
and _TEXT_VERBOSITY_RANK[verbosity] > _TEXT_VERBOSITY_RANK["low"]
|
||||
):
|
||||
text_config["verbosity"] = "low"
|
||||
return [f"output_shaper:text_verbosity:{verbosity}->low"]
|
||||
lowered = lower_text_verbosity_value(verbosity)
|
||||
if lowered is not None:
|
||||
text_config["verbosity"] = lowered
|
||||
return [f"output_shaper:text_verbosity:{verbosity}->{lowered}"]
|
||||
return []
|
||||
|
||||
|
||||
|
|
|
|||
58
tests/test_output_effort_policy.py
Normal file
58
tests/test_output_effort_policy.py
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
"""Tests for pure output effort policy decisions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from headroom.proxy.output_effort_policy import (
|
||||
LEGACY_THINKING_FLOOR,
|
||||
can_create_openai_text_verbosity,
|
||||
clamp_legacy_thinking_budget,
|
||||
lower_effort_value,
|
||||
lower_text_verbosity_value,
|
||||
)
|
||||
|
||||
|
||||
def test_lower_effort_value_lowers_known_higher_effort_to_target() -> None:
|
||||
assert lower_effort_value("xhigh", "low") == "low"
|
||||
assert lower_effort_value("max", "medium") == "medium"
|
||||
|
||||
|
||||
def test_lower_effort_value_keeps_lower_equal_unknown_or_non_string_values() -> None:
|
||||
assert lower_effort_value("low", "medium") is None
|
||||
assert lower_effort_value("medium", "medium") is None
|
||||
assert lower_effort_value("turbo", "low") is None
|
||||
assert lower_effort_value("high", "turbo") is None
|
||||
assert lower_effort_value(None, "low") is None
|
||||
|
||||
|
||||
def test_clamp_legacy_thinking_budget_only_clamps_enabled_over_floor() -> None:
|
||||
assert (
|
||||
clamp_legacy_thinking_budget(
|
||||
thinking_type="enabled",
|
||||
budget_tokens=32_000,
|
||||
)
|
||||
== LEGACY_THINKING_FLOOR
|
||||
)
|
||||
assert (
|
||||
clamp_legacy_thinking_budget(
|
||||
thinking_type="enabled",
|
||||
budget_tokens=LEGACY_THINKING_FLOOR,
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert clamp_legacy_thinking_budget(thinking_type="adaptive", budget_tokens=32_000) is None
|
||||
assert clamp_legacy_thinking_budget(thinking_type="enabled", budget_tokens="32000") is None
|
||||
|
||||
|
||||
def test_can_create_openai_text_verbosity_only_for_gpt5_family() -> None:
|
||||
assert can_create_openai_text_verbosity("gpt-5")
|
||||
assert can_create_openai_text_verbosity("GPT-5.1")
|
||||
assert not can_create_openai_text_verbosity("gpt-4o")
|
||||
assert not can_create_openai_text_verbosity(None)
|
||||
|
||||
|
||||
def test_lower_text_verbosity_value_lowers_existing_verbose_values() -> None:
|
||||
assert lower_text_verbosity_value("medium") == "low"
|
||||
assert lower_text_verbosity_value("high") == "low"
|
||||
assert lower_text_verbosity_value("low") is None
|
||||
assert lower_text_verbosity_value("chatty") is None
|
||||
assert lower_text_verbosity_value(None) is None
|
||||
Loading…
Add table
Add a link
Reference in a new issue