mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy/gemini): preserve non-text content across the compression round-trip (#2079)
## Description Two related content-loss bugs in the Gemini `contents[]` <-> `messages[]` compression round-trip. Both drop or misplace real user content that entries with **non-text** parts should carry through untouched. They share the same theme (non-text preservation), so they're bundled here as two commits. ### 1. Google batch handler restores preserved entries by the wrong index (`handlers/batch.py`) The `batchGenerateContent` handler restored preserved (non-text) entries with the raw-index loop that commit #836 (`_rebuild_gemini_contents`) replaced in the three non-batch Gemini handlers: ```python for orig_idx, original_content in preserved_contents.items(): if orig_idx < len(optimized_contents): optimized_contents[orig_idx] = original_content ``` `preserved_indices` are indices into the **original** `contents[]`, but `optimized_contents` is a **shorter** list (text-less entries produce no message). Indexing `optimized_contents` by `orig_idx` overwrites the wrong entry and drops any preserved entry whose original index is past the optimized length. For: ```python [user text, model functionCall, user functionResponse, model text] ``` the batch was forwarded to Google as **two** entries: the model's answer overwritten by the functionCall, and the functionResponse dropped. Unlike `gemini.py` there is no `if optimized_messages != messages` gate, so it runs on every mixed batch item. **Fix:** use the shared `_rebuild_gemini_contents` interleaving helper. ### 2. Code-execution parts not detected as non-text (`handlers/gemini.py`) `_has_non_text_parts` only recognized `inlineData`/`fileData`/`functionCall`/`functionResponse`. Gemini's code-execution feature emits `executableCode` and `codeExecutionResult` parts, echoed back in `contents[]` on later turns. Because they weren't detected: - a mixed `text`+`executableCode` entry lost its code payload (only the text survived the round-trip); - a text-less `executableCode`+`codeExecutionResult` entry was treated as a phantom in `_rebuild_gemini_contents` — it consumed the next optimized message, dropping the whole code turn and shifting a following user turn into the model's role slot (corrupting role alternation). **Fix:** add both keys to the non-text detection so those entries are preserved verbatim. Closes: no issue filed — both found while auditing the Gemini contents<->messages round-trip. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/proxy/handlers/batch.py`: use `_rebuild_gemini_contents` instead of the raw-index restore loop. - `headroom/proxy/handlers/gemini.py`: recognize `executableCode` / `codeExecutionResult` in `_has_non_text_parts`. - `tests/test_proxy_handlers_batch.py`: add `test_handle_google_batch_create_preserves_functioncall_response_order`, driving the handler with the **real** Gemini converters (the existing batch tests stub them, which hid the bug); mix `GeminiHandlerMixin` into the shared `DummyBatchHandler` so `_rebuild_gemini_contents` is available. - `tests/test_google_multimodal.py`: extend the parametrized `test_each_non_text_key_detected` to the two new keys, and add `test_code_execution_entry_survives`. ## Testing - [x] New regression tests added (`tests/test_proxy_handlers_batch.py`, `tests/test_google_multimodal.py`) - [x] Linting/formatting clean — run with the CI-pinned `ruff==0.15.17` - [ ] Full `pytest` deferred to CI (local-OOM reason below). ```text $ uvx ruff@0.15.17 check headroom/proxy/handlers/batch.py headroom/proxy/handlers/gemini.py \ tests/test_proxy_handlers_batch.py tests/test_google_multimodal.py All checks passed! ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.10, headroom from this branch. Importing `headroom` pulls in the torch/transformers stack and a full `pytest` gets OOM-killed on this box, so I verified the interleaving/detection with dependency-free scripts (replicating the Gemini converters, the old loop, and `_rebuild_gemini_contents`) and left the full pytest to CI. - Exact command / steps: ran two standalone scripts. Script 1 rebuilds a Gemini batch request with `preserved_indices` holding a `functionCall`/`functionResponse` pair and compares the old raw-index loop against `_rebuild_gemini_contents`. Script 2 feeds a `codeExecutionResult` entry through `_has_non_text_parts` and the preserve path with and without the two new allowlist keys. Also ran `uvx ruff@0.15.17 check` on the changed files and tests. - Observed result: the old batch loop drops the `functionResponse` and overwrites the answer (4 parts collapse to 2); `_rebuild_gemini_contents` keeps all 4. Without the new keys the code-execution entry is dropped/shifted (2 parts, code absent); with them it survives intact (3 parts, code present). Lint clean. See the two blocks below. Batch fix (bug #1): ```text preserved_indices: [1, 2] OLD result parts: ['text', 'functionCall'] len 2 NEW result parts: ['text', 'functionCall', 'functionResponse', 'text'] len 4 GEMINI BATCH REBUILD FIX VERIFIED (old drops response + overwrites answer; new keeps all 4) ``` Code-execution fix (bug #2): ```text (b) OLD len=2 NEW len=3 (a) OLD has code=False NEW has code=True GEMINI CODE-EXECUTION PRESERVE FIX VERIFIED (old drops/shifts; new keeps intact) ``` - Not tested: a live Google/Gemini round-trip (handlers stubbed, as the existing tests do). Full local `pytest` deferred to CI (OOM, per above). ## 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 — ran lint + standalone logic checks; full pytest deferred to CI (local OOM, disclosed above) - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes - Two small behavioral changes (one loop -> shared helper, two keys added to an allowlist) plus regression tests; no new dependencies. Both complete/extend the non-text preservation the non-batch handlers already do (the #836 line). - @JerrettDavis tagging you since you reviewed the recent Gemini fixes. Both of these drop content (functionResponse/images on batch; code-execution on the normal round-trip), so they seemed worth surfacing together. Thanks. --------- Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
This commit is contained in:
parent
113894600c
commit
4056117d90
5 changed files with 153 additions and 7 deletions
|
|
@ -102,6 +102,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
|
||||
### Bug Fixes
|
||||
|
||||
* **proxy/batch:** stop corrupting Google `batchGenerateContent` requests whose contents interleave text turns with text-less entries (functionCall/functionResponse/images). The batch handler restored preserved (non-text) entries with the raw-index loop that #836 replaced everywhere else — indexing the shorter `optimized_contents` (text-less entries produce no message) by the original `contents[]` index, which overwrites the wrong entry and drops any preserved entry whose original index is past the optimized length. A request like `[user text, model functionCall, user functionResponse, model text]` was forwarded to Google as two entries: the model's answer overwritten by the functionCall and the functionResponse dropped. The batch handler now uses the shared `_rebuild_gemini_contents` interleaving helper, so all entries survive in order.
|
||||
* **proxy/gemini:** preserve Gemini code-execution parts (`executableCode` / `codeExecutionResult`) across the compression round-trip. `_has_non_text_parts` only recognized `inlineData`/`fileData`/`functionCall`/`functionResponse`, so a content entry carrying code-execution parts was not marked as preserved. A mixed `text`+`executableCode` entry lost its code payload (only the text survived), and a text-less code-execution entry was treated as a phantom that dropped the entire turn and shifted a neighboring message into the wrong role slot. Both keys are now recognized so those entries are preserved verbatim.
|
||||
* **cache/ccr:** don't evict a live entry when a duplicate hash is re-stored at capacity. `CompressionStore.store` ran `_evict_if_needed()` before checking whether the key already existed, so re-storing an already-present hash while the store was full evicted the oldest *distinct* entry to "make room" and then merely overwrote the existing key in place — no room was ever needed. The store dropped below `max_entries` and a live, never-retrieved entry was destroyed, so its `<<ccr:...>>` marker (still in the conversation) resolved to a 404. The CCR mirror bridge re-stores the same `explicit_hash` every turn a marker is re-encountered, so this fired routinely. Eviction now runs only for a genuinely new key.
|
||||
* **tokenizers:** recurse into a native `tool_result` whose content is a list of blocks instead of JSON-serializing it. `_count_content_parts` counted a `tool_result` with list content via `_count_serialized` (json.dumps + sample), so a base64 image nested in a tool result (computer-use / MCP screenshot tools) was priced as text — a ~50-200x overcount (a ~200KB screenshot read as ~70K tokens instead of ~1600). It now recurses into the nested blocks, matching the sibling Strands `toolResult` branch, so the image is priced structurally. The overcount made a single screenshot appear to blow past the model's context window and triggered unnecessary/over-aggressive compression.
|
||||
* **tokenizers:** price dense scripts (CJK/Kana/Hangul) in the fixed-ratio estimator path. `EstimatingTokenCounter.count_text` applied the CJK correction only on the auto-detect path; its fixed-ratio early return divided by the Latin ratio with no adjustment. The registry builds every provider-calibrated counter with a fixed ratio (Anthropic 3.5, Google 4.0, Cohere 4.0, Moonshot 3.1), and the Anthropic/Gemini proxy handlers count via `get_tokenizer(model).count_messages`, so a CJK-heavy context read as ~40-55% of its true token size — it could fall under the size/backpressure gates and skip compression, and every `x-headroom-tokens-before` metric for CJK traffic was materially wrong. The fixed-ratio path now applies the same dense-script split.
|
||||
|
|
|
|||
|
|
@ -225,10 +225,17 @@ class BatchHandlerMixin:
|
|||
optimized_messages
|
||||
)
|
||||
|
||||
# Restore preserved content entries that had non-text parts
|
||||
for orig_idx, original_content in preserved_contents.items():
|
||||
if orig_idx < len(optimized_contents):
|
||||
optimized_contents[orig_idx] = original_content
|
||||
# Restore preserved (non-text) entries at their ORIGINAL positions.
|
||||
# preserved_indices are indices into the original contents[], but
|
||||
# optimized_contents lives in a shorter index space (text-less
|
||||
# entries produced no message), so indexing it by orig_idx
|
||||
# overwrites the wrong entry and drops any preserved entry whose
|
||||
# original index is >= len(optimized_contents). Use the shared
|
||||
# interleaving helper the non-batch Gemini handlers already use
|
||||
# (#836).
|
||||
optimized_contents = self._rebuild_gemini_contents(
|
||||
contents, preserved_indices, preserved_contents, optimized_contents
|
||||
)
|
||||
|
||||
# Create compressed batch request
|
||||
compressed_req_content = {**req_content, "contents": optimized_contents}
|
||||
|
|
|
|||
|
|
@ -58,6 +58,8 @@ class GeminiHandlerMixin:
|
|||
- fileData: File references (URI + MIME type)
|
||||
- functionCall: Function calls from model
|
||||
- functionResponse: Responses to function calls
|
||||
- executableCode / codeExecutionResult: Gemini code-execution parts,
|
||||
echoed back in contents[] on later turns
|
||||
|
||||
Args:
|
||||
content: A single Gemini content entry with 'parts' list.
|
||||
|
|
@ -69,7 +71,14 @@ class GeminiHandlerMixin:
|
|||
for part in parts:
|
||||
if any(
|
||||
key in part
|
||||
for key in ("inlineData", "fileData", "functionCall", "functionResponse")
|
||||
for key in (
|
||||
"inlineData",
|
||||
"fileData",
|
||||
"functionCall",
|
||||
"functionResponse",
|
||||
"executableCode",
|
||||
"codeExecutionResult",
|
||||
)
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -173,7 +173,16 @@ class TestHasNonTextParts:
|
|||
|
||||
@pytest.mark.parametrize(
|
||||
"non_text_key",
|
||||
["inlineData", "fileData", "functionCall", "functionResponse"],
|
||||
[
|
||||
"inlineData",
|
||||
"fileData",
|
||||
"functionCall",
|
||||
"functionResponse",
|
||||
# Gemini code-execution parts, echoed back in contents[] on later
|
||||
# turns; previously not detected, so they were dropped on round-trip.
|
||||
"executableCode",
|
||||
"codeExecutionResult",
|
||||
],
|
||||
)
|
||||
def test_each_non_text_key_detected(self, proxy, non_text_key):
|
||||
"""Each non-text part type is correctly detected."""
|
||||
|
|
@ -657,6 +666,30 @@ class TestRebuildGeminiContents:
|
|||
assert result[0]["parts"][0]["text"] == "Hello, world!"
|
||||
assert result[1]["parts"][0]["text"] == "Hello! How can I help you today?"
|
||||
|
||||
def test_code_execution_entry_survives(self, proxy):
|
||||
"""A text-less code-execution entry (executableCode + codeExecutionResult)
|
||||
between two text turns must survive the round-trip at its position, and
|
||||
not shift a neighboring turn. Before the fix it was not detected as
|
||||
non-text, so it was dropped and the following user turn was misplaced."""
|
||||
code_entry = {
|
||||
"role": "model",
|
||||
"parts": [
|
||||
{"executableCode": {"language": "PYTHON", "code": "x = 1"}},
|
||||
{"codeExecutionResult": {"outcome": "OUTCOME_OK", "output": "1"}},
|
||||
],
|
||||
}
|
||||
contents = [
|
||||
{"role": "user", "parts": [{"text": "Question 1"}]},
|
||||
code_entry,
|
||||
{"role": "user", "parts": [{"text": "Question 2"}]},
|
||||
]
|
||||
|
||||
result = self._round_trip(proxy, contents)
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[1] == code_entry # preserved verbatim, in place
|
||||
assert result[2]["parts"][0]["text"] == "Question 2"
|
||||
|
||||
def test_function_call_sequence_preserved(self, proxy):
|
||||
"""functionCall and functionResponse entries must survive and appear at correct positions."""
|
||||
contents = [
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from types import SimpleNamespace
|
|||
import pytest
|
||||
|
||||
from headroom.proxy.handlers import batch as batch_module
|
||||
from headroom.proxy.handlers.gemini import GeminiHandlerMixin
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
|
|
@ -72,7 +73,10 @@ class FakeMetrics:
|
|||
self.failed_calls.append(kwargs)
|
||||
|
||||
|
||||
class DummyBatchHandler(batch_module.BatchHandlerMixin):
|
||||
class DummyBatchHandler(batch_module.BatchHandlerMixin, GeminiHandlerMixin):
|
||||
# GeminiHandlerMixin supplies the real _rebuild_gemini_contents (and the
|
||||
# other content helpers); the two converter methods below intentionally
|
||||
# override the mixin's for the stub-based tests.
|
||||
OPENAI_API_URL = "https://openai.example"
|
||||
GEMINI_API_URL = "https://gemini.example"
|
||||
|
||||
|
|
@ -939,6 +943,97 @@ async def test_handle_google_batch_create_covers_passthrough_revert_and_store_fa
|
|||
assert optimized["systemInstruction"] == {"parts": [{"text": "sys"}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_google_batch_create_preserves_functioncall_response_order(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A batch request that interleaves text turns with text-less
|
||||
functionCall/functionResponse entries must reach Google with all entries
|
||||
intact and in order. The old raw-index restore loop overwrote the model's
|
||||
answer with the functionCall and dropped the functionResponse."""
|
||||
|
||||
class RealConvHandler(batch_module.BatchHandlerMixin, GeminiHandlerMixin):
|
||||
# Real Gemini converters + _rebuild_gemini_contents (no stubs), so the
|
||||
# actual index interleaving runs.
|
||||
GEMINI_API_URL = "https://gemini.example"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.http_client = FakeHttpClient()
|
||||
self.metrics = FakeMetrics()
|
||||
self.config = SimpleNamespace(
|
||||
optimize=True, ccr_inject_tool=False, ccr_inject_system_instructions=False
|
||||
)
|
||||
self.openai_provider = SimpleNamespace(get_context_limit=lambda m: 8192)
|
||||
# No-op pipeline: return the messages unchanged, no token inflation.
|
||||
self.openai_pipeline = SimpleNamespace(
|
||||
apply=lambda **kw: SimpleNamespace(
|
||||
messages=kw["messages"], timing={}, tokens_before=100, tokens_after=100
|
||||
)
|
||||
)
|
||||
self.captured_body: dict | None = None
|
||||
|
||||
async def _next_request_id(self) -> str:
|
||||
return "req-1"
|
||||
|
||||
async def _record_request_outcome(self, outcome) -> None: # noqa: ANN001
|
||||
pass
|
||||
|
||||
def _extract_tags(self, headers: dict) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
async def _run_compression_in_executor(self, fn, *, timeout): # noqa: ANN001, ANN201
|
||||
return fn()
|
||||
|
||||
async def _store_google_batch_context(self, *a, **k) -> None: # noqa: ANN002, ANN003
|
||||
pass
|
||||
|
||||
async def _retry_request(self, method, url, headers, body, **kwargs): # noqa: ANN001, ANN201
|
||||
# Capture the (in-place mutated) forwarded batch body for assertions.
|
||||
self.captured_body = body
|
||||
return FakeResponse(status_code=200, content=b"{}", json_data={"name": "batches/1"})
|
||||
|
||||
handler = RealConvHandler()
|
||||
|
||||
contents = [
|
||||
{"role": "user", "parts": [{"text": "What's the weather in Paris?"}]},
|
||||
{
|
||||
"role": "model",
|
||||
"parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"functionResponse": {"name": "get_weather", "response": {"temp_c": 18}}}],
|
||||
},
|
||||
{"role": "model", "parts": [{"text": "It's 18C and cloudy in Paris."}]},
|
||||
]
|
||||
batch_body = {
|
||||
"batch": {
|
||||
"input_config": {
|
||||
"requests": {"requests": [{"request": {"contents": contents}, "metadata": {}}]}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async def payload(request): # noqa: ANN001, ANN201
|
||||
return batch_body
|
||||
|
||||
monkeypatch.setattr("headroom.proxy.helpers._read_request_json", payload)
|
||||
|
||||
resp = await handler.handle_google_batch_create(FakeRequest("{}"), "gemini-pro")
|
||||
assert resp.status_code == 200
|
||||
|
||||
out = handler.captured_body["batch"]["input_config"]["requests"]["requests"][0]["request"][
|
||||
"contents"
|
||||
]
|
||||
# All four entries survive in order. The old loop produced only two, dropping
|
||||
# the functionResponse and overwriting the model answer with the functionCall.
|
||||
assert len(out) == 4
|
||||
assert "text" in out[0]["parts"][0]
|
||||
assert out[1]["parts"][0].get("functionCall", {}).get("name") == "get_weather"
|
||||
assert out[2]["parts"][0].get("functionResponse", {}).get("name") == "get_weather"
|
||||
assert "Paris" in out[3]["parts"][0]["text"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_google_batch_passthrough_without_body_and_query_variants() -> None:
|
||||
handler = DummyBatchHandler()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue