diff --git a/headroom/integrations/litellm_callback.py b/headroom/integrations/litellm_callback.py index 1cf152c59..7256d85b8 100644 --- a/headroom/integrations/litellm_callback.py +++ b/headroom/integrations/litellm_callback.py @@ -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 diff --git a/headroom/proxy/output_shaper.py b/headroom/proxy/output_shaper.py index cad23783e..e37d72432 100644 --- a/headroom/proxy/output_shaper.py +++ b/headroom/proxy/output_shaper.py @@ -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, diff --git a/headroom/proxy/output_turn_policy.py b/headroom/proxy/output_turn_policy.py new file mode 100644 index 000000000..045d9ec4f --- /dev/null +++ b/headroom/proxy/output_turn_policy.py @@ -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 diff --git a/tests/test_output_turn_policy.py b/tests/test_output_turn_policy.py new file mode 100644 index 000000000..4ba58661d --- /dev/null +++ b/tests/test_output_turn_policy.py @@ -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 + )