From 36f8c9463f3a447b35bcbd4e8588b9c883ef73fe Mon Sep 17 00:00:00 2001 From: AlexsanderHamir Date: Tue, 25 Nov 2025 12:13:19 -0800 Subject: [PATCH] fix: resolve type checking errors for lazy-loaded functions in budget_manager - Add type stubs and @overload decorators for cost_per_token and completion_cost - Refactor lazy loading system with centralized registry for better maintainability - Add comprehensive documentation for adding new lazy-loaded functions - Fixes 'Any? not callable' errors at lines 111, 139, and 146 in budget_manager.py --- litellm/__init__.py | 82 ++++++++++++++++++++++++++++++++++----------- 1 file changed, 63 insertions(+), 19 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 8e8815b125..a513619c4d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -20,6 +20,8 @@ from typing import ( Literal, get_args, TYPE_CHECKING, + Tuple, + overload, ) from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams from litellm.types.integrations.datadog import DatadogInitParams @@ -1503,8 +1505,20 @@ def set_global_gitlab_config(config: Dict[str, Any]) -> None: global_gitlab_config = config -# Lazy import for cost_calculator functions to avoid loading the module at import time -# This significantly reduces memory usage when importing litellm +# ============================================================================ +# LAZY LOADING SYSTEM +# ============================================================================ +# This system allows heavy modules to be loaded only when accessed, +# significantly reducing initial import time and memory usage. +# +# To add a new lazy-loaded function/class: +# 1. Add the import handler function (e.g., _lazy_import_xxx) +# 2. Add entries to _LAZY_LOAD_REGISTRY below +# 3. Add type stubs in the TYPE_CHECKING block +# 4. Add @overload decorators for type checking +# ============================================================================ + + def _lazy_import_cost_calculator(name: str) -> Any: """Lazy import for cost_calculator functions.""" from .cost_calculator import ( @@ -1513,22 +1527,17 @@ def _lazy_import_cost_calculator(name: str) -> Any: response_cost_calculator as _response_cost_calculator, ) - # Map names to imported functions _cost_functions = { "completion_cost": _completion_cost, "cost_per_token": _cost_per_token, "response_cost_calculator": _response_cost_calculator, } - # Cache the imported function in the module namespace func = _cost_functions[name] - globals()[name] = func - + globals()[name] = func # Cache for future access return func -# Lazy import for litellm_logging to avoid loading the module at import time -# This significantly reduces memory usage when importing litellm def _lazy_import_litellm_logging(name: str) -> Any: """Lazy import for litellm_logging module.""" try: @@ -1537,35 +1546,70 @@ def _lazy_import_litellm_logging(name: str) -> Any: modify_integration as _modify_integration, ) - # Map names to imported objects _logging_objects = { "Logging": _Logging, "modify_integration": _modify_integration, } - # Cache the imported object in the module namespace obj = _logging_objects[name] - globals()[name] = obj - + globals()[name] = obj # Cache for future access return obj except Exception as e: - # If lazy import fails, raise a more informative error raise AttributeError( f"module {__name__!r} has no attribute {name!r}. " f"Lazy import failed: {e}" ) from e +# Registry mapping lazy-loaded names to their import handlers +# Add new lazy-loaded items here for easy maintenance +_LAZY_LOAD_REGISTRY: Dict[str, Callable[[str], Any]] = { + # Cost calculator functions + "completion_cost": _lazy_import_cost_calculator, + "cost_per_token": _lazy_import_cost_calculator, + "response_cost_calculator": _lazy_import_cost_calculator, + # Logging objects + "Logging": _lazy_import_litellm_logging, + "modify_integration": _lazy_import_litellm_logging, +} + + +# Type stubs for lazy-loaded functions/classes to help type checkers +# Add type annotations here for new lazy-loaded items +if TYPE_CHECKING: + # Cost calculator functions + cost_per_token: Callable[..., Tuple[float, float]] + completion_cost: Callable[..., float] + response_cost_calculator: Any + # Logging objects + Logging: Any + modify_integration: Any + + +# Type overloads for __getattr__ to provide proper type hints +# Add @overload decorators here for new lazy-loaded items with specific types +@overload +def __getattr__(name: Literal["cost_per_token"]) -> Callable[..., Tuple[float, float]]: + ... + + +@overload +def __getattr__(name: Literal["completion_cost"]) -> Callable[..., float]: + ... + + +@overload +def __getattr__(name: Literal["response_cost_calculator", "Logging", "modify_integration"]) -> Any: + ... + + def __getattr__(name: str) -> Any: - """Lazy import for cost_calculator and litellm_logging functions. + """Lazy import handler for cost_calculator and litellm_logging functions. This allows these heavy modules to be loaded only when accessed, reducing initial import time and memory usage. """ - if name in ("completion_cost", "response_cost_calculator", "cost_per_token"): - return _lazy_import_cost_calculator(name) - - if name in ("Logging", "modify_integration"): - return _lazy_import_litellm_logging(name) + if name in _LAZY_LOAD_REGISTRY: + return _LAZY_LOAD_REGISTRY[name](name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}")