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

[Anai-Guo]

Summary

Closes #26181 (re-PR of #26187 — that PR was auto-closed when the daily-dated litellm_oss_staging_04_25_2026 base branch was retired in favor of litellm_internal_staging).

When Vertex AI (Gemini) streams fragmented SSE chunks, handle_accumulated_json_chunk previously:

1. Concatenated every chunk into a single growing string (self.accumulated_json += chunk, O(n) copy each call → O(n²) total)
2. Called json.loads() on the full accumulated buffer after every chunk

For multi-MB tool-call payloads or long code generations, json.loads — a single CPython C call that holds the GIL — blocks for seconds at a time. In the default single-process uvicorn deployment this freezes the asyncio event loop, liveness probes time out, and kubelet restarts the pod.

The reporter (@6matt) confirmed this with 7 consecutive py-spy thread dumps spanning ~70 seconds, all showing MainThread stuck in json.loads → json.decoder.decode → json.decoder.raw_decode from handle_accumulated_json_chunk.

Fix

Two changes:

1. List-based accumulation — replace self.accumulated_json += chunk with self.accumulated_json_chunks.append(chunk). O(n²) → O(n) total copy work.
2. Completeness heuristic — only attempt "".join(...) + json.loads() when the latest shard's last non-whitespace character is } or ]. Avoids repeatedly parsing the growing buffer on every chunk. The empty-string flush from __next__ / __anext__ at StopIteration still forces one final parse attempt, so partial-completion semantics are unchanged.

Scope

• One file (litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py).
• 10 additions / 10 deletions (per the original perf(streaming): eliminate O(n²) json.loads in accumulated-JSON paths (Vertex, Anthropic, SageMaker) #26187 diff against the same call sites).
• No public API change; field rename only (accumulated_json: str → accumulated_json_chunks: list).
• Parallel fixes for litellm/llms/anthropic/chat/handler.py and litellm/llms/sagemaker/common_utils.py will follow in separate PRs after this lands, so each provider can be reviewed independently.

Test Plan

• Same-shape unit tests as perf(streaming): eliminate O(n²) json.loads in accumulated-JSON paths (Vertex, Anthropic, SageMaker) #26187 (Gemini path): chunked JSON arrives as multiple shards, only one json.loads call after the closing brace, final state cleared.
• Add the tests in a follow-up commit if the maintainers prefer them bundled rather than split — let me know.

🤖 Generated with Claude Code (AI-assisted, human reviewed)

[greptile-apps]

Greptile Summary

Replaces the O(n²) string-concatenation pattern in ModelResponseIterator.handle_accumulated_json_chunk with list-based shard accumulation and a closing-brace heuristic that skips json.loads on chunks that have not yet closed the JSON envelope, preventing event-loop stalls on large Vertex AI / Gemini streaming payloads.

• accumulated_json: str is replaced by accumulated_json_chunks: list[str]; list append is O(1) vs the previous O(n) per chunk, and json.loads is only called when the latest shard ends with } or ], with a final forced attempt at StopIteration.
• The two StopIteration / StopAsyncIteration guards in __next__ and __anext__ are updated to test the non-empty list instead of the non-empty string.
• No test files are changed; the PR description's unit-test checkbox is explicitly left unchecked.

Confidence Score: 4/5

The change is narrow (one file, one method) and the heuristic is safe for Gemini's object/array envelope format; the main gap is the absence of inline tests.

The accumulation and heuristic logic is correct for Gemini's response format, and the StopIteration flush path still forces a final parse attempt. The flush path does introduce a minor behavioural divergence from the original (appending an empty string on each call rather than a no-op concatenation), and no automated tests ship with the PR to guard against regressions in the multi-shard assembly logic.

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py — specifically the flush path in handle_accumulated_json_chunk and the missing test coverage.

Important Files Changed

Filename	Overview
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py	Replaces O(n²) string concatenation + repeated json.loads with list-based accumulation and a closing-brace heuristic; logic is sound for Gemini's object/array response envelope, but no tests are included and the flush path introduces a minor list-growth divergence from the original no-op behaviour.

Comments Outside Diff (1)

1. litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py, line 3488-3510 (link)

   No tests included in this PR

   The test plan in the PR description explicitly leaves the unit-test checkbox unchecked, with a promise to add them in a follow-up. The repository's review rules require any PR claiming to fix an issue to include evidence of resolution (passing tests or a detailed summary of changes). The summary is present, but the automated tests that would verify the O(n²) regression is fixed and that multi-shard JSON is assembled correctly are absent. Adding at least the "same-shape unit tests as perf(streaming): eliminate O(n²) json.loads in accumulated-JSON paths (Vertex, Anthropic, SageMaker) #26187" inline — covering chunked JSON arriving in multiple shards, exactly one json.loads call after the closing brace, and buffer cleared on success — would satisfy this requirement.

   Rule Used: What: Ensure that any PR claiming to fix an issue ... (source)

   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!

_{Reviews (1): Last reviewed commit: "perf(vertex/gemini): eliminate O(n²) jso..." | Re-trigger Greptile}

[greptile-apps]

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.

[greptile-apps]

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!

[codecov]

Codecov Report

❌ Patch coverage is 41.66667% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ex_ai/gemini/vertex_and_google_ai_studio_gemini.py	41.66%	7 Missing ⚠️

📢 Thoughts on this report? Let us know!

[hclsys]

list-accumulate + join-on-demand is the right shape, and you updated both next and anext flush guards to accumulated_json_chunks — easy spot to miss one and silently break the async final-flush, so good.

one thing on the heuristic: _stripped[-1] not in ('}',']') -> return None means a shard ending in } inside a string (e.g. ..."x": "a}") still triggers a join+json.loads that fails and returns None — harmless, just a wasted parse. and a genuinely complete envelope always ends in } for gemini SSE so it's never wrongly deferred. only real effect is the multi-shard tool-call case now parses once at the end instead of incrementally — intended, since the whole point is to stop re-parsing the growing buffer, right?

[Anai-Guo]

Yep, exactly that. The whole point is to stop re-json.loads-ing the growing buffer on every chunk — multi-shard tool-call payloads now parse once at the closing brace instead of N times, and the incremental parses were pure wasted work since a partial envelope can't deserialize anyway.

On the "a}"-inside-a-string case: agreed it's harmless. The heuristic only decides when to attempt a parse, never whether the result is valid — a shard ending in a } that's inside a string triggers one json.loads that raises, we catch it and return None, and the caller keeps accumulating until the real envelope closes. Worst case is one wasted parse, never a wrong/early return. And since Gemini SSE always closes the envelope with a real }/], a genuinely complete object is never deferred.

Happy to add the same-shape unit tests as #26187 inline here (multi-shard assembly, exactly-one-parse-after-close, buffer-cleared-on-success) to close greptile's P2 — just say the word.

[Anai-Guo]

Superseded by #31297 (merged 2026-06-26), which resolved the same O(n^2) re-parse of the accumulated Gemini stream. Closing — thanks!