From 12639b7ccfff372f43b6fc8da8d1ba4dd5f9662d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 24 Mar 2025 21:29:29 -0700 Subject: [PATCH 1/4] fix sagemaker streaming error --- litellm/llms/sagemaker/common_utils.py | 48 +++++++++++++++++--------- 1 file changed, 32 insertions(+), 16 deletions(-) diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 9884f420c3a..48cddf8064b 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -128,21 +128,25 @@ class AWSEventStreamDecoder: async for chunk in iterator: event_stream_buffer.add_data(chunk) for event in event_stream_buffer: - message = self._parse_message_from_event(event) - if message: - verbose_logger.debug("sagemaker parsed chunk bytes %s", message) - # remove data: prefix and "\n\n" at the end - message = ( - litellm.CustomStreamWrapper._strip_sse_data_from_chunk(message) - or "" - ) - message = message.replace("\n\n", "") + try: + message = self._parse_message_from_event(event) + if message: + verbose_logger.debug( + "sagemaker parsed chunk bytes %s", message + ) + # remove data: prefix and "\n\n" at the end + message = ( + litellm.CustomStreamWrapper._strip_sse_data_from_chunk( + message + ) + or "" + ) + message = message.replace("\n\n", "") - # Accumulate JSON data - accumulated_json += message + # Accumulate JSON data + accumulated_json += message - # Try to parse the accumulated JSON - try: + # Try to parse the accumulated JSON _data = json.loads(accumulated_json) if self.is_messages_api: yield self._chunk_parser_messages_api(chunk_data=_data) @@ -150,9 +154,19 @@ class AWSEventStreamDecoder: yield self._chunk_parser(chunk_data=_data) # Reset accumulated_json after successful parsing accumulated_json = "" - except json.JSONDecodeError: - # If it's not valid JSON yet, continue to the next event - continue + except json.JSONDecodeError: + # If it's not valid JSON yet, continue to the next event + continue + except UnicodeDecodeError as e: + verbose_logger.warning( + f"UnicodeDecodeError: {e}. Attempting to combine with next event." + ) + continue + except Exception as e: + verbose_logger.error( + f"Error parsing message: {e}. Attempting to combine with next event." + ) + continue # Handle any remaining data after the iterator is exhausted if accumulated_json: @@ -168,6 +182,8 @@ class AWSEventStreamDecoder: f"Warning: Unparseable JSON data remained: {accumulated_json}" ) yield None + except Exception as e: + verbose_logger.error(f"Final error parsing accumulated JSON: {e}") def _parse_message_from_event(self, event) -> Optional[str]: response_dict = event.to_response_dict() From bc6682753776873d05194f2fcfb89294e50d3e2c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 16:04:38 -0700 Subject: [PATCH 2/4] test_aiter_bytes_valid_chunk_followed_by_unicode_error --- .../sagemaker/test_sagemaker_common_utils.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) create mode 100644 tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py diff --git a/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py b/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py new file mode 100644 index 00000000000..e19c41f1ac3 --- /dev/null +++ b/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py @@ -0,0 +1,93 @@ +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) +from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder + + +@pytest.mark.asyncio +async def test_aiter_bytes_unicode_decode_error(): + """ + Test that AWSEventStreamDecoder.aiter_bytes() does not raise an error when encountering invalid UTF-8 bytes. (UnicodeDecodeError) + + + Ensures stream processing continues despite the error. + """ + # Create an instance of AWSEventStreamDecoder + decoder = AWSEventStreamDecoder(model="test-model") + + # Create a mock event that will trigger a UnicodeDecodeError + mock_event = MagicMock() + mock_event.to_response_dict.return_value = { + "status_code": 200, + "headers": {}, + "body": b"\xff\xfe", # Invalid UTF-8 bytes + } + + # Create a mock EventStreamBuffer that yields our mock event + mock_buffer = MagicMock() + mock_buffer.__iter__.return_value = [mock_event] + + # Mock the EventStreamBuffer class + with patch("botocore.eventstream.EventStreamBuffer", return_value=mock_buffer): + # Create an async generator that yields some test bytes + async def mock_iterator(): + yield b"" + + # Process the stream + chunks = [] + async for chunk in decoder.aiter_bytes(mock_iterator()): + if chunk is not None: + print("chunk=", chunk) + chunks.append(chunk) + + # Verify that processing continued despite the error + # The chunks list should be empty since we only sent invalid data + assert len(chunks) == 0 + + +@pytest.mark.asyncio +async def test_aiter_bytes_valid_chunk_followed_by_unicode_error(): + """ + Test that valid chunks are processed correctly even when followed by Unicode decode errors. + This ensures errors don't corrupt or prevent processing of valid data that came before. + """ + decoder = AWSEventStreamDecoder(model="test-model") + + # Create two mock events - first valid, then invalid + mock_valid_event = MagicMock() + mock_valid_event.to_response_dict.return_value = { + "status_code": 200, + "headers": {}, + "body": json.dumps({"token": {"text": "hello"}}).encode(), # Valid data first + } + + mock_invalid_event = MagicMock() + mock_invalid_event.to_response_dict.return_value = { + "status_code": 200, + "headers": {}, + "body": b"\xff\xfe", # Invalid UTF-8 bytes second + } + + # Create a mock EventStreamBuffer that yields valid event first, then invalid + mock_buffer = MagicMock() + mock_buffer.__iter__.return_value = [mock_valid_event, mock_invalid_event] + + with patch("botocore.eventstream.EventStreamBuffer", return_value=mock_buffer): + + async def mock_iterator(): + yield b"test_bytes" + + chunks = [] + async for chunk in decoder.aiter_bytes(mock_iterator()): + if chunk is not None: + chunks.append(chunk) + + # Verify we got our valid chunk despite the subsequent error + assert len(chunks) == 1 + assert chunks[0]["text"] == "hello" # Verify the content of the valid chunk From ca4ed9ff2efe00c103f9d427e209b74a0f90603f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 16:05:10 -0700 Subject: [PATCH 3/4] ref issue --- tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py b/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py index e19c41f1ac3..3b53107148b 100644 --- a/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py +++ b/tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py @@ -17,6 +17,8 @@ async def test_aiter_bytes_unicode_decode_error(): Ensures stream processing continues despite the error. + + Relevant issue: https://github.com/BerriAI/litellm/issues/9165 """ # Create an instance of AWSEventStreamDecoder decoder = AWSEventStreamDecoder(model="test-model") @@ -56,6 +58,8 @@ async def test_aiter_bytes_valid_chunk_followed_by_unicode_error(): """ Test that valid chunks are processed correctly even when followed by Unicode decode errors. This ensures errors don't corrupt or prevent processing of valid data that came before. + + Relevant issue: https://github.com/BerriAI/litellm/issues/9165 """ decoder = AWSEventStreamDecoder(model="test-model") From 3d0d5e1473bbf2e4fd0e6a6e0561439bc31ba8a0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 31 Mar 2025 16:12:31 -0700 Subject: [PATCH 4/4] run black formatting --- litellm/_redis.py | 2 +- litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/_redis.py b/litellm/_redis.py index 5c3e7776d6e..b2624d4280a 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -214,7 +214,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis: # Set up the Sentinel client sentinel = redis.Sentinel( - sentinel_nodes, + sentinel_nodes, socket_timeout=0.1, password=sentinel_password, ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index df9ae0ea570..9becb807584 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3,7 +3,7 @@ MCP Client Manager This class is responsible for managing MCP SSE clients. -This is a Proxy +This is a Proxy """ import asyncio