diff --git a/litellm/__init__.py b/litellm/__init__.py index 389ee4d479..f88b6bb3ce 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index a639e5b5b5..db7c90d7d8 100644 --- a/litellm/_lazy_imports.py +++ b/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() diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index 15eb7a5be6..c137e6ddbc 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/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: