mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
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:
parent
e8e5b47fc2
commit
729b005e4e
3 changed files with 186 additions and 29 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -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 ---")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue