[Refactor] lazy imports: Use per-attribute lazy imports and extract shared constants (#17994)

litellm/__init__.py

@@ -1563,41 +1563,24 @@ if TYPE_CHECKING:
 
 def __getattr__(name: str) -> Any:
     """Lazy import handler for cost_calculator and litellm_logging functions."""
-    # Lazy load cost_calculator functions
-    _cost_calculator_names = (
-        "completion_cost",
-        "cost_per_token",
-        "response_cost_calculator",
+    from ._lazy_imports import (
+        COST_CALCULATOR_NAMES,
+        LITELLM_LOGGING_NAMES,
+        UTILS_NAMES,
     )
-    if name in _cost_calculator_names:
+    
+    # Lazy load cost_calculator functions
+    if name in COST_CALCULATOR_NAMES:
         from ._lazy_imports import _lazy_import_cost_calculator
         return _lazy_import_cost_calculator(name)
 
     # Lazy load litellm_logging functions
-    _litellm_logging_names = (
-        "Logging",
-        "modify_integration",
-    )
-    if name in _litellm_logging_names:
+    if name in LITELLM_LOGGING_NAMES:
         from ._lazy_imports import _lazy_import_litellm_logging
         return _lazy_import_litellm_logging(name)
 
     # Lazy load utils functions
-    _utils_names = (
-        "exception_type", "get_optional_params", "get_response_string", "token_counter",
-        "create_pretrained_tokenizer", "create_tokenizer", "supports_function_calling",
-        "supports_web_search", "supports_url_context", "supports_response_schema",
-        "supports_parallel_function_calling", "supports_vision", "supports_audio_input",
-        "supports_audio_output", "supports_system_messages", "supports_reasoning",
-        "get_litellm_params", "acreate", "get_max_tokens", "get_model_info",
-        "register_prompt_template", "validate_environment", "check_valid_key",
-        "register_model", "encode", "decode", "_calculate_retry_after", "_should_retry",
-        "get_supported_openai_params", "get_api_base", "get_first_chars_messages",
-        "ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse",
-        "TranscriptionResponse", "TextCompletionResponse", "get_provider_fields",
-        "ModelResponseListIterator", "get_valid_models",
-    )
-    if name in _utils_names:
+    if name in UTILS_NAMES:
         from ._lazy_imports import _lazy_import_utils
         return _lazy_import_utils(name)
     

litellm/_lazy_imports.py

@@ -5,6 +5,35 @@ def _get_litellm_globals() -> dict:
     """Helper to get the globals dictionary of the litellm module."""
     return sys.modules["litellm"].__dict__
 
+# Cost calculator names that support lazy loading via _lazy_import_cost_calculator
+COST_CALCULATOR_NAMES = (
+    "completion_cost",
+    "cost_per_token",
+    "response_cost_calculator",
+)
+
+# Litellm logging names that support lazy loading via _lazy_import_litellm_logging
+LITELLM_LOGGING_NAMES = (
+    "Logging",
+    "modify_integration",
+)
+
+# Utils names that support lazy loading via _lazy_import_utils
+UTILS_NAMES = (
+    "exception_type", "get_optional_params", "get_response_string", "token_counter",
+    "create_pretrained_tokenizer", "create_tokenizer", "supports_function_calling",
+    "supports_web_search", "supports_url_context", "supports_response_schema",
+    "supports_parallel_function_calling", "supports_vision", "supports_audio_input",
+    "supports_audio_output", "supports_system_messages", "supports_reasoning",
+    "get_litellm_params", "acreate", "get_max_tokens", "get_model_info",
+    "register_prompt_template", "validate_environment", "check_valid_key",
+    "register_model", "encode", "decode", "_calculate_retry_after", "_should_retry",
+    "get_supported_openai_params", "get_api_base", "get_first_chars_messages",
+    "ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse",
+    "TranscriptionResponse", "TextCompletionResponse", "get_provider_fields",
+    "ModelResponseListIterator", "get_valid_models",
+)
+
 # Lazy import for utils module - imports only the requested item by name.
 # Note: PLR0915 (too many statements) is suppressed because the many if statements
 # are intentional - each attribute is imported individually only when requested,
@@ -218,42 +247,35 @@ def _lazy_import_utils(name: str) -> Any:  # noqa: PLR0915
 def _lazy_import_cost_calculator(name: str) -> Any:
     """Lazy import for cost_calculator functions."""
     _globals = _get_litellm_globals()
-    from .cost_calculator import (
-        completion_cost as _completion_cost,
-        cost_per_token as _cost_per_token,
-        response_cost_calculator as _response_cost_calculator,
-    )
+    if name == "completion_cost":
+        from .cost_calculator import completion_cost as _completion_cost
+        _globals["completion_cost"] = _completion_cost
+        return _completion_cost
     
-    _cost_functions = {
-        "completion_cost": _completion_cost,
-        "cost_per_token": _cost_per_token,
-        "response_cost_calculator": _response_cost_calculator,
-    }
+    if name == "cost_per_token":
+        from .cost_calculator import cost_per_token as _cost_per_token
+        _globals["cost_per_token"] = _cost_per_token
+        return _cost_per_token
     
-    func = _cost_functions[name]
-    _globals[name] = func
-    return func
+    if name == "response_cost_calculator":
+        from .cost_calculator import response_cost_calculator as _response_cost_calculator
+        _globals["response_cost_calculator"] = _response_cost_calculator
+        return _response_cost_calculator
+    
+    raise AttributeError(f"Cost calculator lazy import: unknown attribute {name!r}")
 
 
 def _lazy_import_litellm_logging(name: str) -> Any:
     """Lazy import for litellm_logging module."""
     _globals = _get_litellm_globals()
-    try:
-        from litellm.litellm_core_utils.litellm_logging import (
-            Logging as _Logging,
-            modify_integration as _modify_integration,
-        )
-        
-        _logging_objects = {
-            "Logging": _Logging,
-            "modify_integration": _modify_integration,
-        }
-        
-        obj = _logging_objects[name]
-        _globals[name] = obj
-        return obj
-    except Exception as e:
-        raise AttributeError(
-            f"module 'litellm' has no attribute {name!r}. "
-            f"Lazy import failed: {e}"
-        ) from e
+    if name == "Logging":
+        from litellm.litellm_core_utils.litellm_logging import Logging as _Logging
+        _globals["Logging"] = _Logging
+        return _Logging
+    
+    if name == "modify_integration":
+        from litellm.litellm_core_utils.litellm_logging import modify_integration as _modify_integration
+        _globals["modify_integration"] = _modify_integration
+        return _modify_integration
+    
+    raise AttributeError(f"Litellm logging lazy import: unknown attribute {name!r}")

tests/test_litellm/test_lazy_imports.py

@@ -0,0 +1,91 @@
+"""Simple tests for lazy import functionality."""
+
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import litellm
+from litellm._lazy_imports import (
+    COST_CALCULATOR_NAMES,
+    LITELLM_LOGGING_NAMES,
+    UTILS_NAMES,
+    _lazy_import_cost_calculator,
+    _lazy_import_litellm_logging,
+    _lazy_import_utils,
+)
+
+
+def _clear_names_from_globals(names: tuple):
+    """Clear all names from litellm globals."""
+    for name in names:
+        if name in litellm.__dict__:
+            del litellm.__dict__[name]
+
+
+def _verify_only_requested_name_imported(name: str, all_names: tuple):
+    """Verify that only the requested name is in globals, not the others."""
+    for other_name in all_names:
+        if other_name != name:
+            assert other_name not in litellm.__dict__, f"{other_name} should not be imported when importing {name}"
+
+
+def test_cost_calculator_lazy_imports():
+    """Test that all cost calculator functions can be lazy imported."""
+    # Test each name individually - only that name should be imported
+    for name in COST_CALCULATOR_NAMES:
+        # Clear all names before importing just one
+        _clear_names_from_globals(COST_CALCULATOR_NAMES)
+        
+        func = _lazy_import_cost_calculator(name)
+        assert func is not None
+        assert callable(func)
+        assert name in litellm.__dict__
+        
+        # Verify only the requested name is in globals, not the others
+        _verify_only_requested_name_imported(name, COST_CALCULATOR_NAMES)
+
+
+def test_litellm_logging_lazy_imports():
+    """Test that all litellm_logging items can be lazy imported."""
+    # Test each name individually - only that name should be imported
+    for name in LITELLM_LOGGING_NAMES:
+        # Clear all names before importing just one
+        _clear_names_from_globals(LITELLM_LOGGING_NAMES)
+        
+        item = _lazy_import_litellm_logging(name)
+        assert item is not None
+        assert name in litellm.__dict__
+        
+        # Verify only the requested name is in globals, not the others
+        _verify_only_requested_name_imported(name, LITELLM_LOGGING_NAMES)
+
+
+def test_utils_lazy_imports():
+    """Test that all utils functions can be lazy imported."""
+    # Test each name individually - only that name should be imported
+    for name in UTILS_NAMES:
+        # Clear all names before importing just one
+        _clear_names_from_globals(UTILS_NAMES)
+        
+        attr = _lazy_import_utils(name)
+        assert attr is not None
+        assert name in litellm.__dict__
+        
+        # Verify only the requested name is in globals, not the others
+        _verify_only_requested_name_imported(name, UTILS_NAMES)
+
+
+def test_unknown_attribute_raises_error():
+    """Test that unknown attributes raise AttributeError."""
+    with pytest.raises(AttributeError):
+        _lazy_import_cost_calculator("unknown")
+    
+    with pytest.raises(AttributeError):
+        _lazy_import_litellm_logging("unknown")
+    
+    with pytest.raises(AttributeError):
+        _lazy_import_utils("unknown")
+