From 64176efb73a005e835e15037a35a990c9fa52b1c Mon Sep 17 00:00:00 2001 From: Tai An Date: Tue, 21 Apr 2026 12:28:18 -0700 Subject: [PATCH] =?UTF-8?q?perf:=20replace=20O(n=C2=B2)=20accumulated=5Fjs?= =?UTF-8?q?on=20with=20list=20+=20completeness=20heuristic=20in=20gemini?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../vertex_and_google_ai_studio_gemini.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) 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 cd27b4c362a..d2fba8b7295 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 @@ -3117,7 +3117,7 @@ class ModelResponseIterator: self.streaming_response = streaming_response self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json" - self.accumulated_json = "" + self.accumulated_json_chunks: list = [] self.sent_first_chunk = False self.logging_obj = logging_obj self.response_headers = response_headers or {} @@ -3310,16 +3310,16 @@ class ModelResponseIterator: 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) + _stripped = message.rstrip() + if not _stripped or _stripped[-1] not in ('}', ']'): + return None + _full_json = "".join(self.accumulated_json_chunks) 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( @@ -3346,7 +3346,7 @@ class ModelResponseIterator: 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 except ValueError as e: @@ -3368,7 +3368,7 @@ class ModelResponseIterator: 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: