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.
2026-08-02 16:21:54 +00:00 · 2026-05-06 17:39:38 -07:00
parent 854456f58e
commit a67b7a7e87
8 changed files with 568 additions and 78 deletions
@@ -100,4 +100,5 @@ STABILIZATION_TODO.md
 **/playwright-report
 **/*.storageState.json
 **/coverage
-test-config
+test-config
+.vscode
@@ -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."""
@@ -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()
@@ -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:
@@ -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
@@ -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/<resolved model id>/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()
@@ -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
@@ -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():
    """