refactor: lazy load bedrock types (#18053)

litellm/__init__.py

@@ -26,7 +26,6 @@ from typing import (
 )
 from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams
 from litellm.types.integrations.datadog import DatadogInitParams
-from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES
 from litellm.types.utils import (
     ImageObject,
     BudgetConfig,
@@ -292,7 +291,7 @@ AZURE_DEFAULT_API_VERSION = "2025-02-01-preview"  # this is updated to the lates
 ### DEFAULT WATSONX API VERSION ###
 WATSONX_DEFAULT_API_VERSION = "2024-03-13"
 ### COHERE EMBEDDINGS DEFAULT TYPE ###
-COHERE_DEFAULT_EMBEDDING_INPUT_TYPE: COHERE_EMBEDDING_INPUT_TYPES = "search_document"
+COHERE_DEFAULT_EMBEDDING_INPUT_TYPE: "COHERE_EMBEDDING_INPUT_TYPES" = "search_document"
 ### CREDENTIALS ###
 credential_list: List[CredentialItem] = []
 ### GUARDRAILS ###
@@ -1515,6 +1514,7 @@ if TYPE_CHECKING:
     from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
     from litellm.caching.caching import Cache
     from litellm.caching.llm_caching_handler import LLMClientCache
+    from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES
 
     # Cost calculator functions
     cost_per_token: Callable[..., Tuple[float, float]]
@@ -1568,6 +1568,7 @@ def __getattr__(name: str) -> Any:
         UTILS_NAMES,
         TOKEN_COUNTER_NAMES,
         LLM_CLIENT_CACHE_NAMES,
+        BEDROCK_TYPES_NAMES,
         CACHING_NAMES,
         HTTP_HANDLER_NAMES,
     )
@@ -1592,6 +1593,11 @@ def __getattr__(name: str) -> Any:
         from ._lazy_imports import _lazy_import_token_counter
         return _lazy_import_token_counter(name)
     
+    # Lazy load Bedrock type aliases
+    if name in BEDROCK_TYPES_NAMES:
+        from ._lazy_imports import _lazy_import_bedrock_types
+        return _lazy_import_bedrock_types(name)
+    
     # Lazy load LLM client cache and its singleton
     if name in LLM_CLIENT_CACHE_NAMES:
         from ._lazy_imports import _lazy_import_llm_client_cache

litellm/_lazy_imports.py

@@ -45,6 +45,11 @@ LLM_CLIENT_CACHE_NAMES = (
     "in_memory_llm_clients_cache",
 )
 
+# Bedrock type names that support lazy loading via _lazy_import_bedrock_types
+BEDROCK_TYPES_NAMES = (
+    "COHERE_EMBEDDING_INPUT_TYPES",
+)
+
 # Caching / cache classes that support lazy loading via _lazy_import_caching
 CACHING_NAMES = (
     "Cache",
@@ -305,6 +310,21 @@ def _lazy_import_token_counter(name: str) -> Any:
     raise AttributeError(f"Token counter lazy import: unknown attribute {name!r}")
 
 
+def _lazy_import_bedrock_types(name: str) -> Any:
+    """Lazy import for Bedrock type aliases."""
+    _globals = _get_litellm_globals()
+
+    if name == "COHERE_EMBEDDING_INPUT_TYPES":
+        from litellm.types.llms.bedrock import (
+            COHERE_EMBEDDING_INPUT_TYPES as _COHERE_EMBEDDING_INPUT_TYPES,
+        )
+
+        _globals["COHERE_EMBEDDING_INPUT_TYPES"] = _COHERE_EMBEDDING_INPUT_TYPES
+        return _COHERE_EMBEDDING_INPUT_TYPES
+
+    raise AttributeError(f"Bedrock types lazy import: unknown attribute {name!r}")
+
+
 def _lazy_import_caching(name: str) -> Any:
     """Lazy import for caching module classes."""
     _globals = _get_litellm_globals()

tests/test_litellm/test_lazy_imports.py

@@ -14,12 +14,14 @@ from litellm._lazy_imports import (
     UTILS_NAMES,
     TOKEN_COUNTER_NAMES,
     CACHING_NAMES,
+    BEDROCK_TYPES_NAMES,
     LLM_CLIENT_CACHE_NAMES,
     HTTP_HANDLER_NAMES,
     _lazy_import_cost_calculator,
     _lazy_import_litellm_logging,
     _lazy_import_utils,
     _lazy_import_token_counter,
+    _lazy_import_bedrock_types,
     _lazy_import_caching,
     _lazy_import_llm_client_cache,
     _lazy_import_http_handlers,
@@ -113,6 +115,18 @@ def test_token_counter_lazy_imports():
         _verify_only_requested_name_imported(name, TOKEN_COUNTER_NAMES)
 
 
+def test_bedrock_types_lazy_imports():
+    """Test that Bedrock type aliases can be lazy imported."""
+    for name in BEDROCK_TYPES_NAMES:
+        _clear_names_from_globals(BEDROCK_TYPES_NAMES)
+
+        alias = _lazy_import_bedrock_types(name)
+        assert alias is not None
+        assert name in litellm.__dict__
+
+        _verify_only_requested_name_imported(name, BEDROCK_TYPES_NAMES)
+
+
 def test_llm_client_cache_lazy_imports():
     """Test that LLM client cache class and singleton can be lazy imported."""
     for name in LLM_CLIENT_CACHE_NAMES: