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

* 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.

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."""

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(
+        parsed_response = self.parser.parse(response_dict, response_stream_shape)
-            response_dict, BEDROCK_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 @@ 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.
+@functools.lru_cache(maxsize=1)
-BEDROCK_RESPONSE_STREAM_SHAPE = _load_bedrock_response_stream_shape()
+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(
+        parsed_response = self.parser.parse(response_dict, response_stream_shape)
-            response_dict, BEDROCK_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 @@ 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(
+        parsed_response = self.parser.parse(response_dict, response_stream_shape)
-            response_dict, SAGEMAKER_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 @@ 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, (
+    loaded_shape = get_bedrock_response_stream_shape()
-        "BEDROCK_RESPONSE_STREAM_SHAPE is None — botocore may not be installed"
+    assert (
-    )
+        loaded_shape is not None
-    shape: StructureShape = BEDROCK_RESPONSE_STREAM_SHAPE  # remove Optional
+    ), "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 —
+    Repeated calls must return the identical cached object.
-    confirming the constant is not re-loaded per import.
     """
-    from litellm.llms.bedrock.chat.invoke_handler import (
+    from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape
-        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
+    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)
 

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, (
+    shape = get_sagemaker_response_stream_shape()
-        "SAGEMAKER_RESPONSE_STREAM_SHAPE is None — botocore may not be installed"
+    assert (
-    )
+        shape is not None
-    shape: StructureShape = SAGEMAKER_RESPONSE_STREAM_SHAPE  # remove Optional
+    ), "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
+    additional botocore Loader calls — the shape is cached after first access.
-    time and reused.
     """
-    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_a = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder)
-    decoder_b = AWSEventStreamDecoder(model="test-model-b")
+    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),
+    When get_sagemaker_response_stream_shape() returns None (botocore unavailable),
-    _parse_message_from_event must raise ValueError before touching the
+    _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)