perf(vertex): use list buffer for accumulated JSON streaming shards

[Sanjays2402]

Relevant issues

Fixes #31861

Pre-Submission checklist

• I have added meaningful tests
• My PR passes lint/format for the changed lines and unit tests pass locally
• My PR's scope is as isolated as possible; it only solves 1 specific problem

Type

🚄 Infrastructure (performance)

What / Why

ModelResponseIterator.handle_accumulated_json_chunk (Vertex/Gemini streaming) accumulated partial JSON with self.accumulated_json += message on every SSE shard. Because self also holds a reference to that string, CPython cannot take the in-place concatenation fast path and must copy the entire prior buffer on every shard — making buffer assembly O(n²) in total payload size.

For a multi-MB Gemini tool-call response split across thousands of ~2 KB shards, this copies gigabytes before the single json.loads, showing up as elevated asyncio event-loop latency proportional to payload size.

(This is the buffer-assembly companion to the json.loads-frequency fix in #31297 / #26181, which addressed the parse side of the same function.)

The fix

Store shards in a list[str] (accumulated_json_chunks) and "".join() once when the completeness heuristic passes — O(n) total instead of O(n²).

• The tail check (only attempt json.loads once a shard closes a JSON value) now inspects the last shard, which is O(1), instead of rstrip()-ing the whole joined buffer on every fragment.
• The end-of-stream flush (chunk == "") always attempts a parse, so a complete buffer is never dropped even if the final shard were whitespace-only (correctness safeguard).
• A backward-compatible accumulated_json str property (getter joins, setter replaces) preserves the previous read/write interface for any caller/subclass referencing it directly — including the existing truthiness checks in __next__/__anext__.

Test

Added a TestAccumulatedJsonBuffer class in the existing Vertex/Gemini test file (mocked, no API calls):

• test_sharded_json_accumulates_and_parses — a JSON object split across 23 shards parses on the final shard and resets the buffer
• test_incomplete_buffer_returns_none_without_reset — fragments that don't close a value keep accumulating
• test_flush_parses_buffer_with_trailing_whitespace_shard — the flush safeguard recovers a complete buffer even when the last shard is whitespace-only
• test_accumulated_json_property_backward_compat — the str property still reads/writes the list buffer

uv run pytest tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py -q
# 137 passed  (127 pre-existing + 4 new + 6 in the streaming tool-call file)

Proof of Fix

Microbenchmark driving handle_accumulated_json_chunk with 4000 shards of 2 KB each (the issue's reproducer scenario), measuring buffer assembly:

BEFORE (str += per shard):  ~596 ms
AFTER  (list buffer):         ~4 ms

~130–200x faster buffer assembly, and it no longer holds the GIL/event loop proportional to payload size. Correctness is covered by the 137 passing unit tests above (the sharded-parse test asserts the exact same parsed output as before).

[codecov]

Codecov Report

❌ Patch coverage is 94.44444% with 1 line in your changes missing coverage. Please review.

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

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR replaces the O(n²) str += shard accumulation in ModelResponseIterator.handle_accumulated_json_chunk with a list[str] buffer joined once when the completeness heuristic passes, eliminating the quadratic copy cost for large Gemini streaming responses. A backward-compatible accumulated_json property (getter joins, setter replaces) preserves the previous attribute interface for any existing callers.

• The last-shard heuristic (last_shard[-1] not in \"}]\") correctly replaces the whole-buffer rstrip() check with an O(1) peek, and the end-of-stream flush (chunk=\"\") ensures a complete buffer is never silently dropped when the final SSE shard is whitespace-only.
• Four new mock-only unit tests cover shard accumulation, incomplete-buffer deferral, the whitespace-shard flush safeguard, and the backward-compat property.
• The two StopIteration/StopAsyncIteration flush guards still call self.accumulated_json (the property getter, O(n) join) for a simple truthiness check; self.accumulated_json_chunks would be O(1) and semantically equivalent.

Confidence Score: 4/5

Safe to merge — the change is isolated to the Vertex/Gemini streaming buffer, logic is correct, and all 137 unit tests pass.

The optimization is well-implemented and the flush safeguard correctly handles the edge case where the final shard is whitespace-only. The only remaining issue is that the StopIteration/StopAsyncIteration flush guards call the property getter (an O(n) join) just to test truthiness, which is a small inconsistency in a PR explicitly targeting unnecessary buffer copies. No correctness bugs were found.

The two flush guards in __next__ and __anext__ (lines 3380 and 3402) are the only spots worth a second look.

Important Files Changed

Filename	Overview
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py	Core perf fix: replaces O(n²) str+= shard accumulation with a list buffer; backward-compat property added. The last-shard heuristic and end-of-stream flush logic are correct. Minor: the StopIteration guards still call the O(n) property getter for a truthiness check when bool(accumulated_json_chunks) would be O(1).
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py	Adds four mock-only unit tests covering shard accumulation, incomplete-buffer deferral, whitespace-shard flush, and the backward-compat property. All tests are appropriately mocked with no real network calls.

Comments Outside Diff (2)

1. litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py, line 3379-3382 (link)

   Both flush guards invoke the accumulated_json property getter — which calls "".join(self.accumulated_json_chunks) — merely to test truthiness. Since only non-empty strings are ever appended to the list, bool(accumulated_json_chunks) is semantically identical and avoids the O(n) join.

   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!

2. litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py, line 3401-3404 (link)

   Same O(n) getter call in the async flush guard — replace with the direct list check for consistency.

   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): use list buffer for accumu..." | Re-trigger Greptile}

[deepanshululla]

Vertex's function changed since this opened (now uses raw_decode). Opened #36610: a shared helper generalizing the fix to cover Anthropic too.

[Sanjays2402]

The Vertex-specific implementation here has been superseded by #36610. I checked that PR: it adds the shared JSONFragmentAccumulator and migrates this Vertex path (plus Anthropic), with the focused accumulator, Vertex, and Anthropic tests. I am closing this narrower PR so the fix is tracked in one place.