refactor(bedrock/sagemaker): switch to lazy loading for response stre…

litellm/llms/bedrock/chat/invoke_agent/transformation.py

@@ -299,9 +299,9 @@ def _extract_headers_from_event(self, event) -> InvokeAgentEventHeaders:
             )
 
     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."""

litellm/llms/bedrock/chat/invoke_handler.py

@@ -68,9 +68,9 @@
 
 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 @@ async def aiter_bytes(
                     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 @@ def _parse_message_from_event(self, event) -> Optional[str]:
                 ),
             )
         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()

litellm/llms/bedrock/common_utils.py

@@ -4,6 +4,7 @@
 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 @@ def __init__(self):
         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 @@ def _parse_message_from_event(self, event) -> Optional[str]:
                 ),
             )
         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()

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 @@ async def aiter_bytes(
                 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 @@ def _parse_message_from_event(self, event) -> Optional[str]:
                 ),
             )
         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}")

tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py

@@ -14,18 +14,46 @@
 
 
 # --------------------------------------------------------------------------- #
-# 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():
+    """
+    get_bedrock_response_stream_shape() loads from botocore at most once per process.
     """
-    BEDROCK_RESPONSE_STREAM_SHAPE is resolved at module import time.
+    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)