diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 624190a0b61..8dd0dc19b81 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3333,27 +3333,37 @@ class ModelResponseIterator: return self.chunk_parser(chunk=json_chunk) - def handle_accumulated_json_chunk(self, chunk: str) -> Optional["ModelResponseStream"]: - chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" - message = chunk.replace("\n\n", "") + def handle_accumulated_json_chunk(self, chunk: str, is_final: bool = False) -> Optional["ModelResponseStream"]: + message = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" + self.accumulated_json = (self.accumulated_json + message.replace("\n\n", "")).strip() - self.accumulated_json += message - - # json.loads on the whole buffer after every fragment is O(n^2) and - # holds the GIL, freezing the event loop for seconds on large responses - # (https://github.com/BerriAI/litellm/issues/26181). A complete Gemini - # chunk is a JSON object/array, so only attempt the parse once the - # buffer's last non-whitespace byte can close one. - stripped = self.accumulated_json.rstrip() - if not stripped or stripped[-1] not in "}]": + # Mid-stream, defer parsing until the buffer's last byte can close a value: + # attempting a parse after every fragment of one large object is O(n^2) and + # holds the GIL, freezing the event loop. At end of stream (is_final) no more + # data is coming, so drain whatever complete values remain regardless of the + # trailing byte, otherwise a complete leading value sitting behind a truncated + # trailing one would be silently dropped. + if not is_final and (not self.accumulated_json or self.accumulated_json[-1] not in "}]"): return None - try: - _data = json.loads(self.accumulated_json) - self.accumulated_json = "" # reset after successful parsing - return self.chunk_parser(chunk=_data) - except json.JSONDecodeError: - return None + # Peel one complete JSON value from the front of the buffer and keep the + # unconsumed tail. Running json.loads over the whole buffer would fail + # forever once it held more than one concatenated value ("Extra data") while + # never resetting the buffer, so the buffer grew without bound and pinned the + # core. raw_decode reports where the value ended, so concatenated values drain + # one call at a time. A leading non-dict value (never emitted by Gemini in + # practice) is consumed and skipped so it cannot block the dict values behind it. + decoder = json.JSONDecoder() + while self.accumulated_json: + try: + raw_value = decoder.raw_decode(self.accumulated_json) + except json.JSONDecodeError: + return None + decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode -> tuple[Any,int] + self.accumulated_json = self.accumulated_json[end_index:].strip() + if isinstance(decoded, dict): + return self.chunk_parser(chunk=decoded) + return None def _common_chunk_parsing_logic(self, chunk: str) -> Optional["ModelResponseStream"]: try: @@ -3378,7 +3388,9 @@ class ModelResponseIterator: chunk = self.response_iterator.__next__() except StopIteration: if self.chunk_type == "accumulated_json" and self.accumulated_json: - return self.handle_accumulated_json_chunk(chunk="") + result = self.handle_accumulated_json_chunk(chunk="", is_final=True) + if result is not None: + return result raise StopIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") @@ -3400,7 +3412,9 @@ class ModelResponseIterator: chunk = await self.async_response_iterator.__anext__() except StopAsyncIteration: if self.chunk_type == "accumulated_json" and self.accumulated_json: - return self.handle_accumulated_json_chunk(chunk="") + result = self.handle_accumulated_json_chunk(chunk="", is_final=True) + if result is not None: + return result raise StopAsyncIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 3723976191b..5871644ca5f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5407,3 +5407,116 @@ def test_process_candidates_merges_thought_signatures_and_server_side_tools(): fields = model_response.choices[-1].message.provider_specific_fields assert fields["thought_signatures"] == ["sig-text"] assert fields["server_side_tool_invocations"][0]["id"] == "tool-1" + + +def _accumulating_gemini_iterator(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=[], sync_stream=True, logging_obj=MagicMock() + ) + iterator.chunk_type = "accumulated_json" + return iterator + + +def test_accumulated_json_chunk_multi_value_buffer_does_not_wedge(): + """Two complete Gemini objects buffered together must both surface. + + A whole-buffer json.loads raises "Extra data" on concatenated values and, since + the buffer was never reset on failure, returned None forever while growing without + bound. Peeling one value from the front keeps the remainder for the next call. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + + first = iterator.handle_accumulated_json_chunk(chunk=obj + obj) + assert first is not None + assert first.choices[0].delta.content == "a" + + second = iterator.handle_accumulated_json_chunk(chunk="") + assert second is not None + assert second.choices[0].delta.content == "a" + + assert iterator.accumulated_json.strip() == "" + + +def test_accumulated_json_end_of_stream_drains_all_buffered_values(): + """End of stream must drain every buffered value and then terminate. + + With concatenated values a whole-buffer parse never succeeds, so __next__ kept + returning None without shrinking the buffer - an unrecoverable per-request spin. + The bounded loop asserts the iterator both surfaces all values and terminates. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = obj + obj + obj + + out = [] + terminated = False + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + terminated = True + break + if chunk is not None: + out.append(chunk) + + assert terminated, "iterator did not terminate - accumulated buffer wedged" + assert len(out) == 3 + assert iterator.accumulated_json.strip() == "" + + +def test_accumulated_json_end_of_stream_surfaces_leading_value_before_truncated_tail(): + """A complete leading value must survive a truncated trailing value at end of stream. + + The mid-stream perf guard only inspects the buffer's last byte, so a complete leading + object followed by a truncated one (a server that cut the stream mid-object, last byte + not a closer) would keep the guard from ever parsing and drop the complete value. At end + of stream the drain ignores that guard, surfaces the complete value, and discards only + the truncated tail. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = obj + '{"candidates":' + + out = [] + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + break + if chunk is not None: + out.append(chunk) + + assert len(out) == 1 + assert out[0].choices[0].delta.content == "a" + + +def test_accumulated_json_skips_non_dict_leading_value(): + """A non-dict value at the front must not block the dict values behind it. + + raw_decode advances past a decoded value, so a leading non-dict (a JSON array or scalar, + which Gemini never emits but a malformed stream could) must be consumed and skipped. If + the drain stopped on it, the trailing objects would be lost at end of stream. + """ + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.response_iterator = iter([]) + iterator.accumulated_json = "[1, 2]" + obj + + out = [] + for _ in range(100): + try: + chunk = iterator.__next__() + except StopIteration: + break + if chunk is not None: + out.append(chunk) + + assert len(out) == 1 + assert out[0].choices[0].delta.content == "a"