From 729b005e4e00ef43f6187baf7fcd8f2737c764b8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 12 Jun 2026 20:19:30 +0530 Subject: [PATCH] fix(google_genai): preserve complete SSE events in Vertex/Gemini image streaming (#30270) * fix(google_genai): preserve complete SSE events in image streaming Use iter_lines/aiter_lines instead of byte chunking so large inlineData base64 payloads from Vertex/Gemini streamGenerateContent are not split across events, which caused truncated JSON and SDK parse failures. Co-authored-by: Cursor * fix(google_genai): buffer SSE lines until event delimiter Assemble multi-field SSE events on blank-line boundaries instead of terminating each field line individually. Co-authored-by: Cursor * fix(tests): update google_ai_studio mocks from aiter_bytes to aiter_lines Streaming iterator was changed to use iter_lines/aiter_lines instead of iter_bytes/aiter_bytes. Update the two mocked streaming responses in test_google_ai_studio.py to match. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- litellm/google_genai/streaming_iterator.py | 54 ++++++-- .../test_google_genai_streaming_iterator.py | 128 ++++++++++++++++++ .../test_google_ai_studio.py | 33 ++--- 3 files changed, 186 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 3e97b480779..a8d0e5976f0 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -18,6 +18,42 @@ 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") + + +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) + + class BaseGoogleGenAIGenerateContentStreamingIterator: """ Base class for Google GenAI Generate Content streaming iterators that provides common logic @@ -91,18 +127,17 @@ class GoogleGenAIGenerateContentStreamingIterator( self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider - # Store the iterator once to avoid multiple stream consumption - self.stream_iterator = response.iter_bytes() + # 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() def __iter__(self): return self def __next__(self): try: - # Get the next chunk from the stored iterator - chunk = next(self.stream_iterator) + chunk = _next_google_genai_sse_chunk(self.stream_iterator) self.collected_chunks.append(chunk) - # Just yield raw bytes return chunk except StopIteration: raise StopIteration @@ -147,18 +182,17 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator( self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider - # Store the async iterator once to avoid multiple stream consumption - self.stream_iterator = response.aiter_bytes() + # 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() def __aiter__(self): return self async def __anext__(self): try: - # Get the next chunk from the stored async iterator - chunk = await self.stream_iterator.__anext__() + chunk = await _anext_google_genai_sse_chunk(self.stream_iterator) self.collected_chunks.append(chunk) - # Just yield raw bytes return chunk except StopAsyncIteration: await self._handle_async_streaming_logging() 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 new file mode 100644 index 00000000000..d74a05ec59c --- /dev/null +++ b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py @@ -0,0 +1,128 @@ +import json +from unittest.mock import MagicMock + +import pytest + +from litellm.google_genai.streaming_iterator import ( + AsyncGoogleGenAIGenerateContentStreamingIterator, + GoogleGenAIGenerateContentStreamingIterator, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +def _large_inline_data_event() -> str: + payload = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/jpeg", + "data": "A" * 20000, + } + } + ] + } + } + ] + } + return f"data: {json.dumps(payload)}" + + +@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() + + async def _aiter_lines(): + yield _large_inline_data_event() + + 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 ( + json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0][ + "inlineData" + ]["mimeType"] + == "image/jpeg" + ) + + +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 _aiter_lines(): + yield "event: message" + yield 'data: {"text":"hi"}' + yield "" + + 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' + + +@pytest.mark.asyncio +async def test_async_streaming_iterator_forwards_sse_comment_events(): + mock_response = MagicMock() + + async def _aiter_lines(): + yield ": keepalive" + yield "" + + 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" diff --git a/tests/unified_google_tests/test_google_ai_studio.py b/tests/unified_google_tests/test_google_ai_studio.py index afe237a4e5b..3e40fa41089 100644 --- a/tests/unified_google_tests/test_google_ai_studio.py +++ b/tests/unified_google_tests/test_google_ai_studio.py @@ -74,12 +74,6 @@ async def test_mock_stream_generate_content_with_tools(): }, } - # Convert to bytes as expected by the streaming iterator - raw_chunks = [ - f"data: {json.dumps(mock_response_chunk)}\n\n".encode(), - b"data: [DONE]\n\n", - ] - # Mock the HTTP handler with unittest.mock.patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", @@ -90,12 +84,15 @@ async def test_mock_stream_generate_content_with_tools(): mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - # Mock the aiter_bytes method to return our chunks as bytes - async def mock_aiter_bytes(): - for chunk in raw_chunks: - yield chunk + # 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_response.aiter_bytes = mock_aiter_bytes + mock_response.aiter_lines = mock_aiter_lines mock_post.return_value = mock_response print( @@ -328,9 +325,6 @@ async def test_validate_post_request_parameters(): } ] - # Mock response for the HTTP request - raw_chunks = [b"data: [DONE]\n\n"] - # Mock the HTTP handler to capture the request with unittest.mock.patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", @@ -341,12 +335,13 @@ async def test_validate_post_request_parameters(): mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - # Mock the aiter_bytes method - async def mock_aiter_bytes(): - for chunk in raw_chunks: - yield chunk + # 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_response.aiter_bytes = mock_aiter_bytes + mock_response.aiter_lines = mock_aiter_lines mock_post.return_value = mock_response print("\n--- Testing POST request parameters validation ---")