From feb6924efe5ff86f83dd88e7ae4411d40ef2398e Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 16 Dec 2025 06:42:17 -0800 Subject: [PATCH] [Refactor] litellm/init.py: lazy load .types.utils (#18054) --- litellm/__init__.py | 37 +++++------ litellm/_lazy_imports.py | 81 +++++++++++++++++++++++++ tests/test_litellm/test_lazy_imports.py | 20 ++++++ 3 files changed, 120 insertions(+), 18 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index f88b6bb3ce..1628e42a33 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -26,14 +26,6 @@ from typing import ( ) from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams from litellm.types.integrations.datadog import DatadogInitParams -from litellm.types.utils import ( - ImageObject, - BudgetConfig, - all_litellm_params, - all_litellm_params as _litellm_completion_params, - CredentialItem, - PriorityReservationDict, -) # maintain backwards compatibility for root param. from litellm._logging import ( set_verbose, _turn_on_debug, @@ -94,11 +86,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, LiteLLM_UpperboundKeyGenerateParams, ) -from litellm.types.utils import ( - StandardKeyGenerationConfig, - LlmProviders, - SearchProviders, -) +from litellm.types.utils import LlmProviders from litellm.types.utils import PriorityReservationSettings from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager @@ -293,7 +281,7 @@ WATSONX_DEFAULT_API_VERSION = "2024-03-13" ### COHERE EMBEDDINGS DEFAULT TYPE ### COHERE_DEFAULT_EMBEDDING_INPUT_TYPE: "COHERE_EMBEDDING_INPUT_TYPES" = "search_document" ### CREDENTIALS ### -credential_list: List[CredentialItem] = [] +credential_list: List["CredentialItem"] = [] ### GUARDRAILS ### llamaguard_model_name: Optional[str] = None openai_moderations_model_name: Optional[str] = None @@ -368,7 +356,7 @@ aws_sqs_callback_params: Optional[Dict] = None generic_logger_headers: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None upperbound_key_generate_params: Optional[LiteLLM_UpperboundKeyGenerateParams] = None -key_generation_settings: Optional[StandardKeyGenerationConfig] = None +key_generation_settings: Optional["StandardKeyGenerationConfig"] = None default_internal_user_params: Optional[Dict] = None default_team_params: Optional[Union[DefaultTeamSSOParams, Dict]] = None default_team_settings: Optional[List] = None @@ -377,7 +365,7 @@ default_max_internal_user_budget: Optional[float] = None max_internal_user_budget: Optional[float] = None max_ui_session_budget: Optional[float] = 10 # $10 USD budgets for UI Chat sessions internal_user_budget_duration: Optional[str] = None -tag_budget_config: Optional[Dict[str, BudgetConfig]] = None +tag_budget_config: Optional[Dict[str, "BudgetConfig"]] = None max_end_user_budget: Optional[float] = None max_end_user_budget_id: Optional[str] = None disable_end_user_cost_tracking: Optional[bool] = None @@ -400,7 +388,9 @@ public_agent_groups: Optional[List[str]] = None # Old format: { "displayName": "url" } (for backward compatibility) public_model_groups_links: Dict[str, Union[str, Dict[str, Any]]] = {} #### REQUEST PRIORITIZATION ####### -priority_reservation: Optional[Dict[str, Union[float, PriorityReservationDict]]] = None +priority_reservation: Optional[ + Dict[str, Union[float, "PriorityReservationDict"]] +] = None priority_reservation_settings: "PriorityReservationSettings" = ( PriorityReservationSettings() ) @@ -1469,7 +1459,6 @@ from . import rag ### CUSTOM LLMs ### from .types.llms.custom_llm import CustomLLMItem -from .types.utils import GenericStreamingChunk custom_provider_map: List[CustomLLMItem] = [] _custom_providers: List[str] = ( @@ -1515,6 +1504,12 @@ if TYPE_CHECKING: from litellm.caching.caching import Cache from litellm.caching.llm_caching_handler import LLMClientCache from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES + from litellm.types.utils import ( + BudgetConfig, + CredentialItem, + PriorityReservationDict, + StandardKeyGenerationConfig, + ) # Cost calculator functions cost_per_token: Callable[..., Tuple[float, float]] @@ -1569,6 +1564,7 @@ def __getattr__(name: str) -> Any: TOKEN_COUNTER_NAMES, LLM_CLIENT_CACHE_NAMES, BEDROCK_TYPES_NAMES, + TYPES_UTILS_NAMES, CACHING_NAMES, HTTP_HANDLER_NAMES, ) @@ -1598,6 +1594,11 @@ def __getattr__(name: str) -> Any: from ._lazy_imports import _lazy_import_bedrock_types return _lazy_import_bedrock_types(name) + # Lazy load common types.utils symbols + if name in TYPES_UTILS_NAMES: + from ._lazy_imports import _lazy_import_types_utils + return _lazy_import_types_utils(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 db7c90d7d8..65104af6ed 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -50,6 +50,20 @@ BEDROCK_TYPES_NAMES = ( "COHERE_EMBEDDING_INPUT_TYPES", ) +# Common types from litellm.types.utils that support lazy loading via +# _lazy_import_types_utils +TYPES_UTILS_NAMES = ( + "ImageObject", + "BudgetConfig", + "all_litellm_params", + "_litellm_completion_params", + "CredentialItem", + "PriorityReservationDict", + "StandardKeyGenerationConfig", + "SearchProviders", + "GenericStreamingChunk", +) + # Caching / cache classes that support lazy loading via _lazy_import_caching CACHING_NAMES = ( "Cache", @@ -325,6 +339,73 @@ def _lazy_import_bedrock_types(name: str) -> Any: raise AttributeError(f"Bedrock types lazy import: unknown attribute {name!r}") +def _lazy_import_types_utils(name: str) -> Any: + """Lazy import for common types and constants from litellm.types.utils.""" + _globals = _get_litellm_globals() + + if name == "ImageObject": + from .types.utils import ImageObject as _ImageObject + + _globals["ImageObject"] = _ImageObject + return _ImageObject + + if name == "BudgetConfig": + from .types.utils import BudgetConfig as _BudgetConfig + + _globals["BudgetConfig"] = _BudgetConfig + return _BudgetConfig + + if name == "all_litellm_params": + from .types.utils import all_litellm_params as _all_litellm_params + + _globals["all_litellm_params"] = _all_litellm_params + return _all_litellm_params + + if name == "_litellm_completion_params": + from .types.utils import all_litellm_params as _all_litellm_params + + _globals["_litellm_completion_params"] = _all_litellm_params + return _all_litellm_params + + if name == "CredentialItem": + from .types.utils import CredentialItem as _CredentialItem + + _globals["CredentialItem"] = _CredentialItem + return _CredentialItem + + if name == "PriorityReservationDict": + from .types.utils import ( + PriorityReservationDict as _PriorityReservationDict, + ) + + _globals["PriorityReservationDict"] = _PriorityReservationDict + return _PriorityReservationDict + + if name == "StandardKeyGenerationConfig": + from .types.utils import ( + StandardKeyGenerationConfig as _StandardKeyGenerationConfig, + ) + + _globals["StandardKeyGenerationConfig"] = _StandardKeyGenerationConfig + return _StandardKeyGenerationConfig + + if name == "SearchProviders": + from .types.utils import SearchProviders as _SearchProviders + + _globals["SearchProviders"] = _SearchProviders + return _SearchProviders + + if name == "GenericStreamingChunk": + from .types.utils import ( + GenericStreamingChunk as _GenericStreamingChunk, + ) + + _globals["GenericStreamingChunk"] = _GenericStreamingChunk + return _GenericStreamingChunk + + raise AttributeError(f"Types utils 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 c137e6ddbc..d6e4994741 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -15,6 +15,7 @@ from litellm._lazy_imports import ( TOKEN_COUNTER_NAMES, CACHING_NAMES, BEDROCK_TYPES_NAMES, + TYPES_UTILS_NAMES, LLM_CLIENT_CACHE_NAMES, HTTP_HANDLER_NAMES, _lazy_import_cost_calculator, @@ -22,6 +23,7 @@ from litellm._lazy_imports import ( _lazy_import_utils, _lazy_import_token_counter, _lazy_import_bedrock_types, + _lazy_import_types_utils, _lazy_import_caching, _lazy_import_llm_client_cache, _lazy_import_http_handlers, @@ -127,6 +129,18 @@ def test_bedrock_types_lazy_imports(): _verify_only_requested_name_imported(name, BEDROCK_TYPES_NAMES) +def test_types_utils_lazy_imports(): + """Test that common types.utils symbols can be lazy imported.""" + for name in TYPES_UTILS_NAMES: + _clear_names_from_globals(TYPES_UTILS_NAMES) + + obj = _lazy_import_types_utils(name) + assert obj is not None + assert name in litellm.__dict__ + + _verify_only_requested_name_imported(name, TYPES_UTILS_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: @@ -171,3 +185,9 @@ def test_unknown_attribute_raises_error(): with pytest.raises(AttributeError): _lazy_import_llm_client_cache("unknown") + with pytest.raises(AttributeError): + _lazy_import_bedrock_types("unknown") + + with pytest.raises(AttributeError): + _lazy_import_types_utils("unknown") +