scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

perf(sagemaker): replace O(n²) accumulated_json string concat with list + completeness heuristic

Browse source
This commit is contained in:
Tai An 2026-04-21 12:27:14 -07:00
parent 26fcbc93e5
commit 89e8a5211b
1 changed files with 24 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/llms/sagemaker/common_utils.py
View file

@ -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: