scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(vertex): incrementally parse accumulated Gemini stream JSON to prevent multi-value wedge (#34320)

Browse source 
The accumulated-JSON fallback ran json.loads over the whole buffer after every fragment and, on failure, kept the buffer without resetting it. A buffer that ever held more than one concatenated JSON value could never parse (json raises on trailing data), so it returned None on every subsequent chunk while growing without bound - an unrecoverable per-request CPU spin. Parse one value at a time from the front with raw_decode and keep the remainder, draining trailing values on later calls and at end of stream.
This commit is contained in:
Yassin Kortam 2026-07-24 11:07:02 -07:00 • committed by GitHub
parent 692b22655e
commit f6a1050cbf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 147 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

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

113
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

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