mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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
This commit is contained in:
mateo-berri
parent
17b88719a2
commit
393d252eef
4 changed files with 176 additions and 161 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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 ---")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue