From 393d252eef286a2def644f2c8f0a3e03e0b24f68 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 17 Jun 2026 03:31:28 +0000 Subject: [PATCH] fix(google_genai): preserve raw bytes in Gemini streamGenerateContent SSE framing The Gemini generate-content streaming iterators framed events with httpx aiter_lines/iter_lines, which split on str.splitlines boundaries. That set includes U+2028, U+2029, U+0085 and form feed, characters Gemini emits raw inside data: JSON for thinking and text content. A single SSE event therefore got sliced mid-payload and rejoined with a newline, producing invalid JSON and google.genai.errors.UnknownApiResponseError "Failed to parse response as JSON". The failure was intermittent because it only triggered when the response happened to carry one of those separators, which correlates with large thinking responses. Buffer the raw byte stream and split only on real SSE frame delimiters (\r\n\r\n, \n\n, \r\r). This keeps payload bytes intact regardless of content and still reassembles large inlineData blobs across chunk boundaries, the case the previous aiter_lines change was meant to fix. Fixes LIT-3775 --- any-discipline-budget.json | 4 +- litellm/google_genai/streaming_iterator.py | 107 +++++---- .../test_google_genai_streaming_iterator.py | 203 ++++++++++-------- .../test_google_ai_studio.py | 23 +- 4 files changed, 176 insertions(+), 161 deletions(-) diff --git a/any-discipline-budget.json b/any-discipline-budget.json index d78b15e3653..f086bbdeaf6 100644 --- a/any-discipline-budget.json +++ b/any-discipline-budget.json @@ -308,8 +308,8 @@ "slack": 90 }, "litellm/google_genai/streaming_iterator.py": { - "baseline": 56, - "slack": 28 + "baseline": 41, + "slack": 21 }, "litellm/images/main.py": { "baseline": 326, diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index a8d0e5976f0..f9116656e8e 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -17,41 +17,26 @@ else: GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging() - -def _encode_google_genai_sse_event(event_lines: List[str]) -> bytes: - return ("\n".join(event_lines) + "\n\n").encode("utf-8") +# Buffer raw bytes and split only on real SSE frame delimiters. httpx +# ``aiter_lines`` splits on ``str.splitlines`` boundaries (U+2028, U+2029, +# U+0085, form feed, ...) which Gemini emits raw inside ``data:`` JSON, so it +# would slice an event mid-payload and corrupt the JSON the SDK then parses. +_SSE_FRAME_DELIMITERS: Tuple[bytes, ...] = (b"\r\n\r\n", b"\n\n", b"\r\r") -def _next_google_genai_sse_chunk(line_iter) -> bytes: - event_lines: List[str] = [] - while True: - try: - line = next(line_iter) - except StopIteration: - if event_lines: - return _encode_google_genai_sse_event(event_lines) - raise - if line == "": - if event_lines: - return _encode_google_genai_sse_event(event_lines) - continue - event_lines.append(line) - - -async def _anext_google_genai_sse_chunk(line_iter) -> bytes: - event_lines: List[str] = [] - while True: - try: - line = await line_iter.__anext__() - except StopAsyncIteration: - if event_lines: - return _encode_google_genai_sse_event(event_lines) - raise - if line == "": - if event_lines: - return _encode_google_genai_sse_event(event_lines) - continue - event_lines.append(line) +def _split_sse_frame(buffer: bytes) -> Tuple[Optional[bytes], bytes]: + """Pop the first complete SSE frame (delimiter included) from ``buffer``.""" + frame_starts = ( + (position, delimiter) + for delimiter in _SSE_FRAME_DELIMITERS + if (position := buffer.find(delimiter)) != -1 + ) + first = min(frame_starts, key=lambda item: item[0], default=None) + if first is None: + return None, buffer + position, delimiter = first + frame_end = position + len(delimiter) + return buffer[:frame_end], buffer[frame_end:] class BaseGoogleGenAIGenerateContentStreamingIterator: @@ -127,20 +112,26 @@ class GoogleGenAIGenerateContentStreamingIterator( self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider - # Gemini streamGenerateContent uses SSE line framing; iter_lines keeps - # large inlineData payloads (e.g. image/jpeg) intact within one event. - self.stream_iterator = response.iter_lines() + self.stream_iterator = response.iter_bytes() + self._buffer: bytes = b"" def __iter__(self): return self - def __next__(self): - try: - chunk = _next_google_genai_sse_chunk(self.stream_iterator) - self.collected_chunks.append(chunk) - return chunk - except StopIteration: - raise StopIteration + def __next__(self) -> bytes: + while True: + frame, self._buffer = _split_sse_frame(self._buffer) + if frame is not None: + self.collected_chunks.append(frame) + return frame + try: + self._buffer += next(self.stream_iterator) + except StopIteration: + if self._buffer: + frame, self._buffer = self._buffer, b"" + self.collected_chunks.append(frame) + return frame + raise def __aiter__(self): return self @@ -182,18 +173,24 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator( self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider - # Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps - # large inlineData payloads (e.g. image/jpeg) intact within one event. - self.stream_iterator = response.aiter_lines() + self.stream_iterator = response.aiter_bytes() + self._buffer: bytes = b"" def __aiter__(self): return self - async def __anext__(self): - try: - chunk = await _anext_google_genai_sse_chunk(self.stream_iterator) - self.collected_chunks.append(chunk) - return chunk - except StopAsyncIteration: - await self._handle_async_streaming_logging() - raise StopAsyncIteration + async def __anext__(self) -> bytes: + while True: + frame, self._buffer = _split_sse_frame(self._buffer) + if frame is not None: + self.collected_chunks.append(frame) + return frame + try: + self._buffer += await self.stream_iterator.__anext__() + except StopAsyncIteration: + if self._buffer: + frame, self._buffer = self._buffer, b"" + self.collected_chunks.append(frame) + return frame + await self._handle_async_streaming_logging() + raise diff --git a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py index d74a05ec59c..36246bb8942 100644 --- a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py @@ -10,119 +10,142 @@ from litellm.google_genai.streaming_iterator import ( from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -def _large_inline_data_event() -> str: +def _sse_event(payload: dict) -> bytes: + return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n".encode("utf-8") + + +def _chunk_bytes(data: bytes, size: int) -> list[bytes]: + return [data[i : i + size] for i in range(0, len(data), size)] + + +def _build_iterator(stream_bytes, *, is_async: bool): + mock_response = MagicMock() + if is_async: + + async def _aiter_bytes(): + for chunk in stream_bytes: + yield chunk + + mock_response.aiter_bytes = _aiter_bytes + cls = AsyncGoogleGenAIGenerateContentStreamingIterator + else: + mock_response.iter_bytes.return_value = iter(stream_bytes) + cls = GoogleGenAIGenerateContentStreamingIterator + + return cls( + response=mock_response, + model="gemini-3.1-pro-preview", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + +def _parse_event(frame: bytes) -> dict: + assert frame.startswith(b"data: ") + assert frame.endswith(b"\n\n") + return json.loads(frame[len(b"data: ") : -2]) + + +@pytest.mark.asyncio +async def test_async_iterator_preserves_unicode_line_separators(): + """Regression for LIT-3775: U+2028/U+2029/U+0085 inside ``data:`` JSON. + + Gemini emits these raw in thinking/text content. httpx ``aiter_lines`` + splits on them (``str.splitlines``) and the event got rejoined with ``\\n``, + producing invalid JSON the SDK could not parse. + """ + text = "first
second
third…fourth" + payload = {"candidates": [{"content": {"parts": [{"text": text}]}}]} + wire = _sse_event(payload) + # Split the event into small byte chunks so a separator never aligns with a + # chunk boundary; this is what the upstream HTTP stream looks like. + iterator = _build_iterator(_chunk_bytes(wire, 7), is_async=True) + + frames = [frame async for frame in iterator] + + assert len(frames) == 1 + parsed = _parse_event(frames[0]) + assert parsed["candidates"][0]["content"]["parts"][0]["text"] == text + + +@pytest.mark.asyncio +async def test_async_iterator_reassembles_inline_data_split_across_chunks(): + """Large inlineData payloads must stay in one event even when chunked.""" payload = { "candidates": [ { "content": { "parts": [ - { - "inlineData": { - "mimeType": "image/jpeg", - "data": "A" * 20000, - } - } + {"inlineData": {"mimeType": "image/jpeg", "data": "A" * 20000}} ] } } ] } - return f"data: {json.dumps(payload)}" + wire = _sse_event(payload) + iterator = _build_iterator(_chunk_bytes(wire, 1024), is_async=True) + + frames = [frame async for frame in iterator] + + assert len(frames) == 1 + parsed = _parse_event(frames[0]) + inline = parsed["candidates"][0]["content"]["parts"][0]["inlineData"] + assert inline["mimeType"] == "image/jpeg" + assert inline["data"] == "A" * 20000 -@pytest.mark.asyncio -async def test_async_streaming_iterator_yields_complete_sse_events(): - """Large inlineData must not be split across byte-chunk boundaries.""" - mock_response = MagicMock() +def test_sync_iterator_reassembles_inline_data_split_across_chunks(): + payload = { + "candidates": [ + { + "content": { + "parts": [ + {"inlineData": {"mimeType": "image/jpeg", "data": "A" * 20000}} + ] + } + } + ] + } + wire = _sse_event(payload) + iterator = _build_iterator(_chunk_bytes(wire, 1024), is_async=False) - async def _aiter_lines(): - yield _large_inline_data_event() + frames = list(iterator) - mock_response.aiter_lines = _aiter_lines - - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, - model="gemini-3.1-flash-image-preview", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), - generate_content_provider_config=MagicMock(), - litellm_metadata={}, - custom_llm_provider="gemini", - ) - - chunk = await iterator.__anext__() - assert chunk.startswith(b"data: ") - assert chunk.endswith(b"\n\n") + assert len(frames) == 1 assert ( - json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0][ - "inlineData" - ]["mimeType"] - == "image/jpeg" + _parse_event(frames[0])["candidates"][0]["content"]["parts"][0]["inlineData"][ + "data" + ] + == "A" * 20000 ) -def test_sync_streaming_iterator_yields_complete_sse_events(): - mock_response = MagicMock() - mock_response.iter_lines.return_value = iter([_large_inline_data_event()]) - - iterator = GoogleGenAIGenerateContentStreamingIterator( - response=mock_response, - model="gemini-3.1-flash-image-preview", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), - generate_content_provider_config=MagicMock(), - litellm_metadata={}, - custom_llm_provider="gemini", - ) - - chunk = next(iterator) - assert chunk.startswith(b"data: ") - assert chunk.endswith(b"\n\n") - assert json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][ - 0 - ]["inlineData"]["data"].startswith("A") - - @pytest.mark.asyncio -async def test_async_streaming_iterator_preserves_multi_field_sse_event(): - mock_response = MagicMock() +async def test_async_iterator_splits_multiple_events(): + first = _sse_event({"candidates": [{"content": {"parts": [{"text": "one"}]}}]}) + second = _sse_event({"candidates": [{"content": {"parts": [{"text": "two"}]}}]}) + # Feed both events as a single byte blob to prove the buffer splits on the + # frame delimiter rather than relying on chunk boundaries. + iterator = _build_iterator([first + second], is_async=True) - async def _aiter_lines(): - yield "event: message" - yield 'data: {"text":"hi"}' - yield "" + frames = [frame async for frame in iterator] - mock_response.aiter_lines = _aiter_lines - - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, - model="gemini-test", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), - generate_content_provider_config=MagicMock(), - litellm_metadata={}, - custom_llm_provider="gemini", - ) - - chunk = await iterator.__anext__() - assert chunk == b'event: message\ndata: {"text":"hi"}\n\n' + assert [ + _parse_event(f)["candidates"][0]["content"]["parts"][0]["text"] for f in frames + ] == [ + "one", + "two", + ] @pytest.mark.asyncio -async def test_async_streaming_iterator_forwards_sse_comment_events(): - mock_response = MagicMock() +async def test_async_iterator_preserves_multi_field_and_comment_events(): + multi_field = b'event: message\ndata: {"text":"hi"}\n\n' + comment = b": keepalive\n\n" + iterator = _build_iterator([multi_field, comment], is_async=True) - async def _aiter_lines(): - yield ": keepalive" - yield "" + frames = [frame async for frame in iterator] - mock_response.aiter_lines = _aiter_lines - - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, - model="gemini-test", - logging_obj=MagicMock(spec=LiteLLMLoggingObj), - generate_content_provider_config=MagicMock(), - litellm_metadata={}, - custom_llm_provider="gemini", - ) - - chunk = await iterator.__anext__() - assert chunk == b": keepalive\n\n" + assert frames == [multi_field, comment] diff --git a/tests/unified_google_tests/test_google_ai_studio.py b/tests/unified_google_tests/test_google_ai_studio.py index 3e40fa41089..23cae03afb9 100644 --- a/tests/unified_google_tests/test_google_ai_studio.py +++ b/tests/unified_google_tests/test_google_ai_studio.py @@ -84,15 +84,12 @@ async def test_mock_stream_generate_content_with_tools(): mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - # Mock aiter_lines: yield one line at a time (no trailing newlines), - # with a blank line between events, matching httpx aiter_lines behaviour. - async def mock_aiter_lines(): - yield f"data: {json.dumps(mock_response_chunk)}" - yield "" - yield "data: [DONE]" - yield "" + # Mock aiter_bytes: yield raw SSE frames, matching httpx aiter_bytes. + async def mock_aiter_bytes(): + yield f"data: {json.dumps(mock_response_chunk)}\n\n".encode("utf-8") + yield b"data: [DONE]\n\n" - mock_response.aiter_lines = mock_aiter_lines + mock_response.aiter_bytes = mock_aiter_bytes mock_post.return_value = mock_response print( @@ -335,13 +332,11 @@ async def test_validate_post_request_parameters(): mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - # Mock aiter_lines: yield one line at a time (no trailing newlines), - # with a blank line between events, matching httpx aiter_lines behaviour. - async def mock_aiter_lines(): - yield "data: [DONE]" - yield "" + # Mock aiter_bytes: yield raw SSE frames, matching httpx aiter_bytes. + async def mock_aiter_bytes(): + yield b"data: [DONE]\n\n" - mock_response.aiter_lines = mock_aiter_lines + mock_response.aiter_bytes = mock_aiter_bytes mock_post.return_value = mock_response print("\n--- Testing POST request parameters validation ---")