Merge aaa41a35cb into c0ca4a51f7

litellm/litellm_core_utils/custom_logger_registry.py

@@ -8,6 +8,8 @@ Example:
     "prometheus" -> PrometheusLogger
 """
 
+import functools
+from collections.abc import Mapping
 from typing import Final
 
 from litellm import _custom_logger_compatible_callbacks_literal
@@ -53,22 +55,46 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i
     VectorStorePreCallHook,
 )
 from litellm.integrations.zerobus import ZerobusLogger
-from litellm.proxy.hooks.dynamic_rate_limiter import (  # noqa: F401  # legacy module exports
-    PROXY_DynamicRateLimitHandler,
-    _PROXY_DynamicRateLimitHandler,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
-)
-from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (  # noqa: F401  # legacy module exports
-    PROXY_DynamicRateLimitHandlerV3,
-    _PROXY_DynamicRateLimitHandlerV3,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
-)
 
 
-class CustomLoggerRegistry:
-    """
-    Registry mapping the callback class string to the class type.
-    """
+def _enterprise_callback_classes() -> Mapping[str, type]:
+    try:
+        from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import (
+            PagerDutyAlerting,
+        )
+        from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import (
+            ResendEmailLogger,
+        )
+        from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import (
+            SendGridEmailLogger,
+        )
+        from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import (
+            SMTPEmailLogger,
+        )
 
-    CALLBACK_CLASS_STR_TO_CLASS_TYPE = {
+        from litellm.integrations.generic_api.generic_api_callback import (
+            GenericAPILogger,
+        )
+    except ImportError:
+        return {}
+    return {
+        "pagerduty": PagerDutyAlerting,
+        "generic_api": GenericAPILogger,
+        "resend_email": ResendEmailLogger,
+        "sendgrid_email": SendGridEmailLogger,
+        "smtp_email": SMTPEmailLogger,
+    }
+
+
+_ENTERPRISE_CALLBACK_CLASSES: Final = _enterprise_callback_classes()
+
+
+@functools.cache
+def _callback_class_str_to_class_type() -> Mapping[str, type]:
+    from litellm.proxy.hooks.dynamic_rate_limiter import PROXY_DynamicRateLimitHandler
+    from litellm.proxy.hooks.dynamic_rate_limiter_v3 import PROXY_DynamicRateLimitHandlerV3
+
+    return {
         "lago": LagoLogger,
         "openmeter": OpenMeterLogger,
         "braintrust": BraintrustLogger,
@@ -119,36 +145,39 @@ class CustomLoggerRegistry:
         "vantage": VantageLogger,
         "posthog": PostHogLogger,
         "newrelic": NewRelicLogger,
+        **_ENTERPRISE_CALLBACK_CLASSES,
     }
 
-    try:
-        from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import (
-            PagerDutyAlerting,
-        )
-        from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import (
-            ResendEmailLogger,
-        )
-        from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import (
-            SendGridEmailLogger,
-        )
-        from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import (
-            SMTPEmailLogger,
-        )
 
-        from litellm.integrations.generic_api.generic_api_callback import (
-            GenericAPILogger,
-        )
+class _LazyCallbackClassMap:
+    def __get__(self, instance: object, owner: type | None = None) -> Mapping[str, type]:
+        return _callback_class_str_to_class_type()
 
-        enterprise_loggers = {
-            "pagerduty": PagerDutyAlerting,
-            "generic_api": GenericAPILogger,
-            "resend_email": ResendEmailLogger,
-            "sendgrid_email": SendGridEmailLogger,
-            "smtp_email": SMTPEmailLogger,
-        }
-        CALLBACK_CLASS_STR_TO_CLASS_TYPE.update(enterprise_loggers)
-    except ImportError:
-        pass  # enterprise not installed
+
+_LEGACY_RATE_LIMITER_EXPORTS: Final = frozenset(("PROXY_DynamicRateLimitHandler", "_PROXY_DynamicRateLimitHandler"))
+_LEGACY_RATE_LIMITER_V3_EXPORTS: Final = frozenset(
+    ("PROXY_DynamicRateLimitHandlerV3", "_PROXY_DynamicRateLimitHandlerV3")
+)
+
+
+def __getattr__(name: str) -> object:
+    if name in _LEGACY_RATE_LIMITER_EXPORTS:
+        from litellm.proxy.hooks import dynamic_rate_limiter
+
+        return getattr(dynamic_rate_limiter, name)
+    if name in _LEGACY_RATE_LIMITER_V3_EXPORTS:
+        from litellm.proxy.hooks import dynamic_rate_limiter_v3
+
+        return getattr(dynamic_rate_limiter_v3, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+class CustomLoggerRegistry:
+    """
+    Registry mapping the callback class string to the class type.
+    """
+
+    CALLBACK_CLASS_STR_TO_CLASS_TYPE = _LazyCallbackClassMap()
 
     @classmethod
     def get_callback_str_from_class_type(cls, class_type: type) -> str | None:

tests/unit/litellm_core_utils/test_custom_logger_registry.py

@@ -1,6 +1,7 @@
 import asyncio
 import importlib
 import os
+import sys
 from collections.abc import AsyncIterator
 from datetime import datetime
 from typing import Final, Literal
@@ -406,3 +407,93 @@ def setup_and_teardown():
         if hasattr(litellm, "in_memory_llm_clients_cache"):
             litellm.in_memory_llm_clients_cache.flush_cache()
     yield
+
+
+@pytest.mark.parametrize(
+    ("callback_name", "module_path", "class_name"),
+    (
+        ("dynamic_rate_limiter", "litellm.proxy.hooks.dynamic_rate_limiter", "PROXY_DynamicRateLimitHandler"),
+        ("dynamic_rate_limiter_v3", "litellm.proxy.hooks.dynamic_rate_limiter_v3", "PROXY_DynamicRateLimitHandlerV3"),
+    ),
+)
+def test_lazy_proxy_callback_classes_are_real_classes_in_the_registry(
+    callback_name: str, module_path: str, class_name: str
+) -> None:
+    proxy_class: Final = getattr(importlib.import_module(module_path), class_name)
+
+    assert CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE[callback_name] is proxy_class
+    assert CustomLoggerRegistry.get_class_type_for_custom_logger_name(callback_name) is proxy_class
+    assert CustomLoggerRegistry.get_callback_str_from_class_type(proxy_class) == callback_name
+    assert CustomLoggerRegistry.get_all_callback_strs_from_class_type(proxy_class) == [callback_name]
+
+
+def test_registry_import_leaves_proxy_rate_limiter_modules_unloaded(monkeypatch: pytest.MonkeyPatch) -> None:
+    registry_module_name: Final = "litellm.litellm_core_utils.custom_logger_registry"
+    limiter_module_names: Final = (
+        "litellm.proxy.hooks.dynamic_rate_limiter",
+        "litellm.proxy.hooks.dynamic_rate_limiter_v3",
+    )
+    for module_name in (registry_module_name, *limiter_module_names):
+        parent_name, _, attribute_name = module_name.rpartition(".")
+        monkeypatch.setattr(importlib.import_module(parent_name), attribute_name, importlib.import_module(module_name))
+        monkeypatch.delitem(sys.modules, module_name)
+
+    fresh_registry_module: Final = importlib.import_module(registry_module_name)
+
+    assert fresh_registry_module.CustomLoggerRegistry is not CustomLoggerRegistry
+    assert [name for name in limiter_module_names if name in sys.modules] == []
+
+
+@pytest.mark.parametrize(
+    ("legacy_name", "module_path"),
+    (
+        ("PROXY_DynamicRateLimitHandler", "litellm.proxy.hooks.dynamic_rate_limiter"),
+        ("_PROXY_DynamicRateLimitHandler", "litellm.proxy.hooks.dynamic_rate_limiter"),
+        ("PROXY_DynamicRateLimitHandlerV3", "litellm.proxy.hooks.dynamic_rate_limiter_v3"),
+        ("_PROXY_DynamicRateLimitHandlerV3", "litellm.proxy.hooks.dynamic_rate_limiter_v3"),
+    ),
+)
+def test_legacy_rate_limiter_exports_resolve_to_proxy_objects(legacy_name: str, module_path: str) -> None:
+    registry_module: Final = importlib.import_module("litellm.litellm_core_utils.custom_logger_registry")
+
+    assert getattr(registry_module, legacy_name) is getattr(importlib.import_module(module_path), legacy_name)
+
+
+def test_registry_module_still_raises_for_unknown_attributes() -> None:
+    registry_module: Final = importlib.import_module("litellm.litellm_core_utils.custom_logger_registry")
+
+    with pytest.raises(AttributeError):
+        getattr(registry_module, "PROXY_NotARealHandler")
+
+
+def test_registry_maps_every_callback_name_to_a_class() -> None:
+    non_classes: Final = [
+        name
+        for name, callback_class in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE.items()
+        if not isinstance(callback_class, type)
+    ]
+
+    assert non_classes == []
+
+
+def test_registry_builds_without_enterprise_callbacks_when_enterprise_is_missing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    registry_module_name: Final = "litellm.litellm_core_utils.custom_logger_registry"
+    monkeypatch.setitem(sys.modules, "litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty", None)
+    monkeypatch.setattr(
+        litellm.litellm_core_utils, "custom_logger_registry", importlib.import_module(registry_module_name)
+    )
+    monkeypatch.delitem(sys.modules, registry_module_name)
+
+    fresh_mapping: Final = importlib.import_module(
+        registry_module_name
+    ).CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE
+
+    assert fresh_mapping.keys() == CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE.keys() - {
+        "pagerduty",
+        "generic_api",
+        "resend_email",
+        "sendgrid_email",
+        "smtp_email",
+    }