mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy/openai): apply output shaping on /v1/chat/completions (#2328)
## Description Fixes #2302. Output shaping (`HEADROOM_OUTPUT_SHAPER=1`) verbosity steering is wired into the Anthropic `/v1/messages` handler and the OpenAI `/v1/responses` handler, but never into `handle_openai_chat`. OpenAI-compatible clients that route through `/v1/chat/completions` — GitHub Copilot CLI, opencode, older SDKs — therefore got zero output savings, and `headroom output-savings` reported: ``` No shaped requests recorded yet. ``` `handle_openai_chat` referenced verbosity only for cache-key construction, never for actual shaping. The shared helpers (`OutputShaperSettings`, `resolve_verbosity_level`, `assign_arm`, `classify_turn`) existed but were not called from the chat path. ## Fix Run the same shaping block the Anthropic handler already uses, at the end of `handle_openai_chat` (after every other body mutation, before the upstream forward, skipped under `x-headroom-bypass`): - conversation-stable holdout via `assign_arm(conversation_key_from_body(body), holdout)` — `conversation_key_from_body` already reads `messages`, so it works unchanged for a chat body; - stratum labelling on the transforms channel so the outcome funnel feeds the output-savings ledger from the chat path; - for the treatment arm, verbosity steering via a new `shape_openai_chat_request`. The one genuinely new piece is a chat-specific steering injector. Anthropic carries the system prompt in a top-level `system` field and Responses in `instructions`; **chat/completions carries it as a `role: "system"` message inside `messages`**, which neither existing injector touches. `apply_openai_chat_verbosity_steering`: - appends the byte-stable steering block to the tail of the last `system`/`developer` message (idempotent via the `<headroom_output_shaping>` sentinel, and it swaps cleanly when the level changes); - handles both string content and the content-part list form (`[{"type": "text", ...}]`); - inserts a `role: "system"` message at the front only when the request has no system message. Because a whole conversation is stably treatment or control and the block text is fixed per level, a treatment conversation's steering is byte-stable across turns, so the provider prefix cache is not thrashed. Effort routing is intentionally not applied on this path — `route_effort` writes Anthropic-shaped `output_config`/thinking config with no portable chat/completions equivalent — so only the token-reducing verbosity lever runs. Mutating `body` in place is enough on this path; the outbound request serializes `body` fresh, so no body-mutation tracker is needed. ## Type of Change - [x] 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 - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/proxy/output_steering.py`: add `apply_openai_chat_verbosity_steering` (inject the steering block into the chat `messages` system prompt). - `headroom/proxy/output_shaper.py`: add `shape_openai_chat_request` (verbosity-only chat shaper) and export both new names. - `headroom/proxy/handlers/openai.py`: run the holdout/stratum + shaping block at the end of `handle_openai_chat`, mirroring the Anthropic handler and respecting bypass. - `tests/test_output_steering.py`: cover the injector (append, idempotency, level swap, insert-when-absent, list content, level-0 no-op). - `tests/test_output_shaper.py`: cover `shape_openai_chat_request` (disabled no-op, applies steering, level override, stable second pass). - `CHANGELOG.md`: Bug Fixes entry. ## Testing - [ ] 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 $ uvx ruff@0.15.17 check headroom/proxy/output_steering.py headroom/proxy/output_shaper.py headroom/proxy/handlers/openai.py tests/test_output_steering.py tests/test_output_shaper.py All checks passed! $ uvx ruff@0.15.17 format --check <same files> 5 files already formatted $ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/output_steering.py headroom/proxy/output_shaper.py Success: no issues found in 2 source files ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx mypy@1.20.2`. A full `pytest` OOMs this box (ML-stack import), so I reproduced the injector with a dependency-free script and left the full pytest to CI. - Exact command / steps: replicated `apply_openai_chat_verbosity_steering` (and the `steering_text`/`replace_or_append_steering_block` primitives it uses) and exercised: an existing string system message, an existing content-part list, no system message, re-apply at the same level, and a level swap. - Observed result: the steering block is appended to the system message while user turns and message order are untouched; re-applying at the same level is a no-op; a level change replaces the block (exactly one remains); a request with no system message gets one inserted at the front; level 0 is a no-op. The added unit tests assert the same through `shape_openai_chat_request`. - Not tested: a live Copilot CLI `/v1/chat/completions` round trip; the added tests drive the pure shaper/injector directly, matching the existing `test_output_shaper.py` / `test_output_steering.py` patterns. ## 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 - [ ] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes The "unit tests pass locally" box is unchecked because a local pytest run imports the ML stack and OOMs this box; the added tests are pure (no ML imports) and run under the normal CI pytest job, and the injector behavior is corroborated by the standalone proof above. Effort routing on chat/completions is deliberately out of scope here (no portable equivalent to the Anthropic effort levers); this PR restores the verbosity-steering savings the issue reports as missing, and effort routing for chat can follow separately if wanted. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
a63d235e7e
commit
1b8c11ebfb
5 changed files with 253 additions and 0 deletions
|
|
@ -3361,6 +3361,63 @@ class OpenAIHandlerMixin:
|
|||
# `max_completion_tokens`.
|
||||
_normalize_openai_max_tokens(body)
|
||||
|
||||
# Output shaping (opt-in via HEADROOM_OUTPUT_SHAPER): verbosity steering
|
||||
# on the chat system message. Runs after every other body mutation so the
|
||||
# turn classifier sees the final messages, and respects the same bypass
|
||||
# as compression. OpenAI-compatible clients that route through
|
||||
# /v1/chat/completions (GitHub Copilot CLI, opencode, older SDKs) never
|
||||
# reached the shaper before, so they saw zero output savings (#2302).
|
||||
# Mutating `body` in place is sufficient here — the outbound request
|
||||
# serializes `body` fresh, so no body-mutation tracker is needed.
|
||||
if not _bypass:
|
||||
from headroom.proxy import runtime_env
|
||||
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_openai_chat_request,
|
||||
)
|
||||
|
||||
_shaper_settings = OutputShaperSettings.from_env()
|
||||
if _shaper_settings.enabled:
|
||||
# Conversation-stable holdout: a whole conversation is treatment
|
||||
# or control, which keeps the A/B comparison clean and the
|
||||
# provider prefix cache stable (the steering block never flips
|
||||
# mid-conversation).
|
||||
_holdout = 0.0
|
||||
try:
|
||||
_holdout = float(runtime_env.getenv("HEADROOM_OUTPUT_HOLDOUT", "0") or "0")
|
||||
except ValueError:
|
||||
_holdout = 0.0
|
||||
_arm = assign_arm(conversation_key_from_body(body), _holdout)
|
||||
_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 transforms channel so the outcome
|
||||
# funnel feeds the output-savings ledger from the chat path too.
|
||||
transforms_applied.append(stratum_label(_arm, _stratum))
|
||||
if _arm == "treatment":
|
||||
_level, _src = resolve_verbosity_level(_shaper_settings)
|
||||
_shape_result = shape_openai_chat_request(
|
||||
body, _shaper_settings, level_override=_level
|
||||
)
|
||||
if _shape_result.changed:
|
||||
transforms_applied.extend(_shape_result.labels or [])
|
||||
logger.info(
|
||||
f"[{request_id}] OutputShaper(chat, L{_level}/{_src}): "
|
||||
f"{_shape_result.labels}"
|
||||
)
|
||||
|
||||
# Route through LiteLLM/any-llm backend if configured
|
||||
if self.anthropic_backend is not None:
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ from headroom.proxy.output_effort_policy import (
|
|||
lower_text_verbosity_value,
|
||||
)
|
||||
from headroom.proxy.output_steering import (
|
||||
apply_openai_chat_verbosity_steering,
|
||||
apply_openai_responses_verbosity_steering,
|
||||
apply_verbosity_steering,
|
||||
replace_or_append_steering_block,
|
||||
|
|
@ -76,6 +77,7 @@ __all__ = [
|
|||
"OutputShaperSettings",
|
||||
"ShapeResult",
|
||||
"TurnKind",
|
||||
"apply_openai_chat_verbosity_steering",
|
||||
"apply_openai_responses_verbosity_steering",
|
||||
"apply_verbosity_steering",
|
||||
"classify_openai_responses_input",
|
||||
|
|
@ -84,6 +86,7 @@ __all__ = [
|
|||
"route_effort",
|
||||
"route_openai_reasoning_effort",
|
||||
"route_openai_text_verbosity",
|
||||
"shape_openai_chat_request",
|
||||
"shape_openai_responses_request",
|
||||
"shape_request",
|
||||
"steering_text",
|
||||
|
|
@ -352,6 +355,36 @@ def shape_request(
|
|||
return result
|
||||
|
||||
|
||||
def shape_openai_chat_request(
|
||||
body: dict[str, Any],
|
||||
settings: OutputShaperSettings | None = None,
|
||||
level_override: int | None = None,
|
||||
) -> ShapeResult:
|
||||
"""Apply output-shaping levers to an OpenAI chat/completions body in place.
|
||||
|
||||
The chat counterpart of :func:`shape_request`. Chat carries the system
|
||||
prompt as a ``role: "system"`` message, so verbosity steering uses the
|
||||
chat-specific injector. Effort routing is intentionally not applied here:
|
||||
the ``route_effort`` levers write Anthropic-shaped config and there is no
|
||||
portable chat/completions equivalent, so only the verbosity steering lever
|
||||
(the one that reduces output tokens) runs on this path.
|
||||
"""
|
||||
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_openai_chat_verbosity_steering(body, level):
|
||||
result.changed = True
|
||||
result.labels.append(f"output_shaper:verbosity:L{level}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OpenAI Responses format (Codex, /v1/responses HTTP + WebSocket)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -46,6 +46,68 @@ def apply_verbosity_steering(body: dict[str, Any], level: int) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def apply_openai_chat_verbosity_steering(
|
||||
body: dict[str, Any],
|
||||
level: int,
|
||||
) -> bool:
|
||||
"""Append or replace the steering block in an OpenAI chat/completions body.
|
||||
|
||||
OpenAI ``/v1/chat/completions`` carries the system prompt as a
|
||||
``role: "system"`` (or ``"developer"``) message inside ``messages`` rather
|
||||
than a top-level field, so it needs its own injector (the Anthropic
|
||||
``system`` and Responses ``instructions`` variants do not reach it — the
|
||||
root cause of GitHub Copilot CLI seeing zero output savings, #2302).
|
||||
|
||||
The block is appended to the tail of the last system/developer message so a
|
||||
treatment conversation's steering stays byte-stable across turns (and
|
||||
re-applies idempotently via the sentinel). When the request carries no
|
||||
system message at all, one is inserted at the front. Returns True only when
|
||||
the body actually changed.
|
||||
"""
|
||||
text = steering_text(level)
|
||||
if text is None:
|
||||
return False
|
||||
|
||||
messages = body.get("messages")
|
||||
if not isinstance(messages, list):
|
||||
return False
|
||||
|
||||
target: dict[str, Any] | None = None
|
||||
for message in messages:
|
||||
if isinstance(message, dict) and message.get("role") in ("system", "developer"):
|
||||
target = message
|
||||
if target is None:
|
||||
# No system prompt to append to — insert one carrying just the block.
|
||||
messages.insert(0, {"role": "system", "content": text})
|
||||
return True
|
||||
|
||||
content = target.get("content")
|
||||
if content is None:
|
||||
target["content"] = text
|
||||
return True
|
||||
if isinstance(content, str):
|
||||
updated, changed = replace_or_append_steering_block(content, text)
|
||||
if changed:
|
||||
target["content"] = updated
|
||||
return changed
|
||||
if isinstance(content, list):
|
||||
# OpenAI also accepts a content-part list ([{"type": "text", ...}]).
|
||||
for part in content:
|
||||
if (
|
||||
isinstance(part, dict)
|
||||
and part.get("type") == "text"
|
||||
and isinstance(part.get("text"), str)
|
||||
and part["text"].startswith(_STEERING_SENTINEL)
|
||||
):
|
||||
if part["text"] == text:
|
||||
return False
|
||||
part["text"] = text
|
||||
return True
|
||||
content.append({"type": "text", "text": text})
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def apply_openai_responses_verbosity_steering(
|
||||
body: dict[str, Any],
|
||||
level: int,
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from headroom.proxy.output_shaper import (
|
|||
route_effort,
|
||||
route_openai_reasoning_effort,
|
||||
route_openai_text_verbosity,
|
||||
shape_openai_chat_request,
|
||||
shape_openai_responses_request,
|
||||
shape_request,
|
||||
steering_text,
|
||||
|
|
@ -402,3 +403,40 @@ class TestOpenAIResponsesTextVerbosity:
|
|||
assert steering_text(2) in body["instructions"]
|
||||
assert body["reasoning"]["effort"] == "low"
|
||||
assert body["text"]["verbosity"] == "low"
|
||||
|
||||
|
||||
class TestShapeOpenAIChatRequest:
|
||||
def test_disabled_is_noop(self):
|
||||
body = {"messages": [{"role": "system", "content": "Sys."}]}
|
||||
snapshot = copy.deepcopy(body)
|
||||
result = shape_openai_chat_request(body, OutputShaperSettings(enabled=False))
|
||||
assert result.changed is False
|
||||
assert body == snapshot
|
||||
|
||||
def test_enabled_applies_verbosity_steering(self):
|
||||
body = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "Sys."},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
}
|
||||
result = shape_openai_chat_request(body, ENABLED)
|
||||
assert result.changed is True
|
||||
assert result.labels == ["output_shaper:verbosity:L2"]
|
||||
assert steering_text(2) in body["messages"][0]["content"]
|
||||
# User turn is untouched.
|
||||
assert body["messages"][1] == {"role": "user", "content": "hi"}
|
||||
|
||||
def test_level_override_supersedes_settings(self):
|
||||
body = {"messages": [{"role": "system", "content": "Sys."}]}
|
||||
result = shape_openai_chat_request(body, ENABLED, level_override=4)
|
||||
assert result.labels == ["output_shaper:verbosity:L4"]
|
||||
assert steering_text(4) in body["messages"][0]["content"]
|
||||
|
||||
def test_second_pass_is_stable(self):
|
||||
body = {"messages": [{"role": "system", "content": "Sys."}]}
|
||||
shape_openai_chat_request(body, ENABLED)
|
||||
snapshot = copy.deepcopy(body)
|
||||
second = shape_openai_chat_request(body, ENABLED)
|
||||
assert second.changed is False
|
||||
assert body == snapshot
|
||||
|
|
|
|||
|
|
@ -42,3 +42,66 @@ def test_openai_responses_steering_is_idempotent() -> None:
|
|||
snapshot = body.copy()
|
||||
assert apply_openai_responses_verbosity_steering(body, 2) is False
|
||||
assert body == snapshot
|
||||
|
||||
|
||||
def test_openai_chat_steering_appends_to_system_message() -> None:
|
||||
from headroom.proxy.output_steering import apply_openai_chat_verbosity_steering
|
||||
|
||||
body = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
}
|
||||
assert apply_openai_chat_verbosity_steering(body, 2) is True
|
||||
sys_content = body["messages"][0]["content"]
|
||||
assert "You are helpful." in sys_content
|
||||
assert steering_text(2) in sys_content
|
||||
# Other messages and ordering are untouched.
|
||||
assert body["messages"][1] == {"role": "user", "content": "hi"}
|
||||
assert [m["role"] for m in body["messages"]] == ["system", "user"]
|
||||
|
||||
|
||||
def test_openai_chat_steering_is_idempotent_and_swaps_level() -> None:
|
||||
from headroom.proxy.output_steering import apply_openai_chat_verbosity_steering
|
||||
|
||||
body = {"messages": [{"role": "system", "content": "S."}]}
|
||||
assert apply_openai_chat_verbosity_steering(body, 2) is True
|
||||
first = body["messages"][0]["content"]
|
||||
# Same level again: no change.
|
||||
assert apply_openai_chat_verbosity_steering(body, 2) is False
|
||||
assert body["messages"][0]["content"] == first
|
||||
# Different level: replace, still exactly one block.
|
||||
assert apply_openai_chat_verbosity_steering(body, 4) is True
|
||||
swapped = body["messages"][0]["content"]
|
||||
assert steering_text(4) in swapped
|
||||
assert swapped.count("<headroom_output_shaping>") == 1
|
||||
|
||||
|
||||
def test_openai_chat_steering_inserts_system_when_absent() -> None:
|
||||
from headroom.proxy.output_steering import apply_openai_chat_verbosity_steering
|
||||
|
||||
body = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
assert apply_openai_chat_verbosity_steering(body, 3) is True
|
||||
assert body["messages"][0]["role"] == "system"
|
||||
assert body["messages"][0]["content"] == steering_text(3)
|
||||
assert body["messages"][1] == {"role": "user", "content": "hi"}
|
||||
|
||||
|
||||
def test_openai_chat_steering_handles_list_content() -> None:
|
||||
from headroom.proxy.output_steering import apply_openai_chat_verbosity_steering
|
||||
|
||||
body = {"messages": [{"role": "system", "content": [{"type": "text", "text": "base"}]}]}
|
||||
assert apply_openai_chat_verbosity_steering(body, 1) is True
|
||||
parts = body["messages"][0]["content"]
|
||||
assert parts[0] == {"type": "text", "text": "base"}
|
||||
assert parts[1]["type"] == "text"
|
||||
assert parts[1]["text"] == steering_text(1)
|
||||
|
||||
|
||||
def test_openai_chat_steering_level_zero_is_noop() -> None:
|
||||
from headroom.proxy.output_steering import apply_openai_chat_verbosity_steering
|
||||
|
||||
body = {"messages": [{"role": "system", "content": "S."}]}
|
||||
assert apply_openai_chat_verbosity_steering(body, 0) is False
|
||||
assert body["messages"][0]["content"] == "S."
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue