[Refactor] litellm/init.py: lazy load .types.utils (#18054)

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

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()

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")
+