diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 8367c9ee00e..9f84a67c89e 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/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: diff --git a/tests/unit/litellm_core_utils/test_custom_logger_registry.py b/tests/unit/litellm_core_utils/test_custom_logger_registry.py index 5fc21d75e3e..eb5aaa6f187 100644 --- a/tests/unit/litellm_core_utils/test_custom_logger_registry.py +++ b/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", + }