From 0290c7bc00102e13de520d925609ae82c274eb79 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 19 May 2026 11:16:11 +0530 Subject: [PATCH 1/2] fix(proxy): decode bytes and pass-through SSE for Google-native streamGenerateContent (#27444) (#28213) * fix(proxy): decode bytes and pass-through SSE for Google-native streamGenerateContent (#27444) * fix(proxy): address Greptile review on Google-native SSE bytes path Remove unreachable try/except around SSE pass-through yield and add a unit test covering pre-formatted SSE bytes, terminator padding, and non-SSE byte fallback wrapping. Co-authored-by: Cursor --------- Co-authored-by: Tai An Co-authored-by: Cursor --- litellm/proxy/proxy_server.py | 9 +++ tests/test_litellm/proxy/test_proxy_server.py | 60 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5d89d3fa9c5..879914e5ac6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6917,6 +6917,15 @@ async def async_data_generator( # noqa: PLR0915 if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) + elif isinstance(chunk, bytes): + # Some upstream streaming iterators (e.g. AsyncGoogleGenAIGenerateContentStreamingIterator + # for /v1beta/.../streamGenerateContent) yield raw SSE bytes from Gemini. + # Decode to str so the f-string below does not emit a Python b'...' literal, + # and pass already-formatted SSE through unchanged to avoid double "data:" prefix. + chunk = chunk.decode("utf-8", errors="replace") + if chunk.startswith(("data:", "event:", ":")): + yield chunk if chunk.endswith("\n\n") else chunk + "\n\n" + continue elif isinstance(chunk, str) and chunk.startswith("data: "): error_message = chunk break diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 73d53631622..6d10d2a6353 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5065,6 +5065,66 @@ async def test_async_data_generator_uses_direct_stream_fast_path_without_callbac mock_response.aclose.assert_awaited_once() +@pytest.mark.asyncio +async def test_async_data_generator_passes_through_google_native_sse_bytes(): + """ + Google-native streamGenerateContent yields raw SSE bytes; they must not be + re-wrapped as data: b'data: {...}'. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import async_data_generator + from litellm.proxy.utils import ProxyLogging + + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_request_data = { + "model": "gemini-2.0-flash", + "messages": [{"role": "user", "content": "test"}], + } + gemini_event = b'data: {"candidates": [{"content": "hi"}]}\n\n' + gemini_event_without_terminator = b'data: {"candidates": [{"content": "there"}]}' + raw_payload = b'{"partial": true}' + + class MockStream: + def __aiter__(self): + return self._stream() + + async def _stream(self): + yield gemini_event + yield gemini_event_without_terminator + yield raw_payload + + async def aclose(self): + pass + + mock_response = MockStream() + mock_response.aclose = AsyncMock() + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.has_streaming_callbacks.return_value = False + mock_proxy_logging_obj.needs_iterator_wrap.return_value = False + mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock() + mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): + with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): + yielded_data = [] + async for data in async_data_generator( + mock_response, mock_user_api_key_dict, mock_request_data + ): + yielded_data.append(data) + + yielded_text = [ + chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk + for chunk in yielded_data + ] + assert yielded_text[0] == gemini_event.decode("utf-8") + assert yielded_text[1] == gemini_event_without_terminator.decode("utf-8") + "\n\n" + assert yielded_text[2] == f'data: {raw_payload.decode("utf-8")}\n\n' + assert "b'data:" not in "".join(yielded_text) + assert yielded_text[-1] == "data: [DONE]\n\n" + + @pytest.mark.asyncio async def test_async_data_generator_cleanup_on_normal_completion(): """ From cff3e0b75eeb836e9fafbcbf73fd4c968b001def Mon Sep 17 00:00:00 2001 From: harish-berri Date: Mon, 18 May 2026 23:21:04 -0700 Subject: [PATCH 2/2] =?UTF-8?q?refactor(bedrock/sagemaker):=20switch=20to?= =?UTF-8?q?=20lazy=20loading=20for=20response=20stre=E2=80=A6=20(#28189)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(bedrock/sagemaker): switch to lazy loading for response stream shapes - Replace eager loading of BEDROCK_RESPONSE_STREAM_SHAPE and SAGEMAKER_RESPONSE_STREAM_SHAPE with lazy loading via get_bedrock_response_stream_shape() and get_sagemaker_response_stream_shape() respectively. - This change optimizes performance by avoiding unnecessary imports and logging warnings unless the response stream shapes are actually needed. - Update relevant classes and tests to utilize the new lazy loading functions, ensuring consistent behavior across the codebase. * test(bedrock/sagemaker): add fixtures to clear response stream shape cache - Introduced `_reset_bedrock_response_stream_shape_cache` and `_reset_sagemaker_response_stream_shape_cache` fixtures to prevent lru_cache leakage between tests in their respective modules. - Updated tests to utilize these fixtures, ensuring that the response stream shape cache is cleared before and after each test run. - Added `pytest.importorskip("botocore")` to ensure that tests are skipped if the botocore library is not available. --- .../chat/invoke_agent/transformation.py | 4 +- litellm/llms/bedrock/chat/invoke_handler.py | 9 +-- litellm/llms/bedrock/common_utils.py | 25 +++--- litellm/llms/sagemaker/common_utils.py | 20 +++-- .../llms/bedrock/test_bedrock_common_utils.py | 79 ++++++++++++------ .../sagemaker/test_sagemaker_common_utils.py | 81 +++++++++++++------ 6 files changed, 146 insertions(+), 72 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index e4072c24557..c88fa32b6a0 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -299,9 +299,9 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): ) def _get_response_stream_shape(self): - from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE + from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape - return BEDROCK_RESPONSE_STREAM_SHAPE + return get_bedrock_response_stream_shape() def _extract_response_content(self, events: InvokeAgentEventList) -> str: """Extract the final response content from parsed events.""" diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 92ca75db95b..7a9916f1f31 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -68,9 +68,9 @@ from litellm.utils import CustomStreamWrapper, get_secret from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( - BEDROCK_RESPONSE_STREAM_SHAPE, BedrockError, ModelResponseIterator, + get_bedrock_response_stream_shape, get_bedrock_tool_name, ) @@ -1828,7 +1828,8 @@ class AWSEventStreamDecoder: yield self._chunk_parser(chunk_data=_data) def _parse_message_from_event(self, event) -> Optional[str]: - if BEDROCK_RESPONSE_STREAM_SHAPE is None: + response_stream_shape = get_bedrock_response_stream_shape() + if response_stream_shape is None: raise BedrockError( status_code=500, message=( @@ -1837,9 +1838,7 @@ class AWSEventStreamDecoder: ), ) response_dict = event.to_response_dict() - parsed_response = self.parser.parse( - response_dict, BEDROCK_RESPONSE_STREAM_SHAPE - ) + parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: decoded_body = response_dict["body"].decode() diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 0256d5d4b95..4f4729e4019 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -4,6 +4,7 @@ from __future__ import annotations Common utilities used across bedrock chat/embedding/image generation """ +import functools import json import os from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union @@ -963,10 +964,8 @@ def _load_bedrock_response_stream_shape(): """ Load the ResponseStream shape from botocore's bundled bedrock-runtime schema. - Called once at module import time; the result is stored in - ``BEDROCK_RESPONSE_STREAM_SHAPE`` and reused for the process lifetime. Returns ``None`` if botocore is unavailable or the service model cannot be - loaded, so the module still imports cleanly. + loaded. """ try: from botocore.loaders import Loader @@ -977,15 +976,22 @@ def _load_bedrock_response_stream_shape(): return ServiceModel(service_dict).shape_for("ResponseStream") except Exception as e: verbose_logger.warning( - "litellm: could not pre-load bedrock-runtime response stream shape " + "litellm: could not load bedrock-runtime response stream shape " "— Bedrock event-stream decoding will be unavailable. Error: %s", e, ) return None -# Eagerly resolved once per process — avoids per-instance or per-request disk I/O. -BEDROCK_RESPONSE_STREAM_SHAPE = _load_bedrock_response_stream_shape() +@functools.lru_cache(maxsize=1) +def get_bedrock_response_stream_shape(): + """ + Lazily load and cache the bedrock-runtime ResponseStream shape for the process. + + Avoids importing botocore (and logging warnings) unless Bedrock event-stream + decoding is actually needed. + """ + return _load_bedrock_response_stream_shape() class BedrockEventStreamDecoderBase: @@ -999,7 +1005,8 @@ class BedrockEventStreamDecoderBase: self.parser = EventStreamJSONParser() def _parse_message_from_event(self, event) -> Optional[str]: - if BEDROCK_RESPONSE_STREAM_SHAPE is None: + response_stream_shape = get_bedrock_response_stream_shape() + if response_stream_shape is None: raise BedrockError( status_code=500, message=( @@ -1008,9 +1015,7 @@ class BedrockEventStreamDecoderBase: ), ) response_dict = event.to_response_dict() - parsed_response = self.parser.parse( - response_dict, BEDROCK_RESPONSE_STREAM_SHAPE - ) + parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: decoded_body = response_dict["body"].decode() diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 50c8ee4220e..6c15d642f8c 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -1,3 +1,4 @@ +import functools import json from typing import AsyncIterator, Iterator, List, Optional, Union @@ -22,14 +23,22 @@ def _load_sagemaker_response_stream_shape(): ) except Exception as e: verbose_logger.warning( - "litellm: could not pre-load sagemaker-runtime response stream shape " + "litellm: could not load sagemaker-runtime response stream shape " "— SageMaker event-stream decoding will be unavailable. Error: %s", e, ) return None -SAGEMAKER_RESPONSE_STREAM_SHAPE = _load_sagemaker_response_stream_shape() +@functools.lru_cache(maxsize=1) +def get_sagemaker_response_stream_shape(): + """ + Lazily load and cache the sagemaker-runtime stream shape for the process. + + Avoids importing botocore (and logging warnings) unless SageMaker event-stream + decoding is actually needed. + """ + return _load_sagemaker_response_stream_shape() class SagemakerError(BaseLLMException): @@ -207,7 +216,8 @@ class AWSEventStreamDecoder: verbose_logger.error(f"Final error parsing accumulated JSON: {e}") def _parse_message_from_event(self, event) -> Optional[str]: - if SAGEMAKER_RESPONSE_STREAM_SHAPE is None: + response_stream_shape = get_sagemaker_response_stream_shape() + if response_stream_shape is None: raise SagemakerError( status_code=500, message=( @@ -216,9 +226,7 @@ class AWSEventStreamDecoder: ), ) response_dict = event.to_response_dict() - parsed_response = self.parser.parse( - response_dict, SAGEMAKER_RESPONSE_STREAM_SHAPE - ) + parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: raise ValueError(f"Bad response code, expected 200: {response_dict}") diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 8fa9290d3de..c39fb427a01 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -14,18 +14,46 @@ from litellm.llms.bedrock.common_utils import BedrockModelInfo # --------------------------------------------------------------------------- # -# BEDROCK_RESPONSE_STREAM_SHAPE eager-load tests # +# get_bedrock_response_stream_shape lazy-load tests # # --------------------------------------------------------------------------- # -def test_bedrock_response_stream_shape_loaded_at_import(): +@pytest.fixture(autouse=True) +def _reset_bedrock_response_stream_shape_cache(): + """Prevent lru_cache leakage between tests in this module.""" + import litellm.llms.bedrock.common_utils as mod + + mod.get_bedrock_response_stream_shape.cache_clear() + yield + mod.get_bedrock_response_stream_shape.cache_clear() + + +def test_bedrock_response_stream_shape_lazy_loads_once(): """ - BEDROCK_RESPONSE_STREAM_SHAPE is resolved at module import time. + get_bedrock_response_stream_shape() loads from botocore at most once per process. + """ + from unittest.mock import MagicMock, patch + + import litellm.llms.bedrock.common_utils as mod + + sentinel = MagicMock() + with patch.object( + mod, "_load_bedrock_response_stream_shape", return_value=sentinel + ) as mock_load: + assert mod.get_bedrock_response_stream_shape() is sentinel + assert mod.get_bedrock_response_stream_shape() is sentinel + mock_load.assert_called_once() + + +def test_bedrock_response_stream_shape_loaded_on_first_access(): + """ + get_bedrock_response_stream_shape() loads once on first use. In a standard environment with botocore installed it must be non-None. """ - from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE + pytest.importorskip("botocore") + from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape - assert BEDROCK_RESPONSE_STREAM_SHAPE is not None + assert get_bedrock_response_stream_shape() is not None def test_bedrock_response_stream_shape_load_failure_returns_none(): @@ -38,6 +66,7 @@ def test_bedrock_response_stream_shape_load_failure_returns_none(): import litellm.llms.bedrock.common_utils as mod + pytest.importorskip("botocore") with patch( "botocore.loaders.Loader.load_service_model", side_effect=Exception("no data"), @@ -51,31 +80,29 @@ def test_bedrock_response_stream_shape_is_structure_shape(): The loaded shape should be the botocore StructureShape for ResponseStream, not a plain dict or any other type. """ + pytest.importorskip("botocore") from botocore.model import StructureShape - from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE + from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape - assert BEDROCK_RESPONSE_STREAM_SHAPE is not None, ( - "BEDROCK_RESPONSE_STREAM_SHAPE is None — botocore may not be installed" - ) - shape: StructureShape = BEDROCK_RESPONSE_STREAM_SHAPE # remove Optional + loaded_shape = get_bedrock_response_stream_shape() + assert ( + loaded_shape is not None + ), "get_bedrock_response_stream_shape() is None — botocore may not be installed" + shape: StructureShape = loaded_shape assert isinstance(shape, StructureShape) assert shape.name == "ResponseStream" -def test_bedrock_response_stream_shape_same_object_across_imports(): +def test_bedrock_response_stream_shape_same_object_across_calls(): """ - Both bedrock modules that use the shape must reference the identical object — - confirming the constant is not re-loaded per import. + Repeated calls must return the identical cached object. """ - from litellm.llms.bedrock.chat.invoke_handler import ( - BEDROCK_RESPONSE_STREAM_SHAPE as invoke_shape, - ) - from litellm.llms.bedrock.common_utils import ( - BEDROCK_RESPONSE_STREAM_SHAPE as common_shape, - ) + from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape - assert common_shape is invoke_shape + first = get_bedrock_response_stream_shape() + second = get_bedrock_response_stream_shape() + assert first is second def test_bedrock_event_stream_decoder_base_uses_module_shape(): @@ -95,19 +122,23 @@ def test_bedrock_event_stream_decoder_base_uses_module_shape(): def test_bedrock_parse_message_from_event_raises_on_none_shape(): """ - When BEDROCK_RESPONSE_STREAM_SHAPE is None (botocore unavailable), + When get_bedrock_response_stream_shape() returns None (botocore unavailable), _parse_message_from_event must raise BedrockError before touching the botocore parser — not an opaque AttributeError from inside botocore. """ from unittest.mock import MagicMock, patch import litellm.llms.bedrock.common_utils as mod - from litellm.llms.bedrock.common_utils import BedrockError, BedrockEventStreamDecoderBase + from litellm.llms.bedrock.common_utils import ( + BedrockError, + BedrockEventStreamDecoderBase, + ) - decoder = BedrockEventStreamDecoderBase() + decoder = BedrockEventStreamDecoderBase.__new__(BedrockEventStreamDecoderBase) + decoder.parser = MagicMock() mock_event = MagicMock() - with patch.object(mod, "BEDROCK_RESPONSE_STREAM_SHAPE", None): + with patch.object(mod, "get_bedrock_response_stream_shape", return_value=None): with pytest.raises(BedrockError) as exc_info: decoder._parse_message_from_event(mock_event) diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py index 9d7706557b5..7e13459bca1 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py @@ -12,18 +12,46 @@ from litellm.llms.sagemaker.completion.transformation import SagemakerConfig # --------------------------------------------------------------------------- # -# SAGEMAKER_RESPONSE_STREAM_SHAPE eager-load tests # +# get_sagemaker_response_stream_shape lazy-load tests # # --------------------------------------------------------------------------- # -def test_sagemaker_response_stream_shape_loaded_at_import(): +@pytest.fixture(autouse=True) +def _reset_sagemaker_response_stream_shape_cache(): + """Prevent lru_cache leakage between tests in this module.""" + import litellm.llms.sagemaker.common_utils as mod + + mod.get_sagemaker_response_stream_shape.cache_clear() + yield + mod.get_sagemaker_response_stream_shape.cache_clear() + + +def test_sagemaker_response_stream_shape_lazy_loads_once(): """ - SAGEMAKER_RESPONSE_STREAM_SHAPE is resolved at module import time. + get_sagemaker_response_stream_shape() loads from botocore at most once per process. + """ + from unittest.mock import MagicMock, patch + + import litellm.llms.sagemaker.common_utils as mod + + sentinel = MagicMock() + with patch.object( + mod, "_load_sagemaker_response_stream_shape", return_value=sentinel + ) as mock_load: + assert mod.get_sagemaker_response_stream_shape() is sentinel + assert mod.get_sagemaker_response_stream_shape() is sentinel + mock_load.assert_called_once() + + +def test_sagemaker_response_stream_shape_loaded_on_first_access(): + """ + get_sagemaker_response_stream_shape() loads once on first use. In a standard environment with botocore installed it must be non-None. """ - from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE + pytest.importorskip("botocore") + from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape - assert SAGEMAKER_RESPONSE_STREAM_SHAPE is not None + assert get_sagemaker_response_stream_shape() is not None def test_sagemaker_response_stream_shape_load_failure_returns_none(): @@ -36,6 +64,7 @@ def test_sagemaker_response_stream_shape_load_failure_returns_none(): import litellm.llms.sagemaker.common_utils as mod + pytest.importorskip("botocore") with patch( "botocore.loaders.Loader.load_service_model", side_effect=Exception("no data"), @@ -49,14 +78,16 @@ def test_sagemaker_response_stream_shape_is_structure_shape(): The loaded shape should be the botocore StructureShape for InvokeEndpointWithResponseStreamOutput, not a plain dict or any other type. """ + pytest.importorskip("botocore") from botocore.model import StructureShape - from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE + from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape - assert SAGEMAKER_RESPONSE_STREAM_SHAPE is not None, ( - "SAGEMAKER_RESPONSE_STREAM_SHAPE is None — botocore may not be installed" - ) - shape: StructureShape = SAGEMAKER_RESPONSE_STREAM_SHAPE # remove Optional + shape = get_sagemaker_response_stream_shape() + assert ( + shape is not None + ), "get_sagemaker_response_stream_shape() is None — botocore may not be installed" + shape: StructureShape = shape # remove Optional assert isinstance(shape, StructureShape) assert shape.name == "InvokeEndpointWithResponseStreamOutput" @@ -64,29 +95,25 @@ def test_sagemaker_response_stream_shape_is_structure_shape(): def test_sagemaker_response_stream_shape_not_reloaded_on_new_decoder(): """ Creating multiple AWSEventStreamDecoder instances must not trigger - additional botocore Loader calls — the shape is resolved once at import - time and reused. + additional botocore Loader calls — the shape is cached after first access. """ - from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE + from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape - decoder_a = AWSEventStreamDecoder(model="test-model-a") - decoder_b = AWSEventStreamDecoder(model="test-model-b") + decoder_a = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder) + decoder_b = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder) - # Both decoders should use the same pre-loaded shape object (identity check) assert "_response_stream_shape_cache" not in decoder_a.__dict__ assert "_response_stream_shape_cache" not in decoder_b.__dict__ - # The module constant is still the same object - from litellm.llms.sagemaker.common_utils import ( - SAGEMAKER_RESPONSE_STREAM_SHAPE as shape_after, - ) - assert SAGEMAKER_RESPONSE_STREAM_SHAPE is shape_after + first = get_sagemaker_response_stream_shape() + second = get_sagemaker_response_stream_shape() + assert first is second def test_sagemaker_parse_message_from_event_raises_on_none_shape(): """ - When SAGEMAKER_RESPONSE_STREAM_SHAPE is None (botocore unavailable), - _parse_message_from_event must raise ValueError before touching the + When get_sagemaker_response_stream_shape() returns None (botocore unavailable), + _parse_message_from_event must raise SagemakerError before touching the botocore parser — not an opaque AttributeError from inside botocore. """ from unittest.mock import MagicMock, patch @@ -94,10 +121,14 @@ def test_sagemaker_parse_message_from_event_raises_on_none_shape(): import litellm.llms.sagemaker.common_utils as mod from litellm.llms.sagemaker.common_utils import SagemakerError - decoder = AWSEventStreamDecoder(model="test-model") + decoder = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder) + decoder.model = "test-model" + decoder.parser = MagicMock() + decoder.content_blocks = [] + decoder.is_messages_api = None mock_event = MagicMock() - with patch.object(mod, "SAGEMAKER_RESPONSE_STREAM_SHAPE", None): + with patch.object(mod, "get_sagemaker_response_stream_shape", return_value=None): with pytest.raises(SagemakerError) as exc_info: decoder._parse_message_from_event(mock_event)