diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index b3c1d7adb..6dd8e92fe 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -108,7 +108,11 @@ class GeminiHandlerMixin: return result def _gemini_contents_to_messages( - self, contents: list[dict], system_instruction: dict | None = None + self, + contents: list[dict], + system_instruction: dict | None = None, + *, + include_function_responses: bool = False, ) -> tuple[list[dict], set[int]]: """Convert Gemini contents[] format to OpenAI messages[] format for optimization. @@ -119,6 +123,12 @@ class GeminiHandlerMixin: OpenAI format: messages: [{"role": "user", "content": "..."}] + When include_function_responses is True, functionResponse payloads are + additionally emitted as ``role="tool"`` messages so waste-signal + detection can see tool output (#819). That richer list is telemetry-only: + entries with non-text parts stay in preserved_indices and are restored + verbatim, so it must never be used as the compression input. + Returns: Tuple of (messages, preserved_indices) where preserved_indices contains the indices of content entries that have non-text parts (images, function @@ -151,8 +161,29 @@ class GeminiHandlerMixin: if text_parts: messages.append({"role": role, "content": "\n".join(text_parts)}) + if include_function_responses: + for part in parts: + if "functionResponse" not in part: + continue + payload = self._function_response_text(part["functionResponse"]) + if payload: + messages.append({"role": "tool", "content": payload}) + return messages, preserved_indices + @staticmethod + def _function_response_text(function_response: dict) -> str: + """Serialize a functionResponse payload for waste-signal parsing.""" + response = function_response.get("response") + if response is None: + return "" + if isinstance(response, str): + return response + try: + return json.dumps(response, ensure_ascii=False, default=str) + except (TypeError, ValueError): + return str(response) + def _messages_to_gemini_contents(self, messages: list[dict]) -> tuple[list[dict], dict | None]: """Convert OpenAI messages[] format back to Gemini contents[] format. @@ -446,11 +477,17 @@ class GeminiHandlerMixin: try: # Use OpenAI pipeline (similar message format) context_limit = self.openai_provider.get_context_limit(model) + # Richer conversion incl. functionResponse payloads so tool + # output reaches waste-signal detection (#819); telemetry-only. + waste_messages, _ = self._gemini_contents_to_messages( + contents, system_instruction, include_function_responses=True + ) result = self.openai_pipeline.apply( messages=messages, model=model, model_limit=context_limit, context=extract_user_query(messages), + waste_messages=waste_messages, ) if result.messages != messages: optimized_messages = result.messages @@ -792,11 +829,17 @@ class GeminiHandlerMixin: if _decision.should_compress: try: context_limit = self.openai_provider.get_context_limit(model) + # Richer conversion incl. functionResponse payloads so tool + # output reaches waste-signal detection (#819); telemetry-only. + waste_messages, _ = self._gemini_contents_to_messages( + contents, system_instruction, include_function_responses=True + ) result = self.openai_pipeline.apply( messages=messages, model=model, model_limit=context_limit, context=extract_user_query(messages), + waste_messages=waste_messages, ) if result.messages != messages: optimized_messages = result.messages diff --git a/headroom/transforms/pipeline.py b/headroom/transforms/pipeline.py index e278ebe6b..973fc30bf 100644 --- a/headroom/transforms/pipeline.py +++ b/headroom/transforms/pipeline.py @@ -212,11 +212,14 @@ class TransformPipeline: - output_buffer: Output buffer override. - tool_profiles: Per-tool compression profiles. - request_id: Optional request ID for diff artifact. + - waste_messages: Optional richer conversion of the same request + used for waste-signal detection only (never transformed). Returns: Combined TransformResult. """ record_metrics = kwargs.pop("record_metrics", True) + waste_messages = kwargs.pop("waste_messages", None) tokenizer = self._get_tokenizer(model) provider_name = self._provider_name() @@ -430,13 +433,17 @@ class TransformPipeline: transforms=transform_diffs, ) - # Detect waste signals in original messages (only when significant compression) + # Detect waste signals in original messages (only when significant + # compression). Handlers whose wire format carries tool output the + # message conversion drops (e.g. Gemini functionResponse parts, #819) + # pass a richer waste_messages list that is parsed instead — it is + # telemetry-only and never transformed. waste_signals: WasteSignals | None = None if tokens_before > tokens_after and (tokens_before - tokens_after) > 100: try: from ..parser import parse_messages - _, _, waste_signals = parse_messages(messages, tokenizer) + _, _, waste_signals = parse_messages(waste_messages or messages, tokenizer) if waste_signals.total() == 0: waste_signals = None except Exception: diff --git a/tests/test_gemini_function_response_waste.py b/tests/test_gemini_function_response_waste.py new file mode 100644 index 000000000..638db4f8a --- /dev/null +++ b/tests/test_gemini_function_response_waste.py @@ -0,0 +1,203 @@ +"""Gemini functionResponse waste-signal visibility (issue #819). + +Gemini ``functionResponse`` parts are preserved verbatim on the wire (never +compressed), but their payloads previously never reached ``parse_messages``, +so tool output — where most waste lives — contributed nothing to waste +detection on the Gemini paths. + +The fix is telemetry-only: + +1. ``_gemini_contents_to_messages(..., include_function_responses=True)`` + additionally emits each functionResponse payload as a ``role="tool"`` + message. +2. ``TransformPipeline.apply(..., waste_messages=...)`` parses that richer + list for waste signals instead of the transform input. The transform path + and token accounting are untouched. +""" + +from __future__ import annotations + +import json + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from headroom import OpenAIProvider, Tokenizer +from headroom.config import HeadroomConfig +from headroom.parser import parse_messages +from headroom.proxy.server import HeadroomProxy, ProxyConfig +from headroom.transforms.pipeline import TransformPipeline + +_provider = OpenAIProvider() + + +@pytest.fixture +def proxy() -> HeadroomProxy: + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + return HeadroomProxy(config) + + +@pytest.fixture +def tokenizer() -> Tokenizer: + return Tokenizer(_provider.get_token_counter("gpt-4o"), "gpt-4o") + + +def _big_payload(rows: int = 200) -> dict: + return { + "result": [ + {"id": i, "name": f"item_{i}", "status": "ok", "score": i * 3.14} for i in range(rows) + ] + } + + +def _function_response_content(payload: object, name: str = "fetch_data") -> dict: + return { + "role": "user", + "parts": [{"functionResponse": {"name": name, "response": payload}}], + } + + +class TestFunctionResponseConversion: + def test_default_conversion_emits_no_tool_messages(self, proxy): + contents = [ + {"role": "user", "parts": [{"text": "fetch the data"}]}, + _function_response_content(_big_payload()), + ] + messages, preserved = proxy._gemini_contents_to_messages(contents) + assert [m["role"] for m in messages] == ["user"] + assert preserved == {1} + + def test_flag_emits_tool_message_for_dict_response(self, proxy): + payload = _big_payload() + contents = [ + {"role": "user", "parts": [{"text": "fetch the data"}]}, + _function_response_content(payload), + ] + messages, preserved = proxy._gemini_contents_to_messages( + contents, include_function_responses=True + ) + assert [m["role"] for m in messages] == ["user", "tool"] + assert json.loads(messages[1]["content"]) == payload + # preserved_indices semantics unchanged: the entry is still restored + # verbatim on the wire regardless of the telemetry conversion. + assert preserved == {1} + + def test_flag_passes_string_response_through(self, proxy): + contents = [_function_response_content("plain text tool output")] + messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True) + assert messages == [{"role": "tool", "content": "plain text tool output"}] + + def test_flag_skips_missing_response(self, proxy): + contents = [ + {"role": "user", "parts": [{"functionResponse": {"name": "noop"}}]}, + {"role": "user", "parts": [{"functionResponse": {"name": "none", "response": None}}]}, + ] + messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True) + assert messages == [] + + def test_flag_emits_text_before_tool_within_entry(self, proxy): + contents = [ + { + "role": "user", + "parts": [ + {"text": "tool said:"}, + {"functionResponse": {"name": "f", "response": "output"}}, + ], + } + ] + messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True) + assert [m["role"] for m in messages] == ["user", "tool"] + assert messages[0]["content"] == "tool said:" + assert messages[1]["content"] == "output" + + def test_unserializable_response_falls_back_to_str(self, proxy): + circular: dict = {"name": "loop"} + circular["self"] = circular + text = proxy._function_response_text({"response": circular}) + assert "loop" in text + + +class TestFunctionResponseWasteParsing: + def test_function_response_payload_reaches_waste_signals(self, proxy, tokenizer): + contents = [ + {"role": "user", "parts": [{"text": "fetch the data"}]}, + _function_response_content(_big_payload()), + ] + messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True) + blocks, _, waste = parse_messages(messages, tokenizer) + assert any(b.kind == "tool_result" for b in blocks) + assert waste.json_bloat_tokens > 0 + + def test_repeated_function_response_counts_as_reread(self, proxy, tokenizer): + payload = _big_payload() + filler = [{"role": "user", "parts": [{"text": f"working on step {i}"}]} for i in range(5)] + contents = [ + _function_response_content(payload), + *filler, + _function_response_content(payload), + ] + messages, _ = proxy._gemini_contents_to_messages(contents, include_function_responses=True) + _, _, waste = parse_messages(messages, tokenizer) + assert waste.reread_tokens > 0 + + +class TestPipelineWasteMessages: + @staticmethod + def _base_messages() -> list[dict]: + # Compressible enough that the pipeline clears the >100 saved-token + # gate that guards waste-signal detection. + return [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Inspect the data set."}, + {"role": "tool", "content": json.dumps(_big_payload(400)["result"])}, + ] + + def test_waste_messages_override_waste_source(self, tokenizer): + messages = self._base_messages() + extra_tool = {"role": "tool", "content": json.dumps(_big_payload(300))} + + baseline = TransformPipeline(HeadroomConfig()).apply( + [dict(m) for m in messages], model="gpt-4o", model_limit=128000 + ) + enriched = TransformPipeline(HeadroomConfig()).apply( + [dict(m) for m in messages], + model="gpt-4o", + model_limit=128000, + waste_messages=[*messages, extra_tool], + ) + + assert baseline.waste_signals is not None + assert enriched.waste_signals is not None + assert enriched.waste_signals.json_bloat_tokens > baseline.waste_signals.json_bloat_tokens + + def test_waste_messages_do_not_affect_transform_output(self, tokenizer): + messages = self._base_messages() + extra_tool = {"role": "tool", "content": json.dumps(_big_payload(300))} + + baseline = TransformPipeline(HeadroomConfig()).apply( + [dict(m) for m in messages], model="gpt-4o", model_limit=128000 + ) + enriched = TransformPipeline(HeadroomConfig()).apply( + [dict(m) for m in messages], + model="gpt-4o", + model_limit=128000, + waste_messages=[*messages, extra_tool], + ) + + assert enriched.messages == baseline.messages + assert enriched.tokens_before == baseline.tokens_before + assert enriched.tokens_after == baseline.tokens_after + + def test_no_waste_messages_falls_back_to_transform_input(self, tokenizer): + result = TransformPipeline(HeadroomConfig()).apply( + [dict(m) for m in self._base_messages()], model="gpt-4o", model_limit=128000 + ) + assert result.waste_signals is not None + assert result.waste_signals.json_bloat_tokens > 0