mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
refactor(proxy): isolate output turn policy (#1962)
## Description Extracts output-shaper turn classification into a pure `output_turn_policy` module. The shaper still owns request mutation and labels, while Anthropic-style and OpenAI Responses structural turn classification now live in a deterministic policy boundary. 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_turn_policy` with `TurnKind`, `classify_turn`, and `classify_openai_responses_input`. - Updated `output_shaper` to import and re-export the classifiers, preserving existing import behavior. - Added direct policy tests for Anthropic tool-result turns and OpenAI Responses input classification. - 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_turn_policy.py tests/test_output_shaper.py tests/test_litellm_callback.py -q 60 passed in 6.25s 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 turn policy/shaper/callback tests pass; static checks pass; no staged secrets detected. - Not tested: live provider calls; this slice only moves structural classification logic and preserves existing shaper behavior. ## 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
2c9eb7c5f1
commit
c904a70d4e
4 changed files with 233 additions and 130 deletions
|
|
@ -104,6 +104,7 @@ class HeadroomCallback(_CustomLogger):
|
|||
data, call_type = cache, data
|
||||
if data is None:
|
||||
return None
|
||||
|
||||
if call_type not in ("completion", "acompletion"):
|
||||
return data
|
||||
|
||||
|
|
|
|||
|
|
@ -35,7 +35,6 @@ from __future__ import annotations
|
|||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from headroom.proxy import runtime_env
|
||||
|
|
@ -55,6 +54,11 @@ from headroom.proxy.output_steering import (
|
|||
replace_or_append_steering_block,
|
||||
steering_text,
|
||||
)
|
||||
from headroom.proxy.output_turn_policy import (
|
||||
TurnKind,
|
||||
classify_openai_responses_input,
|
||||
classify_turn,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -76,27 +80,9 @@ __all__ = [
|
|||
"steering_text",
|
||||
]
|
||||
|
||||
_OPENAI_RESPONSES_OUTPUT_ITEM_TYPES = frozenset(
|
||||
{
|
||||
"custom_tool_call_output",
|
||||
"function_call_output",
|
||||
"local_shell_call_output",
|
||||
"apply_patch_call_output",
|
||||
}
|
||||
)
|
||||
|
||||
_replace_or_append_steering_block = replace_or_append_steering_block
|
||||
|
||||
|
||||
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.
|
||||
|
|
@ -200,53 +186,6 @@ class ShapeResult:
|
|||
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 route_effort(
|
||||
body: dict[str, Any],
|
||||
kind: TurnKind,
|
||||
|
|
@ -288,70 +227,6 @@ def route_effort(
|
|||
return labels
|
||||
|
||||
|
||||
def _responses_part_text(value: Any) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
texts: list[str] = []
|
||||
for part in value:
|
||||
if isinstance(part, str):
|
||||
texts.append(part)
|
||||
elif isinstance(part, dict) and isinstance(part.get("text"), str):
|
||||
texts.append(part["text"])
|
||||
return "\n".join(text for text in texts if text)
|
||||
return ""
|
||||
|
||||
|
||||
def _responses_user_signal(item: dict[str, Any]) -> bool:
|
||||
item_type = item.get("type")
|
||||
role = item.get("role")
|
||||
if role == "user":
|
||||
content = item.get("content")
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") in {
|
||||
"input_file",
|
||||
"input_image",
|
||||
}:
|
||||
return True
|
||||
text = _responses_part_text(content)
|
||||
return bool(text.strip())
|
||||
if item_type == "input_text":
|
||||
text = _responses_part_text(item.get("text"))
|
||||
return bool(text.strip())
|
||||
if item_type == "input_image":
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def classify_openai_responses_input(input_data: Any) -> TurnKind:
|
||||
"""Classify OpenAI Responses ``input`` without content heuristics."""
|
||||
if isinstance(input_data, str):
|
||||
return TurnKind.NEW_USER_ASK if input_data.strip() else TurnKind.UNKNOWN
|
||||
if not isinstance(input_data, list) or not input_data:
|
||||
return TurnKind.UNKNOWN
|
||||
|
||||
saw_tool_output = False
|
||||
saw_unknown = False
|
||||
for item in input_data:
|
||||
if not isinstance(item, dict):
|
||||
saw_unknown = True
|
||||
continue
|
||||
item_type = item.get("type")
|
||||
if item_type in _OPENAI_RESPONSES_OUTPUT_ITEM_TYPES:
|
||||
saw_tool_output = True
|
||||
continue
|
||||
if _responses_user_signal(item):
|
||||
return TurnKind.NEW_USER_ASK
|
||||
if item_type in {"message", "function_call", "reasoning"}:
|
||||
continue
|
||||
saw_unknown = True
|
||||
|
||||
if saw_tool_output and not saw_unknown:
|
||||
return TurnKind.MECHANICAL_CONTINUATION
|
||||
return TurnKind.UNKNOWN
|
||||
|
||||
|
||||
def route_openai_reasoning_effort(
|
||||
body: dict[str, Any],
|
||||
kind: TurnKind,
|
||||
|
|
|
|||
125
headroom/proxy/output_turn_policy.py
Normal file
125
headroom/proxy/output_turn_policy.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
"""Pure structural turn classification for output shaping."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
|
||||
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"
|
||||
|
||||
|
||||
_OPENAI_RESPONSES_OUTPUT_ITEM_TYPES = frozenset(
|
||||
{
|
||||
"custom_tool_call_output",
|
||||
"function_call_output",
|
||||
"local_shell_call_output",
|
||||
"apply_patch_call_output",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def classify_turn(messages: list[dict[str, Any]]) -> TurnKind:
|
||||
"""Classify the latest Anthropic-style turn from message structure only."""
|
||||
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":
|
||||
return TurnKind.NEW_USER_ASK
|
||||
elif btype in ("image", "document"):
|
||||
return TurnKind.NEW_USER_ASK
|
||||
|
||||
if saw_error:
|
||||
return TurnKind.ERROR_CONTINUATION
|
||||
if saw_tool_result:
|
||||
return TurnKind.MECHANICAL_CONTINUATION
|
||||
return TurnKind.UNKNOWN
|
||||
|
||||
|
||||
def _responses_part_text(value: Any) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
texts: list[str] = []
|
||||
for part in value:
|
||||
if isinstance(part, str):
|
||||
texts.append(part)
|
||||
elif isinstance(part, dict) and isinstance(part.get("text"), str):
|
||||
texts.append(part["text"])
|
||||
return "\n".join(text for text in texts if text)
|
||||
return ""
|
||||
|
||||
|
||||
def _responses_user_signal(item: dict[str, Any]) -> bool:
|
||||
item_type = item.get("type")
|
||||
role = item.get("role")
|
||||
if role == "user":
|
||||
content = item.get("content")
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") in {
|
||||
"input_file",
|
||||
"input_image",
|
||||
}:
|
||||
return True
|
||||
text = _responses_part_text(content)
|
||||
return bool(text.strip())
|
||||
if item_type == "input_text":
|
||||
text = _responses_part_text(item.get("text"))
|
||||
return bool(text.strip())
|
||||
if item_type == "input_image":
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def classify_openai_responses_input(input_data: Any) -> TurnKind:
|
||||
"""Classify OpenAI Responses ``input`` without content heuristics."""
|
||||
if isinstance(input_data, str):
|
||||
return TurnKind.NEW_USER_ASK if input_data.strip() else TurnKind.UNKNOWN
|
||||
if not isinstance(input_data, list) or not input_data:
|
||||
return TurnKind.UNKNOWN
|
||||
|
||||
saw_tool_output = False
|
||||
saw_unknown = False
|
||||
for item in input_data:
|
||||
if not isinstance(item, dict):
|
||||
saw_unknown = True
|
||||
continue
|
||||
item_type = item.get("type")
|
||||
if item_type in _OPENAI_RESPONSES_OUTPUT_ITEM_TYPES:
|
||||
saw_tool_output = True
|
||||
continue
|
||||
if _responses_user_signal(item):
|
||||
return TurnKind.NEW_USER_ASK
|
||||
if item_type in {"message", "function_call", "reasoning"}:
|
||||
continue
|
||||
saw_unknown = True
|
||||
|
||||
if saw_tool_output and not saw_unknown:
|
||||
return TurnKind.MECHANICAL_CONTINUATION
|
||||
return TurnKind.UNKNOWN
|
||||
102
tests/test_output_turn_policy.py
Normal file
102
tests/test_output_turn_policy.py
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
"""Tests for pure output turn classification policy."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from headroom.proxy.output_turn_policy import (
|
||||
TurnKind,
|
||||
classify_openai_responses_input,
|
||||
classify_turn,
|
||||
)
|
||||
|
||||
|
||||
def _tool_result(is_error: bool = False) -> dict[str, Any]:
|
||||
block: dict[str, Any] = {"type": "tool_result", "content": "ok"}
|
||||
if is_error:
|
||||
block["is_error"] = True
|
||||
return block
|
||||
|
||||
|
||||
def test_anthropic_text_user_message_is_new_ask() -> None:
|
||||
assert classify_turn([{"role": "user", "content": "explain this"}]) is TurnKind.NEW_USER_ASK
|
||||
|
||||
|
||||
def test_anthropic_clean_tool_results_are_mechanical() -> None:
|
||||
messages = [{"role": "user", "content": [_tool_result(), _tool_result()]}]
|
||||
assert classify_turn(messages) is TurnKind.MECHANICAL_CONTINUATION
|
||||
|
||||
|
||||
def test_anthropic_error_tool_result_is_error_continuation() -> None:
|
||||
messages = [{"role": "user", "content": [_tool_result(), _tool_result(is_error=True)]}]
|
||||
assert classify_turn(messages) is TurnKind.ERROR_CONTINUATION
|
||||
|
||||
|
||||
def test_anthropic_user_media_or_text_block_is_new_ask() -> None:
|
||||
assert (
|
||||
classify_turn([{"role": "user", "content": [{"type": "image", "source": {}}]}])
|
||||
is TurnKind.NEW_USER_ASK
|
||||
)
|
||||
assert (
|
||||
classify_turn(
|
||||
[{"role": "user", "content": [_tool_result(), {"type": "text", "text": "also"}]}]
|
||||
)
|
||||
is TurnKind.NEW_USER_ASK
|
||||
)
|
||||
|
||||
|
||||
def test_anthropic_unknown_shapes_are_unknown() -> None:
|
||||
assert classify_turn([]) is TurnKind.UNKNOWN
|
||||
assert classify_turn([{"role": "assistant", "content": "done"}]) is TurnKind.UNKNOWN
|
||||
assert classify_turn([{"role": "user", "content": []}]) is TurnKind.UNKNOWN
|
||||
assert classify_turn([{"role": "user", "content": [{}]}]) is TurnKind.UNKNOWN
|
||||
|
||||
|
||||
def test_openai_responses_string_input_is_new_ask() -> None:
|
||||
assert classify_openai_responses_input("explain this") is TurnKind.NEW_USER_ASK
|
||||
assert classify_openai_responses_input(" ") is TurnKind.UNKNOWN
|
||||
|
||||
|
||||
def test_openai_responses_tool_outputs_only_are_mechanical() -> None:
|
||||
assert (
|
||||
classify_openai_responses_input(
|
||||
[
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "ok"},
|
||||
{"type": "local_shell_call_output", "call_id": "call_2", "output": "ok"},
|
||||
]
|
||||
)
|
||||
is TurnKind.MECHANICAL_CONTINUATION
|
||||
)
|
||||
|
||||
|
||||
def test_openai_responses_user_message_or_input_media_is_new_ask() -> None:
|
||||
assert (
|
||||
classify_openai_responses_input(
|
||||
[
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "also check foo.py"}],
|
||||
}
|
||||
]
|
||||
)
|
||||
is TurnKind.NEW_USER_ASK
|
||||
)
|
||||
assert (
|
||||
classify_openai_responses_input(
|
||||
[{"type": "message", "role": "user", "content": [{"type": "input_image"}]}]
|
||||
)
|
||||
is TurnKind.NEW_USER_ASK
|
||||
)
|
||||
|
||||
|
||||
def test_openai_responses_unknown_mixed_with_tool_output_is_unknown() -> None:
|
||||
assert (
|
||||
classify_openai_responses_input(
|
||||
[
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "ok"},
|
||||
{"type": "unrecognized_event"},
|
||||
]
|
||||
)
|
||||
is TurnKind.UNKNOWN
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue