diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index c40aee63bf7..9f84a67c89e 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -8,11 +8,11 @@ Example: "prometheus" -> PrometheusLogger """ -import importlib -from dataclasses import dataclass -from types import MappingProxyType +import functools +from collections.abc import Mapping from typing import Final +from litellm import _custom_logger_compatible_callbacks_literal from litellm.integrations.agentops import AgentOps from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.argilla import ArgillaLogger @@ -56,51 +56,45 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i ) from litellm.integrations.zerobus import ZerobusLogger -_DYNAMIC_RATE_LIMITER_MODULE: Final = "litellm.proxy.hooks.dynamic_rate_limiter" -_DYNAMIC_RATE_LIMITER_V3_MODULE: Final = "litellm.proxy.hooks.dynamic_rate_limiter_v3" -_LEGACY_EXPORT_MODULES: Final = MappingProxyType( - { - "PROXY_DynamicRateLimitHandler": _DYNAMIC_RATE_LIMITER_MODULE, - "_PROXY_DynamicRateLimitHandler": _DYNAMIC_RATE_LIMITER_MODULE, - "PROXY_DynamicRateLimitHandlerV3": _DYNAMIC_RATE_LIMITER_V3_MODULE, - "_PROXY_DynamicRateLimitHandlerV3": _DYNAMIC_RATE_LIMITER_V3_MODULE, + +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, + ) + + 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, } -) -def __getattr__(name: str) -> object: - module: Final = _LEGACY_EXPORT_MODULES.get(name) - if module is None: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - return getattr(importlib.import_module(module), name) +_ENTERPRISE_CALLBACK_CLASSES: Final = _enterprise_callback_classes() -@dataclass(frozen=True, slots=True) -class _LazyCallbackClass: - module: str - name: str +@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 - def resolve(self) -> type: - loaded: Final[object] = getattr(importlib.import_module(self.module), self.name) - if not isinstance(loaded, type): - raise TypeError(f"{self.module}.{self.name} is not a class") - return loaded - - -def _resolve(callback_class: type | _LazyCallbackClass) -> type: - match callback_class: - case _LazyCallbackClass(): - return callback_class.resolve() - case _: - return callback_class - - -class CustomLoggerRegistry: - """ - Registry mapping the callback class string to the class type. - """ - - CALLBACK_CLASS_STR_TO_CLASS_TYPE = { + return { "lago": LagoLogger, "openmeter": OpenMeterLogger, "braintrust": BraintrustLogger, @@ -139,10 +133,8 @@ class CustomLoggerRegistry: "pointfive": PointFiveLogger, "zerobus": ZerobusLogger, "aws_sqs": SQSLogger, - "dynamic_rate_limiter": _LazyCallbackClass(_DYNAMIC_RATE_LIMITER_MODULE, "PROXY_DynamicRateLimitHandler"), - "dynamic_rate_limiter_v3": _LazyCallbackClass( - _DYNAMIC_RATE_LIMITER_V3_MODULE, "PROXY_DynamicRateLimitHandlerV3" - ), + "dynamic_rate_limiter": PROXY_DynamicRateLimitHandler, + "dynamic_rate_limiter_v3": PROXY_DynamicRateLimitHandlerV3, "vector_store_pre_call_hook": VectorStorePreCallHook, "dotprompt": DotpromptManager, "bitbucket": BitBucketPromptManager, @@ -153,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: @@ -199,7 +194,7 @@ class CustomLoggerRegistry: callback_str, callback_class, ) in cls.CALLBACK_CLASS_STR_TO_CLASS_TYPE.items(): - if _resolve(callback_class) == class_type: + if callback_class == class_type: return callback_str return None @@ -220,16 +215,16 @@ class CustomLoggerRegistry: callback_str, callback_class, ) in cls.CALLBACK_CLASS_STR_TO_CLASS_TYPE.items(): - if _resolve(callback_class) == class_type: + if callback_class == class_type: callback_strs.append(callback_str) return callback_strs @classmethod def get_class_type_for_custom_logger_name( cls, - custom_logger_name: str, + custom_logger_name: _custom_logger_compatible_callbacks_literal, ) -> type: """ Get the class type for a given custom logger name """ - return _resolve(cls.CALLBACK_CLASS_STR_TO_CLASS_TYPE[custom_logger_name]) + return cls.CALLBACK_CLASS_STR_TO_CLASS_TYPE[custom_logger_name] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f75d4a954a..089647c6297 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -19331,7 +19331,7 @@ def _hidden_runtime_callback_names(configured_callback_names: frozenset[str]) -> from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry configured_classes: Final = tuple( - CustomLoggerRegistry.get_class_type_for_custom_logger_name(name) + CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE[name] for name in configured_callback_names if name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE ) 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 f028e528418..8c839cb05f0 100644 --- a/tests/unit/litellm_core_utils/test_custom_logger_registry.py +++ b/tests/unit/litellm_core_utils/test_custom_logger_registry.py @@ -416,11 +416,12 @@ def setup_and_teardown(): ("dynamic_rate_limiter_v3", "litellm.proxy.hooks.dynamic_rate_limiter_v3", "PROXY_DynamicRateLimitHandlerV3"), ), ) -def test_lazy_proxy_callback_classes_resolve_in_registry_lookups( +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] @@ -463,3 +464,13 @@ def test_registry_module_still_raises_for_unknown_attributes() -> None: 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 == []