From 89e8a5211be726a41cc66a7ee1c45730da4f3765 Mon Sep 17 00:00:00 2001 From: Tai An Date: Tue, 21 Apr 2026 12:27:14 -0700 Subject: [PATCH] =?UTF-8?q?perf(sagemaker):=20replace=20O(n=C2=B2)=20accum?= =?UTF-8?q?ulated=5Fjson=20string=20concat=20with=20list=20+=20completenes?= =?UTF-8?q?s=20heuristic?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- litellm/llms/sagemaker/common_utils.py | 51 ++++++++++++-------------- 1 file changed, 24 insertions(+), 27 deletions(-) diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index ad6b24d85a3..756842bbc25 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -72,7 +72,7 @@ class AWSEventStreamDecoder: from botocore.eventstream import EventStreamBuffer event_stream_buffer = EventStreamBuffer() - accumulated_json = "" + accumulated_json_chunks: list = [] for chunk in iterator: event_stream_buffer.add_data(chunk) @@ -86,34 +86,32 @@ class AWSEventStreamDecoder: ) message = message.replace("\n\n", "") - # Accumulate JSON data - accumulated_json += message - - # Try to parse the accumulated JSON + accumulated_json_chunks.append(message) + _stripped = message.rstrip() + if not _stripped or _stripped[-1] not in ('}', ']'): + continue + _full_json = "".join(accumulated_json_chunks) try: - _data = json.loads(accumulated_json) + _data = json.loads(_full_json) if self.is_messages_api: yield self._chunk_parser_messages_api(chunk_data=_data) else: yield self._chunk_parser(chunk_data=_data) - # Reset accumulated_json after successful parsing - accumulated_json = "" + accumulated_json_chunks = [] except json.JSONDecodeError: - # If it's not valid JSON yet, continue to the next event continue - # Handle any remaining data after the iterator is exhausted - if accumulated_json: + if accumulated_json_chunks: + _full_json = "".join(accumulated_json_chunks) try: - _data = json.loads(accumulated_json) + _data = json.loads(_full_json) if self.is_messages_api: yield self._chunk_parser_messages_api(chunk_data=_data) else: yield self._chunk_parser(chunk_data=_data) except json.JSONDecodeError: - # Handle or log any unparseable data at the end verbose_logger.error( - f"Warning: Unparseable JSON data remained: {accumulated_json}" + f"Warning: Unparseable JSON data remained: {_full_json}" ) yield None @@ -124,7 +122,7 @@ class AWSEventStreamDecoder: from botocore.eventstream import EventStreamBuffer event_stream_buffer = EventStreamBuffer() - accumulated_json = "" + accumulated_json_chunks: list = [] async for chunk in iterator: event_stream_buffer.add_data(chunk) @@ -144,17 +142,17 @@ class AWSEventStreamDecoder: ) message = message.replace("\n\n", "") - # Accumulate JSON data - accumulated_json += message - - # Try to parse the accumulated JSON - _data = json.loads(accumulated_json) + accumulated_json_chunks.append(message) + _stripped = message.rstrip() + if not _stripped or _stripped[-1] not in ('}', ']'): + continue + _full_json = "".join(accumulated_json_chunks) + _data = json.loads(_full_json) if self.is_messages_api: yield self._chunk_parser_messages_api(chunk_data=_data) else: yield self._chunk_parser(chunk_data=_data) - # Reset accumulated_json after successful parsing - accumulated_json = "" + accumulated_json_chunks = [] except json.JSONDecodeError: # If it's not valid JSON yet, continue to the next event continue @@ -169,18 +167,17 @@ class AWSEventStreamDecoder: ) continue - # Handle any remaining data after the iterator is exhausted - if accumulated_json: + if accumulated_json_chunks: + _full_json = "".join(accumulated_json_chunks) try: - _data = json.loads(accumulated_json) + _data = json.loads(_full_json) if self.is_messages_api: yield self._chunk_parser_messages_api(chunk_data=_data) else: yield self._chunk_parser(chunk_data=_data) except json.JSONDecodeError: - # Handle or log any unparseable data at the end verbose_logger.error( - f"Warning: Unparseable JSON data remained: {accumulated_json}" + f"Warning: Unparseable JSON data remained: {_full_json}" ) yield None except Exception as e: