diff --git a/litellm/__init__.py b/litellm/__init__.py index ef44aa53a13..05968e8a7c0 100644 --- a/litellm/__init__.py +++ b/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) diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 91b16864de1..b87b9c955fb 100644 --- a/litellm/_lazy_imports.py +++ b/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 \ No newline at end of file + 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}") \ No newline at end of file diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py new file mode 100644 index 00000000000..e36acf22eb8 --- /dev/null +++ b/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") +