diff --git a/litellm/__init__.py b/litellm/__init__.py index 69f80123a0..98e5a34bd6 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1,4 +1,6 @@ ### Hide pydantic namespace conflict warnings globally ### +from __future__ import annotations + import warnings warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*") @@ -72,7 +74,6 @@ from litellm.constants import ( DEFAULT_SOFT_BUDGET, DEFAULT_ALLOWED_FAILS, ) -from litellm.types.guardrails import GuardrailItem from litellm.types.secret_managers.main import ( KeyManagementSystem, KeyManagementSettings, @@ -1506,6 +1507,7 @@ if TYPE_CHECKING: PriorityReservationDict, StandardKeyGenerationConfig, ) + from litellm.types.guardrails import GuardrailItem # Cost calculator functions cost_per_token: Callable[..., Tuple[float, float]] @@ -1568,6 +1570,7 @@ def __getattr__(name: str) -> Any: HTTP_HANDLER_NAMES, DOTPROMPT_NAMES, LLM_CONFIG_NAMES, + TYPES_NAMES, ) # Lazy load cost_calculator functions @@ -1628,6 +1631,12 @@ def __getattr__(name: str) -> Any: return _lazy_import_llm_configs(name) + # Lazy load types + if name in TYPES_NAMES: + from ._lazy_imports import _lazy_import_types + + return _lazy_import_types(name) + # Lazy load encoding from main.py to avoid heavy tiktoken import if name == "encoding": from .main import encoding as _encoding diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index ba62355cc8..9da861a102 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -159,6 +159,11 @@ LLM_CONFIG_NAMES = ( "AmazonConverseConfig", ) +# Types that support lazy loading via _lazy_import_types +TYPES_NAMES = ( + "GuardrailItem", +) + # 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, @@ -619,6 +624,21 @@ def _lazy_import_dotprompt(name: str) -> Any: raise AttributeError(f"Dotprompt lazy import: unknown attribute {name!r}") +def _lazy_import_types(name: str) -> Any: + """Lazy import for type classes.""" + _globals = _get_litellm_globals() + + if name == "GuardrailItem": + from litellm.types.guardrails import ( + GuardrailItem as _GuardrailItem, + ) + + _globals["GuardrailItem"] = _GuardrailItem + return _GuardrailItem + + raise AttributeError(f"Types lazy import: unknown attribute {name!r}") + + def _lazy_import_llm_configs(name: str) -> Any: """Lazy import for LLM config classes.""" _globals = _get_litellm_globals() diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index 0b38956441..0eaedaab60 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -31,6 +31,8 @@ from litellm._lazy_imports import ( _lazy_import_dotprompt, LLM_CONFIG_NAMES, _lazy_import_llm_configs, + TYPES_NAMES, + _lazy_import_types, ) @@ -213,6 +215,9 @@ def test_unknown_attribute_raises_error(): with pytest.raises(AttributeError): _lazy_import_llm_configs("unknown") + with pytest.raises(AttributeError): + _lazy_import_types("unknown") + def test_llm_config_lazy_imports(): """Test that LLM config classes can be lazy imported.""" @@ -227,3 +232,17 @@ def test_llm_config_lazy_imports(): _verify_only_requested_name_imported(name, LLM_CONFIG_NAMES) + +def test_types_lazy_imports(): + """Test that type classes can be lazy imported.""" + for name in TYPES_NAMES: + _clear_names_from_globals(TYPES_NAMES) + + obj = _lazy_import_types(name) + assert obj is not None + assert name in litellm.__dict__ + # Type classes should be classes/types + assert isinstance(obj, type), f"{name} should be a class" + + _verify_only_requested_name_imported(name, TYPES_NAMES) +