Skip to content

perf(vertex/gemini): eliminate O(n²) json.loads on accumulated streaming chunks (#26181) #29296

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Anai-Guo wants to merge 2 commits into from
Closed

perf(vertex/gemini): eliminate O(n²) json.loads on accumulated streaming chunks (#26181) #29296

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9892459
perf(vertex/gemini): eliminate O(n²) json.loads on accumulated stream…
Anai-Guo May 29, 2026
92b3913
style: apply black formatting to vertex gemini accumulator
Anai-Guo May 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 21 additions & 10 deletions litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3263,7 +3263,12 @@ def __init__(

self.streaming_response = streaming_response
self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json"
self.accumulated_json = ""
# Buffer SSE shards as a list to avoid O(n²) string concatenation when
# Gemini emits multi-MB tool-call payloads as many small chunks. The
# parser only joins+json.loads when the trailing brace/bracket suggests
# a complete envelope, instead of re-parsing the growing buffer on every
# chunk and freezing the event loop.
self.accumulated_json_chunks: list = []

@greptile-apps greptile-apps Bot May 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 The type annotation list is valid but imprecise — the list holds str elements only. Using list[str] makes the contract explicit for static analysis and readers.

Suggested change
self.accumulated_json_chunks: list = []
self.accumulated_json_chunks: list[str] = []

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

self.sent_first_chunk = False
self.logging_obj = logging_obj
self.response_headers = response_headers or {}
Expand Down Expand Up @@ -3486,16 +3491,22 @@ def handle_accumulated_json_chunk(
chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or ""
message = chunk.replace("\n\n", "")

# Accumulate JSON data
self.accumulated_json += message

# Try to parse the accumulated JSON
self.accumulated_json_chunks.append(message)
# Cheap completeness heuristic: only join + json.loads when the latest
# shard looks like it might have closed the envelope. Avoids parsing
# the growing buffer on every chunk -- the empty-string flush from
# __next__/__anext__ at StopIteration still forces a final attempt.
_stripped = message.rstrip()
if _stripped and _stripped[-1] not in ("}", "]"):
return None
_full_json = "".join(self.accumulated_json_chunks)
if not _full_json:
return None
try:
_data = json.loads(self.accumulated_json)
self.accumulated_json = "" # reset after successful parsing
_data = json.loads(_full_json)
self.accumulated_json_chunks = []
return self.chunk_parser(chunk=_data)
except json.JSONDecodeError:
# If it's not valid JSON yet, continue to the next event
return None

def _common_chunk_parsing_logic(
Expand All @@ -3522,7 +3533,7 @@ def __next__(self):
try:
chunk = self.response_iterator.__next__()
except StopIteration:
if self.chunk_type == "accumulated_json" and self.accumulated_json:
if self.chunk_type == "accumulated_json" and self.accumulated_json_chunks:
return self.handle_accumulated_json_chunk(chunk="")
raise StopIteration
Comment on lines +3536 to 3538

@greptile-apps greptile-apps Bot May 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Unbounded list growth on repeated flush calls

handle_accumulated_json_chunk("") appends an empty string to accumulated_json_chunks every time the flush path is taken. If json.loads fails on the incomplete/invalid accumulated JSON at stream end and the consumer's outer loop calls __next__ again, each call appends another "" — the list grows indefinitely. The original self.accumulated_json += "" was a no-op and did not change the buffer. Functionally the join result is identical (empty strings are transparent to "".join), but this is a silent behavioral divergence from the original: memory usage grows with every repeated call to the flush path. Guarding against repeated flush appends (e.g., by not appending when chunk is empty) would preserve the no-op semantics of the original.

except ValueError as e:
Expand All @@ -3544,7 +3555,7 @@ async def __anext__(self):
try:
chunk = await self.async_response_iterator.__anext__()
except StopAsyncIteration:
if self.chunk_type == "accumulated_json" and self.accumulated_json:
if self.chunk_type == "accumulated_json" and self.accumulated_json_chunks:
return self.handle_accumulated_json_chunk(chunk="")
raise StopAsyncIteration
except ValueError as e:
Expand Down
Loading