From a67b7a7e87f11bed01f9e073125a7f8f180105a2 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Wed, 6 May 2026 17:39:38 -0700 Subject: [PATCH] Refactor Bedrock response stream shape handling (#27257) * Refactor Bedrock response stream shape handling - Introduced a module-level constant `BEDROCK_RESPONSE_STREAM_SHAPE` to cache the response stream shape, eliminating the need for per-instance caching in `BedrockEventStreamDecoderBase`. - Updated relevant methods to utilize the new constant, improving performance by avoiding redundant loading of the shape. - Added tests to ensure the shape is loaded correctly at import time and is consistent across different modules. - Added a new mock server script for testing Bedrock pass-through functionality. * Refactor response parsing for Bedrock and SageMaker - Improved code readability by formatting the parsing method calls in `AWSEventStreamDecoder` for both Bedrock and SageMaker response stream shapes. - Added blank lines for better separation of code blocks in `invoke_handler.py` and `common_utils.py` to enhance maintainability. * Enhance error handling for Bedrock and SageMaker response stream shape loading - Wrapped the loading logic in `_load_bedrock_response_stream_shape` and `_load_sagemaker_response_stream_shape` with try-except blocks to gracefully handle exceptions. - Added logging to warn when the response stream shape cannot be pre-loaded, ensuring the module imports cleanly. - Updated tests to verify that loading failures return `None` instead of propagating exceptions. * Implement error handling for missing response stream shapes in Bedrock and SageMaker - Added checks in `_parse_message_from_event` methods to raise appropriate errors when `BEDROCK_RESPONSE_STREAM_SHAPE` or `SAGEMAKER_RESPONSE_STREAM_SHAPE` is None, ensuring clearer error reporting. - Updated logging messages to reflect the unavailability of event-stream decoding for both Bedrock and SageMaker. - Enhanced unit tests to verify that the correct exceptions are raised when the response stream shapes are not loaded. --- .gitignore | 3 +- .../chat/invoke_agent/transformation.py | 24 +- litellm/llms/bedrock/chat/invoke_handler.py | 34 +-- litellm/llms/bedrock/common_utils.py | 58 ++-- litellm/llms/sagemaker/common_utils.py | 51 ++-- scripts/mock_bedrock_passthrough_target.py | 276 ++++++++++++++++++ .../llms/bedrock/test_bedrock_common_utils.py | 104 +++++++ .../sagemaker/test_sagemaker_common_utils.py | 96 ++++++ 8 files changed, 568 insertions(+), 78 deletions(-) create mode 100644 scripts/mock_bedrock_passthrough_target.py diff --git a/.gitignore b/.gitignore index 59812ed6ed..20355a8e4e 100644 --- a/.gitignore +++ b/.gitignore @@ -100,4 +100,5 @@ STABILIZATION_TODO.md **/playwright-report **/*.storageState.json **/coverage -test-config \ No newline at end of file +test-config +.vscode \ No newline at end of file diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index 4c667b0ce3..e4072c2455 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -299,29 +299,9 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): ) def _get_response_stream_shape(self): - """Get the response stream shape for parsing, reusing existing logic.""" - try: - # Try to reuse the cached shape from the existing decoder - from litellm.llms.bedrock.chat.invoke_handler import ( - get_response_stream_shape, - ) + from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE - return get_response_stream_shape() - except ImportError: - # Fallback: create our own shape - try: - from botocore.loaders import Loader - from botocore.model import ServiceModel - - loader = Loader() - bedrock_service_dict = loader.load_service_model( - "bedrock-runtime", "service-2" - ) - bedrock_service_model = ServiceModel(bedrock_service_dict) - return bedrock_service_model.shape_for("ResponseStream") - except Exception as e: - verbose_logger.warning(f"Could not load response stream shape: {e}") - return None + return 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 9dfada7c41..92ca75db95 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -67,9 +67,13 @@ from litellm.types.utils import ( from litellm.utils import CustomStreamWrapper, get_secret from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name +from ..common_utils import ( + BEDROCK_RESPONSE_STREAM_SHAPE, + BedrockError, + ModelResponseIterator, + get_bedrock_tool_name, +) -_response_stream_shape_cache = None bedrock_tool_name_mappings: InMemoryCache = InMemoryCache( max_size_in_memory=50, default_ttl=600 ) @@ -1391,20 +1395,6 @@ class BedrockLLM(BaseAWSLLM): return None -def get_response_stream_shape(): - global _response_stream_shape_cache - if _response_stream_shape_cache is None: - from botocore.loaders import Loader - from botocore.model import ServiceModel - - loader = Loader() - bedrock_service_dict = loader.load_service_model("bedrock-runtime", "service-2") - bedrock_service_model = ServiceModel(bedrock_service_dict) - _response_stream_shape_cache = bedrock_service_model.shape_for("ResponseStream") - - return _response_stream_shape_cache - - class AWSEventStreamDecoder: def __init__(self, model: str, json_mode: Optional[bool] = False) -> None: from botocore.parsers import EventStreamJSONParser @@ -1838,8 +1828,18 @@ 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: + raise BedrockError( + status_code=500, + message=( + "Bedrock event-stream shape could not be loaded from botocore. " + "Ensure botocore is correctly installed." + ), + ) response_dict = event.to_response_dict() - parsed_response = self.parser.parse(response_dict, get_response_stream_shape()) + parsed_response = self.parser.parse( + response_dict, BEDROCK_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 9a97a134cc..856a525f77 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: import httpx import litellm +from litellm import verbose_logger from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -917,38 +918,57 @@ def get_bedrock_chat_config(model: str): return litellm.AmazonInvokeConfig() +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. + """ + try: + from botocore.loaders import Loader + from botocore.model import ServiceModel + + loader = Loader() + service_dict = loader.load_service_model("bedrock-runtime", "service-2") + return ServiceModel(service_dict).shape_for("ResponseStream") + except Exception as e: + verbose_logger.warning( + "litellm: could not pre-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() + + class BedrockEventStreamDecoderBase: """ Base class for event stream decoding for Bedrock """ - _response_stream_shape_cache = None - def __init__(self): from botocore.parsers import EventStreamJSONParser self.parser = EventStreamJSONParser() - def get_response_stream_shape(self): - if self._response_stream_shape_cache is None: - from botocore.loaders import Loader - from botocore.model import ServiceModel - - loader = Loader() - bedrock_service_dict = loader.load_service_model( - "bedrock-runtime", "service-2" - ) - bedrock_service_model = ServiceModel(bedrock_service_dict) - self._response_stream_shape_cache = bedrock_service_model.shape_for( - "ResponseStream" - ) - - return self._response_stream_shape_cache - def _parse_message_from_event(self, event) -> Optional[str]: + if BEDROCK_RESPONSE_STREAM_SHAPE is None: + raise BedrockError( + status_code=500, + message=( + "Bedrock event-stream shape could not be loaded from botocore. " + "Ensure botocore is correctly installed." + ), + ) response_dict = event.to_response_dict() parsed_response = self.parser.parse( - response_dict, self.get_response_stream_shape() + response_dict, BEDROCK_RESPONSE_STREAM_SHAPE ) if response_dict["status_code"] != 200: diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index ad6b24d85a..50c8ee4220 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -9,7 +9,27 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.types.utils import StreamingChatCompletionChunk -_response_stream_shape_cache = None + +def _load_sagemaker_response_stream_shape(): + try: + from botocore.loaders import Loader + from botocore.model import ServiceModel + + loader = Loader() + service_dict = loader.load_service_model("sagemaker-runtime", "service-2") + return ServiceModel(service_dict).shape_for( + "InvokeEndpointWithResponseStreamOutput" + ) + except Exception as e: + verbose_logger.warning( + "litellm: could not pre-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() class SagemakerError(BaseLLMException): @@ -187,8 +207,18 @@ 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: + raise SagemakerError( + status_code=500, + message=( + "SageMaker event-stream shape could not be loaded from botocore. " + "Ensure botocore is correctly installed." + ), + ) response_dict = event.to_response_dict() - parsed_response = self.parser.parse(response_dict, get_response_stream_shape()) + parsed_response = self.parser.parse( + response_dict, SAGEMAKER_RESPONSE_STREAM_SHAPE + ) if response_dict["status_code"] != 200: raise ValueError(f"Bad response code, expected 200: {response_dict}") @@ -204,20 +234,3 @@ class AWSEventStreamDecoder: return None return chunk.decode() # type: ignore[no-any-return] - - -def get_response_stream_shape(): - global _response_stream_shape_cache - if _response_stream_shape_cache is None: - from botocore.loaders import Loader - from botocore.model import ServiceModel - - loader = Loader() - sagemaker_service_dict = loader.load_service_model( - "sagemaker-runtime", "service-2" - ) - sagemaker_service_model = ServiceModel(sagemaker_service_dict) - _response_stream_shape_cache = sagemaker_service_model.shape_for( - "InvokeEndpointWithResponseStreamOutput" - ) - return _response_stream_shape_cache diff --git a/scripts/mock_bedrock_passthrough_target.py b/scripts/mock_bedrock_passthrough_target.py new file mode 100644 index 0000000000..e993cd99bd --- /dev/null +++ b/scripts/mock_bedrock_passthrough_target.py @@ -0,0 +1,276 @@ +#!/usr/bin/env python3 +""" +Minimal HTTP target for testing LiteLLM **Bedrock pass-through** (`/bedrock/...` on the proxy). + +What it does + - Serves a tiny Converse-shaped JSON (and optional invoke-shaped) response so the proxy can + complete a round trip without calling AWS. + - Does **not** verify SigV4 (Bedrock does); any Authorization header is accepted. + +How to run + uv run python scripts/mock_bedrock_passthrough_target.py --host 127.0.0.1 --port 9999 + +Wire LiteLLM to this host (use **one** of these patterns): + + 1) model_list (recommended) — set the Bedrock runtime base to the mock: + + model_list: + - model_name: mock-bedrock-claude + litellm_params: + model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0 + custom_llm_provider: bedrock + aws_region_name: us-west-2 + api_base: "http://127.0.0.1:9999" + + 2) Environment (see litellm BaseAWSLLM.get_runtime_endpoint):: + + export AWS_BEDROCK_RUNTIME_ENDPOINT="http://127.0.0.1:9999" + +Then call the proxy, e.g. (model_name must match config):: + + curl -sS -X POST "http://127.0.0.1:4000/bedrock/model/mock-bedrock-claude/converse" \ + -H "Authorization: Bearer $LITELLM_KEY" -H "Content-Type: application/json" \ + -d '{"messages":[{"role":"user","content":[{"text":"hi"}]}]}' + +The proxy will forward to: {api_base}/model//converse (SigV4-signed). +This mock implements POST .../converse and returns a minimal valid Converse response. + +Notes + - `invoke-with-response-stream` returns a real **binary** AWS event stream + (`application/vnd.amazon.eventstream`) with Anthropic-style JSON payloads inside each + `PayloadPart`, matching Bedrock's InvokeModelWithResponseStream wire format. See + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModelWithResponseStream.html + and https://docs.aws.amazon.com/awstreams/latest/devguide/message-formats.html + - `converse-stream` is still JSON-only placeholder (different inner event shapes). + - Use real (or any non-empty) AWS creds in the environment of the **proxy**; signing still runs. +""" +from __future__ import annotations + +import argparse +import base64 +import json +from binascii import crc32 +from struct import pack +from typing import Any, Dict, Iterator, List + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse +from starlette.responses import StreamingResponse + +app = FastAPI(title="Mock Bedrock runtime (pass-through test target)") + + +# Minimal structure compatible with Converse: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Converse.html +def _converse_response_body() -> Dict[str, Any]: + return { + "output": { + "message": { + "role": "assistant", + "content": [ + {"text": "mock: ok from mock_bedrock_passthrough_target.py"} + ], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 1, + "outputTokens": 2, + "totalTokens": 3, + }, + } + + +# Minimal invoke (Anthropic messages on bedrock) style — adjust if you test /invoke +def _invoke_response_body() -> Dict[str, Any]: + return { + "id": "msg_mock", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "mock invoke response"}], + "model": "mock", + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 2}, + } + + +def _encode_event_stream_message(headers: Dict[str, str], payload: bytes) -> bytes: + """Single AWS binary event-stream frame (same layout botocore's ``EventStreamBuffer`` parses).""" + header_blob = b"" + for name, value in headers.items(): + nb = name.encode("utf-8") + vb = value.encode("utf-8") + header_blob += bytes([len(nb)]) + nb + bytes([7]) + pack("!H", len(vb)) + vb + headers_length = len(header_blob) + payload_length = len(payload) + total_length = 12 + headers_length + payload_length + 4 + prelude_wo_crc = pack("!II", total_length, headers_length) + prelude_crc_val = crc32(prelude_wo_crc) & 0xFFFFFFFF + prelude = prelude_wo_crc + pack("!I", prelude_crc_val) + wo_msg_crc = prelude + header_blob + payload + msg_crc_val = crc32(wo_msg_crc[8:], prelude_crc_val) & 0xFFFFFFFF + return wo_msg_crc + pack("!I", msg_crc_val) + + +def _bedrock_payload_part(inner_event: Dict[str, Any]) -> bytes: + """Outer JSON expected by bedrock-runtime ``ResponseStream`` / ``PayloadPart``.""" + inner_bytes = json.dumps(inner_event, separators=(",", ":")).encode("utf-8") + outer = { + "chunk": { + "bytes": base64.b64encode(inner_bytes).decode("ascii"), + } + } + return json.dumps(outer, separators=(",", ":")).encode("utf-8") + + +def _anthropic_invoke_stream_events( + model_id: str, assistant_text: str +) -> List[Dict[str, Any]]: + """ + Minimal Anthropic Messages stream events as returned inside Bedrock stream chunks. + Mirrors the sequence Amazon emits for Claude on ``invoke-with-response-stream``. + """ + msg_id = "msg_mock_bedrock_stream" + input_tokens = 3 + output_tokens = max(1, len(assistant_text) // 4) + events: List[Dict[str, Any]] = [ + { + "type": "message_start", + "message": { + "model": model_id, + "id": msg_id, + "type": "message", + "role": "assistant", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": { + "input_tokens": input_tokens, + "output_tokens": 1, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 0, + }, + }, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ] + # Split text into small deltas so downstream streaming behavior is visible. + step = 24 + for i in range(0, len(assistant_text), step): + events.append( + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "text_delta", + "text": assistant_text[i : i + step], + }, + } + ) + events.append({"type": "content_block_stop", "index": 0}) + events.append( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, + } + ) + events.append( + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": input_tokens, + "outputTokenCount": output_tokens, + "invocationLatency": 42, + "firstByteLatency": 10, + }, + } + ) + return events + + +def _iter_invoke_with_response_stream(model_id: str) -> Iterator[bytes]: + text = ( + "mock streaming: ok from scripts/mock_bedrock_passthrough_target.py " + "(invoke-with-response-stream)." + ) + headers = { + ":event-type": "chunk", + ":content-type": "application/json", + ":message-type": "event", + } + for ev in _anthropic_invoke_stream_events(model_id, text): + yield _encode_event_stream_message(headers, _bedrock_payload_part(ev)) + + +@app.get("/health") +def health() -> Dict[str, str]: + return {"status": "ok"} + + +@app.post("/model/{model_path:path}/converse") +async def converse(model_path: str, request: Request) -> JSONResponse: + # Optional: log body for debugging + _ = await request.body() + return JSONResponse(content=_converse_response_body()) + + +@app.post("/model/{model_path:path}/converse-stream") +async def converse_stream(model_path: str, request: Request) -> JSONResponse: + """ + Not a real AWS event stream — returns JSON for quick smoke tests only. + """ + _ = await request.body() + return JSONResponse( + content={ + "note": "This mock does not implement application/vnd.amazon.eventstream; use /converse for basic tests." + } + ) + + +@app.post("/model/{model_path:path}/invoke") +async def invoke(model_path: str, request: Request) -> JSONResponse: + _ = await request.body() + return JSONResponse(content=_invoke_response_body()) + + +@app.post("/model/{model_path:path}/invoke-with-response-stream") +async def invoke_with_response_stream( + model_path: str, request: Request +) -> StreamingResponse: + """ + Binary ``application/vnd.amazon.eventstream`` body compatible with boto3/botocore + ``InvokeModelWithResponseStream`` / LiteLLM's Bedrock invoke streaming path. + """ + _ = await request.body() + return StreamingResponse( + _iter_invoke_with_response_stream(model_id=model_path), + media_type="application/vnd.amazon.eventstream", + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=9999) + args = parser.parse_args() + + import uvicorn + + uvicorn.run(app, host=args.host, port=args.port, log_level="info") + + +if __name__ == "__main__": + main() 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 c356f866b0..8fa9290d3d 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -13,6 +13,110 @@ sys.path.insert( from litellm.llms.bedrock.common_utils import BedrockModelInfo +# --------------------------------------------------------------------------- # +# BEDROCK_RESPONSE_STREAM_SHAPE eager-load tests # +# --------------------------------------------------------------------------- # + + +def test_bedrock_response_stream_shape_loaded_at_import(): + """ + BEDROCK_RESPONSE_STREAM_SHAPE is resolved at module import time. + In a standard environment with botocore installed it must be non-None. + """ + from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE + + assert BEDROCK_RESPONSE_STREAM_SHAPE is not None + + +def test_bedrock_response_stream_shape_load_failure_returns_none(): + """ + If botocore's Loader raises (e.g. missing data files), _load_bedrock_response_stream_shape + should return None rather than propagating the exception, so the module + still imports cleanly. + """ + from unittest.mock import patch + + import litellm.llms.bedrock.common_utils as mod + + with patch( + "botocore.loaders.Loader.load_service_model", + side_effect=Exception("no data"), + ): + shape = mod._load_bedrock_response_stream_shape() + assert shape is None + + +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. + """ + from botocore.model import StructureShape + + from litellm.llms.bedrock.common_utils import 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 + assert isinstance(shape, StructureShape) + assert shape.name == "ResponseStream" + + +def test_bedrock_response_stream_shape_same_object_across_imports(): + """ + Both bedrock modules that use the shape must reference the identical object — + confirming the constant is not re-loaded per import. + """ + 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, + ) + + assert common_shape is invoke_shape + + +def test_bedrock_event_stream_decoder_base_uses_module_shape(): + """ + BedrockEventStreamDecoderBase instances no longer carry their own + per-instance cache — _parse_message_from_event uses the module constant + directly, so there is no instance-level _response_stream_shape_cache attr. + """ + from litellm.llms.bedrock.common_utils import BedrockEventStreamDecoderBase + + decoder_a = BedrockEventStreamDecoderBase() + decoder_b = BedrockEventStreamDecoderBase() + + assert "_response_stream_shape_cache" not in decoder_a.__dict__ + assert "_response_stream_shape_cache" not in decoder_b.__dict__ + + +def test_bedrock_parse_message_from_event_raises_on_none_shape(): + """ + When BEDROCK_RESPONSE_STREAM_SHAPE is 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 + + decoder = BedrockEventStreamDecoderBase() + mock_event = MagicMock() + + with patch.object(mod, "BEDROCK_RESPONSE_STREAM_SHAPE", None): + with pytest.raises(BedrockError) as exc_info: + decoder._parse_message_from_event(mock_event) + + assert exc_info.value.status_code == 500 + assert "botocore" in str(exc_info.value.message).lower() + # The botocore parser must never have been called + mock_event.to_response_dict.assert_not_called() + + def test_deepseek_cris(): """ Test that DeepSeek models with cross-region inference prefix use converse route 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 70a8d86cb1..9d7706557b 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py @@ -11,6 +11,102 @@ from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder from litellm.llms.sagemaker.completion.transformation import SagemakerConfig +# --------------------------------------------------------------------------- # +# SAGEMAKER_RESPONSE_STREAM_SHAPE eager-load tests # +# --------------------------------------------------------------------------- # + + +def test_sagemaker_response_stream_shape_loaded_at_import(): + """ + SAGEMAKER_RESPONSE_STREAM_SHAPE is resolved at module import time. + In a standard environment with botocore installed it must be non-None. + """ + from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE + + assert SAGEMAKER_RESPONSE_STREAM_SHAPE is not None + + +def test_sagemaker_response_stream_shape_load_failure_returns_none(): + """ + If botocore's Loader raises (e.g. missing data files), _load_sagemaker_response_stream_shape + should return None rather than propagating the exception, so the module + still imports cleanly. + """ + from unittest.mock import patch + + import litellm.llms.sagemaker.common_utils as mod + + with patch( + "botocore.loaders.Loader.load_service_model", + side_effect=Exception("no data"), + ): + shape = mod._load_sagemaker_response_stream_shape() + assert shape is None + + +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. + """ + from botocore.model import StructureShape + + from litellm.llms.sagemaker.common_utils import 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 + assert isinstance(shape, StructureShape) + assert shape.name == "InvokeEndpointWithResponseStreamOutput" + + +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. + """ + from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE + + decoder_a = AWSEventStreamDecoder(model="test-model-a") + decoder_b = AWSEventStreamDecoder(model="test-model-b") + + # 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 + + +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 + botocore parser — not an opaque AttributeError from inside botocore. + """ + from unittest.mock import MagicMock, patch + + import litellm.llms.sagemaker.common_utils as mod + from litellm.llms.sagemaker.common_utils import SagemakerError + + decoder = AWSEventStreamDecoder(model="test-model") + mock_event = MagicMock() + + with patch.object(mod, "SAGEMAKER_RESPONSE_STREAM_SHAPE", None): + with pytest.raises(SagemakerError) as exc_info: + decoder._parse_message_from_event(mock_event) + + assert exc_info.value.status_code == 500 + assert "botocore" in str(exc_info.value.message).lower() + # The botocore parser must never have been called + mock_event.to_response_dict.assert_not_called() + + @pytest.mark.asyncio async def test_aiter_bytes_unicode_decode_error(): """