diff --git a/any-discipline-budget.json b/any-discipline-budget.json index d78b15e3653..f086bbdeaf6 100644 --- a/any-discipline-budget.json +++ b/any-discipline-budget.json @@ -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, diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index a8d0e5976f0..f9116656e8e 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -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 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 index d74a05ec59c..36246bb8942 100644 --- a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py @@ -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] diff --git a/tests/unified_google_tests/test_google_ai_studio.py b/tests/unified_google_tests/test_google_ai_studio.py index 3e40fa41089..23cae03afb9 100644 --- a/tests/unified_google_tests/test_google_ai_studio.py +++ b/tests/unified_google_tests/test_google_ai_studio.py @@ -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 ---")