diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index 26a545f8d..fc5efcfde 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -73,6 +73,40 @@ class GeminiHandlerMixin: return True return False + def _rebuild_gemini_contents( + self, + original_contents: list[dict], + preserved_indices: set[int], + preserved_contents: dict[int, dict], + optimized_contents: list[dict], + ) -> list[dict]: + """Interleave preserved (non-text) entries back into optimized_contents at their + original positions. + + preserved_indices uses original contents[] indices, but optimized_contents uses + a different (shorter) index space because entries with no text parts were excluded + from the messages[] sent for compression. Using orig_idx directly to overwrite + optimized_contents[orig_idx] corrupts or silently drops entries. + + This method walks original_contents in order, placing each position with either + the preserved original (for non-text entries) or the next optimized text entry. + """ + opt_iter = iter(optimized_contents) + result: list[dict] = [] + for idx, content in enumerate(original_contents): + had_text = any("text" in p for p in content.get("parts", [])) + if idx in preserved_indices: + result.append(preserved_contents[idx]) + if had_text: + # Entry also produced a message; consume but discard the optimized version + next(opt_iter, None) + else: + opt_entry = next(opt_iter, None) + if opt_entry is not None: + result.append(opt_entry) + # else: dropped by compression — omit + return result + def _gemini_contents_to_messages( self, contents: list[dict], system_instruction: dict | None = None ) -> tuple[list[dict], set[int]]: @@ -501,12 +535,9 @@ class GeminiHandlerMixin: optimized_contents, optimized_system = self._messages_to_gemini_contents( 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 - + optimized_contents = self._rebuild_gemini_contents( + contents, preserved_indices, preserved_contents, optimized_contents + ) body["contents"] = optimized_contents if optimized_system: body["systemInstruction"] = optimized_system @@ -784,9 +815,12 @@ class GeminiHandlerMixin: optimized_contents, optimized_system = self._messages_to_gemini_contents( optimized_messages ) - for orig_idx, original_content in preserved_contents.items(): - if orig_idx < len(optimized_contents): - optimized_contents[orig_idx] = original_content + optimized_contents = self._rebuild_gemini_contents( + contents if isinstance(contents, list) else [], + preserved_indices, + preserved_contents, + optimized_contents, + ) request_payload["contents"] = optimized_contents if not is_antigravity: if optimized_system: @@ -1028,12 +1062,9 @@ class GeminiHandlerMixin: optimized_contents, optimized_system = self._messages_to_gemini_contents( 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 - + optimized_contents = self._rebuild_gemini_contents( + contents, preserved_indices, preserved_contents, optimized_contents + ) body["contents"] = optimized_contents if optimized_system: body["systemInstruction"] = optimized_system diff --git a/tests/test_google_multimodal.py b/tests/test_google_multimodal.py index c3f265885..03e583a8e 100644 --- a/tests/test_google_multimodal.py +++ b/tests/test_google_multimodal.py @@ -627,3 +627,80 @@ class TestParametrizedPreservation: """Parametrized test for preserved indices.""" _, preserved_indices = proxy._gemini_contents_to_messages(contents) assert preserved_indices == expected_indices + + +# ============================================================================= +# Tests for _rebuild_gemini_contents +# ============================================================================= + + +class TestRebuildGeminiContents: + """_rebuild_gemini_contents must re-insert preserved entries at their original positions.""" + + def _round_trip(self, proxy, contents): + """Simulate the full compression round-trip for a given contents list. + + Mimics what the handler does: convert → strip system msg → convert back → rebuild. + """ + messages, preserved_indices = proxy._gemini_contents_to_messages(contents) + preserved_contents = {idx: contents[idx] for idx in preserved_indices} + optimized_contents, _ = proxy._messages_to_gemini_contents(messages) + return proxy._rebuild_gemini_contents( + contents, preserved_indices, preserved_contents, optimized_contents + ) + + def test_text_only_unchanged(self, proxy): + """Text-only round-trip should produce identical contents.""" + contents = [TEXT_ONLY_CONTENT, MODEL_TEXT_CONTENT] + result = self._round_trip(proxy, contents) + assert len(result) == 2 + assert result[0]["parts"][0]["text"] == "Hello, world!" + assert result[1]["parts"][0]["text"] == "Hello! How can I help you today?" + + def test_function_call_sequence_preserved(self, proxy): + """functionCall and functionResponse entries must survive and appear at correct positions.""" + contents = [ + TEXT_ONLY_CONTENT, # idx 0: text + FUNCTION_CALL_CONTENT, # idx 1: functionCall only — no text → preserved + FUNCTION_RESPONSE_CONTENT, # idx 2: functionResponse only — no text → preserved + MODEL_TEXT_CONTENT, # idx 3: text + ] + result = self._round_trip(proxy, contents) + + assert len(result) == 4, f"Expected 4 entries, got {len(result)}: {result}" + # Position 0: original text + assert result[0]["parts"][0].get("text") == "Hello, world!" + # Position 1: functionCall preserved exactly + assert "functionCall" in result[1]["parts"][0], "functionCall missing at position 1" + assert result[1]["parts"][0]["functionCall"]["name"] == "get_weather" + # Position 2: functionResponse preserved exactly + assert "functionResponse" in result[2]["parts"][0], "functionResponse missing at position 2" + # Position 3: text preserved + assert result[3]["parts"][0].get("text") == "Hello! How can I help you today?" + + def test_function_call_at_start(self, proxy): + """Preserved entry at idx=0 must not overwrite idx=0 of optimized_contents.""" + contents = [ + FUNCTION_CALL_CONTENT, # idx 0: no text → preserved + TEXT_ONLY_CONTENT, # idx 1: text + ] + result = self._round_trip(proxy, contents) + + assert len(result) == 2 + assert "functionCall" in result[0]["parts"][0] + assert result[1]["parts"][0].get("text") == "Hello, world!" + + def test_hybrid_entry_uses_original(self, proxy): + """Entry with both text and functionCall keeps the original (with functionCall intact).""" + contents = [ + TEXT_ONLY_CONTENT, + FUNCTION_CALL_WITH_TEXT_CONTENT, # idx 1: has both text and functionCall → preserved + MODEL_TEXT_CONTENT, + ] + result = self._round_trip(proxy, contents) + + assert len(result) == 3 + # Hybrid entry must come back as the original (functionCall retained) + hybrid = result[1] + part_keys = {k for p in hybrid["parts"] for k in p} + assert "functionCall" in part_keys, "functionCall lost from hybrid entry"