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 <cursoragent@cursor.com>

* 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 <cursoragent@cursor.com>

* 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 <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Sameer Kankute 2026-06-12 20:19:30 +05:30 • committed by GitHub
parent e8e5b47fc2
commit 729b005e4e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 186 additions and 29 deletions

View file

@ -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()

View file

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

View file

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