diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 1869328c039..60a03689804 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,7 +1,10 @@ import copy import os from collections.abc import Callable, Iterable -from typing import TYPE_CHECKING, Any, Final, Optional +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias + +from typing_extensions import assert_never import litellm from litellm import get_secret @@ -50,6 +53,66 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +@dataclass(frozen=True, slots=True) +class _CallbackResolvedToClass: + entry: str + loaded: type + tag: Literal["resolved_to_class"] = "resolved_to_class" + + +@dataclass(frozen=True, slots=True) +class _CallbackNotDispatchable: + entry: str + loaded: object + tag: Literal["not_dispatchable"] = "not_dispatchable" + + +_CallbackLoadError: TypeAlias = _CallbackResolvedToClass | _CallbackNotDispatchable + + +def _classify_loaded_callback(entry: str, loaded: object) -> CustomLogger | Callable[..., object] | _CallbackLoadError: + """ + Decide whether what a ``litellm_settings.callbacks`` dotted path resolved to can be dispatched. + + A dotted path only ever runs as a ``CustomLogger`` instance or as a callback function. Anything + else (most commonly a class instead of an instance) used to load without complaint and then be + skipped on every request, with no log line and no error. + """ + if isinstance(loaded, CustomLogger) or (callable(loaded) and not isinstance(loaded, type)): + return loaded + if isinstance(loaded, type): + return _CallbackResolvedToClass(entry=entry, loaded=loaded) + return _CallbackNotDispatchable(entry=entry, loaded=loaded) + + +def _raise_callback_load_error(error: _CallbackLoadError) -> NoReturn: + """The one edge that raises: map a load error onto config load's failure contract.""" + match error: + case _CallbackResolvedToClass(): + module_path: Final = error.entry.rsplit(".", 1)[0] if "." in error.entry else error.entry + raise ValueError( + f"litellm_settings.callbacks entry '{error.entry}' resolved to the class " + f"{error.loaded.__module__}.{error.loaded.__qualname__}, which is neither a " + "CustomLogger instance nor a callable, so the proxy would never run it." + f" Point it at an instance instead, e.g. add `proxy_handler_instance = {error.loaded.__name__}()` to " + f'{module_path} and set `callbacks: ["{module_path}.proxy_handler_instance"]`.' + ) + case _CallbackNotDispatchable(): + raise ValueError( + f"litellm_settings.callbacks entry '{error.entry}' resolved to " + f"{type(error.loaded).__name__} {error.loaded!r}, which is neither a " + "CustomLogger instance nor a callable, so the proxy would never run it." + ) + assert_never(error) + + +def _loaded_callback_or_raise(entry: str, loaded: object) -> CustomLogger | Callable[..., object]: + resolved: Final = _classify_loaded_callback(entry=entry, loaded=loaded) + if isinstance(resolved, _CallbackResolvedToClass | _CallbackNotDispatchable): + _raise_callback_load_error(resolved) + return resolved + + def initialize_callbacks_on_proxy( value: Any, premium_user: bool, @@ -305,9 +368,12 @@ def initialize_callbacks_on_proxy( "%s attempting to import custom calback=%s %s", blue_color_code, callback, reset_color_code ) imported_list.append( - get_instance_fn( - value=callback, - config_file_path=config_file_path, + _loaded_callback_or_raise( + entry=callback, + loaded=get_instance_fn( + value=callback, + config_file_path=config_file_path, + ), ) ) if isinstance(litellm.callbacks, list): @@ -321,9 +387,12 @@ def initialize_callbacks_on_proxy( PrometheusLogger._mount_metrics_endpoint() else: litellm.callbacks = [ - get_instance_fn( - value=value, - config_file_path=config_file_path, + _loaded_callback_or_raise( + entry=value, + loaded=get_instance_fn( + value=value, + config_file_path=config_file_path, + ), ) ] verbose_proxy_logger.debug("%s Initialized Callbacks - %s %s", blue_color_code, litellm.callbacks, reset_color_code) diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index bfd4ffe1593..0c73c3fcf22 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -21,6 +21,10 @@ from litellm.proxy.common_utils.callback_utils import ( strip_callback_config, ) import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging from unittest.mock import patch from litellm.proxy.common_utils.callback_utils import process_callback @@ -491,3 +495,163 @@ def test_strip_callback_config_drops_credential_bearing_slots(): @pytest.mark.parametrize("value", [None, "not-a-dict", 42]) def test_strip_callback_config_passes_through_non_dicts(value): assert strip_callback_config(value) is value + + +# --------------------------------------------------------------------------- +# initialize_callbacks_on_proxy: dotted-path entries must resolve to something +# the request path can actually dispatch +# --------------------------------------------------------------------------- + +_PROBE_MODULE_NAME = "custom_callback_probe" + +_PROBE_MODULE_SOURCE = ''' +from litellm.integrations.custom_logger import CustomLogger + + +class FloorMaxTokens(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + data["max_tokens"] = 16 + return data + + +class NotALogger: + pass + + +def log_event_fn(kwargs, response_obj, start_time, end_time): + return None + + +NOT_A_CALLBACK = "some-plain-string" + +proxy_handler_instance = FloorMaxTokens() +''' + + +@pytest.fixture +def probe_config_path(tmp_path): + """Write a callback module next to a config.yaml, the layout get_instance_fn's file + branch expects, and restore every global the load + dispatch path touches. + + ``ProxyLogging._callback_capabilities_cache`` is keyed on the id()s of the + litellm.callbacks members, so an entry left behind here can be read back by an + unrelated test whose (len, ids) signature happens to collide. + """ + (tmp_path / f"{_PROBE_MODULE_NAME}.py").write_text(_PROBE_MODULE_SOURCE) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else litellm.callbacks + ) + litellm.callbacks = [] + ProxyLogging._callback_capabilities_cache.clear() + try: + yield str(tmp_path / "config.yaml") + finally: + litellm.callbacks = original_callbacks + ProxyLogging._callback_capabilities_cache.clear() + + +def _load_callbacks(value, config_file_path): + initialize_callbacks_on_proxy( + value=value, + premium_user=False, + config_file_path=config_file_path, + litellm_settings={}, + callback_specific_params={}, + ) + + +def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_path): + """A class path loads an object that fails the `isinstance(_callback, CustomLogger)` + dispatch gate in ProxyLogging.pre_call_hook, so the proxy used to boot clean and + silently never run the hook. Config load must fail instead.""" + entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" + + with pytest.raises(ValueError) as exc_info: + _load_callbacks([entry], probe_config_path) + + message = str(exc_info.value) + assert entry in message + assert "the class" in message + assert "FloorMaxTokens" in message + assert f"{_PROBE_MODULE_NAME}.proxy_handler_instance" in message + assert litellm.callbacks == [] + + +@pytest.mark.parametrize( + "attribute, expected_fragment", + [ + ("NotALogger", "the class"), + ("NOT_A_CALLBACK", "str 'some-plain-string'"), + ], +) +def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( + probe_config_path, attribute, expected_fragment +): + entry = f"{_PROBE_MODULE_NAME}.{attribute}" + + with pytest.raises(ValueError) as exc_info: + _load_callbacks([entry], probe_config_path) + + message = str(exc_info.value) + assert entry in message + assert expected_fragment in message + assert litellm.callbacks == [] + + +def test_initialize_callbacks_on_proxy_rejects_class_valued_non_list_value(probe_config_path): + entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" + + with pytest.raises(ValueError) as exc_info: + _load_callbacks(entry, probe_config_path) + + assert entry in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_initialize_callbacks_on_proxy_instance_entry_runs_pre_call_hook(probe_config_path): + """Positive control: the supported shape must still load AND still run. Drives the + real ProxyLogging.pre_call_hook, which is where a class-valued entry goes silent.""" + _load_callbacks([f"{_PROBE_MODULE_NAME}.proxy_handler_instance"], probe_config_path) + + assert len(litellm.callbacks) == 1 + assert isinstance(litellm.callbacks[0], CustomLogger) + + ProxyLogging._callback_capabilities_cache.clear() + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + data = await proxy_logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-probe"), + data={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 1, + "metadata": {}, + }, + call_type="acompletion", + ) + + assert data["max_tokens"] == 16 + + +def test_initialize_callbacks_on_proxy_keeps_known_string_callback(probe_config_path): + """Non-narrowing control: a known callback name never reaches get_instance_fn and + stays a plain string in litellm.callbacks.""" + _load_callbacks(["langfuse"], probe_config_path) + + assert litellm.callbacks == ["langfuse"] + + +def test_initialize_callbacks_on_proxy_accepts_plain_function_callback(probe_config_path): + """Non-narrowing control: litellm.callbacks is typed + `Callable | | CustomLogger`, so a dotted path resolving to a plain + function is a supported shape and must keep loading.""" + _load_callbacks([f"{_PROBE_MODULE_NAME}.log_event_fn"], probe_config_path) + + assert [getattr(cb, "__name__", None) for cb in litellm.callbacks] == ["log_event_fn"] + + +def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_config_path): + _load_callbacks(f"{_PROBE_MODULE_NAME}.proxy_handler_instance", probe_config_path) + + assert len(litellm.callbacks) == 1 + assert isinstance(litellm.callbacks[0], CustomLogger)