fix(google_genai): preserve complete SSE events in Vertex/Gemini image streaming

litellm/google_genai/streaming_iterator.py

@@ -18,6 +18,42 @@
 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 @@ def __init__(
         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 @@ def __init__(
         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()

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"

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 ---")