From 577c22d6a90e25743d22b98d4c55ec819f59cac2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 2 Aug 2026 02:12:52 +0000 Subject: [PATCH] fix(proxy): activate interception callbacks enabled via DB config Enabling websearch_interception, compression_interception or code_interpreter_interception from the Admin UI / POST /config/update persisted the setting and reported success, but the feature never turned on; only the config.yaml boot path worked. _add_callbacks_from_db_config could not build these loggers. They need config params (websearch_interception_params and friends), which _init_custom_logger_compatible_class does not take, so websearch and code_interpreter fell through to appending the bare callback name to litellm.callbacks, and compression (which is in _custom_logger_compatible_callbacks_literal) resolved to None and was dropped entirely. No hooks were registered either way. Resolution now lives in one place: resolve_config_parameterized_callback in callback_utils, used by both initialize_callbacks_on_proxy (YAML boot) and the DB-config path. The DB path installs through install_config_parameterized_callback, which is idempotent across config polls and swaps the installed logger when its params change instead of stacking a second instance. initialize_from_proxy_config now takes Mapping instead of dict, since these config sections are read-only inputs. --- basedpyright-code-budget.json | 2 +- .../code_interpreter_interception/handler.py | 5 +- .../compression_interception/handler.py | 5 +- .../websearch_interception/handler.py | 4 +- litellm/proxy/common_utils/callback_utils.py | 138 ++++++++++++++---- litellm/proxy/proxy_server.py | 12 +- .../proxy/common_utils/test_callback_utils.py | 108 ++++++++++++++ .../proxy/proxy_server/test_proxy_config.py | 134 +++++++++++++++++ type-discipline-budget.json | 2 +- 9 files changed, 370 insertions(+), 40 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index f6dd90077b1..dd94286b4ce 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45324 + "limit": 45322 }, "reportUnknownLambdaType": { "limit": 113 diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index db34f00b051..4e45161487d 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -9,6 +9,7 @@ captured stdout back through the typed agentic loop plan. import json import time import uuid +from collections.abc import Mapping from typing import Any, Literal, TypedDict, cast from pydantic import ValidationError @@ -168,8 +169,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( - litellm_settings: dict[str, Any], - callback_specific_params: dict[str, Any], + litellm_settings: Mapping[str, Any], + callback_specific_params: Mapping[str, Any], ) -> "CodeInterpreterInterceptionLogger": params: CodeInterpreterInterceptionConfig = {} if "code_interpreter_interception_params" in litellm_settings: diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 93765001c94..fa66d2d47a7 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -7,6 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan. import time import uuid +from collections.abc import Mapping from typing import Any, cast from litellm._logging import verbose_logger @@ -100,8 +101,8 @@ class CompressionInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( - litellm_settings: dict[str, Any], - callback_specific_params: dict[str, Any], + litellm_settings: Mapping[str, Any], + callback_specific_params: Mapping[str, Any], ) -> "CompressionInterceptionLogger": compression_params: CompressionInterceptionConfig = {} if "compression_interception_params" in litellm_settings: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 54278afafc4..6dd4875c6f4 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1577,8 +1577,8 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( - litellm_settings: dict[str, Any], - callback_specific_params: dict[str, Any], + litellm_settings: Mapping[str, Any], + callback_specific_params: Mapping[str, Any], ) -> "WebSearchInterceptionLogger": """ Static method to initialize WebSearchInterceptionLogger from proxy config. diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 1eca0eb768c..b18994afc99 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,6 +1,7 @@ import copy import os -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Optional import litellm @@ -46,6 +47,105 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +_NO_CALLBACK_SPECIFIC_PARAMS: Mapping[str, Any] = MappingProxyType({}) + + +def resolve_config_parameterized_callback( + callback: str, + litellm_settings: Mapping[str, Any], + callback_specific_params: Mapping[str, Any] | None = None, +) -> CustomLogger | None: + """ + Build the ``CustomLogger`` for callbacks whose construction needs config params. + + These callbacks cannot be resolved through ``_init_custom_logger_compatible_class`` + because that factory takes no config, so every callsite that turns a callback name + into a logger has to come through here. + + Returns ``None`` when ``callback`` is not one of them, so callers fall through to + their own generic resolution. + """ + callback_params = callback_specific_params or _NO_CALLBACK_SPECIFIC_PARAMS + match callback: + case "compression_interception": + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + return CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_params, + ) + case "code_interpreter_interception": + from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + ) + + return CodeInterpreterInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_params, + ) + case "websearch_interception": + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + return WebSearchInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_params, + ) + case _: + return None + + +def _callback_config_state(logger: CustomLogger) -> tuple[tuple[str, object], ...]: + return tuple( + sorted( + ((name, value) for name, value in vars(logger).items() if not name.startswith("_")), + key=lambda item: item[0], + ) + ) + + +def install_config_parameterized_callback( + callback: str, + litellm_settings: Mapping[str, Any], + callback_specific_params: Mapping[str, Any] | None = None, +) -> bool: + """ + Idempotently install a config-parameterized callback into ``litellm.callbacks``. + + Written for the DB-config path, which re-runs on every config poll: an already + installed logger built from the same params is left alone, and one built from + stale params is replaced rather than stacked. + + Returns ``True`` when ``callback`` was handled here. + """ + resolved = resolve_config_parameterized_callback( + callback=callback, + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, + ) + if resolved is None: + return False + + installed = next( + ( + existing + for existing in litellm.callbacks + if isinstance(existing, CustomLogger) and type(existing) is type(resolved) + ), + None, + ) + if installed is not None: + if _callback_config_state(installed) == _callback_config_state(resolved): + return True + litellm.logging_callback_manager.remove_callback_from_all_lists(installed) + + litellm.logging_callback_manager.add_litellm_callback(resolved) + return True + + def initialize_callbacks_on_proxy( value: Any, premium_user: bool, @@ -65,29 +165,15 @@ def initialize_callbacks_on_proxy( if isinstance(value, list): imported_list: list[Any] = [] for callback in value: # ["presidio", ] - if isinstance(callback, str) and callback == "compression_interception": - from litellm.integrations.compression_interception.handler import ( - CompressionInterceptionLogger, - ) - - compression_interception_obj = CompressionInterceptionLogger.initialize_from_proxy_config( + if isinstance(callback, str): + config_parameterized_obj = resolve_config_parameterized_callback( + callback=callback, litellm_settings=litellm_settings, callback_specific_params=callback_specific_params, ) - imported_list.append(compression_interception_obj) - continue - - if isinstance(callback, str) and callback == "code_interpreter_interception": - from litellm.integrations.code_interpreter_interception.handler import ( - CodeInterpreterInterceptionLogger, - ) - - code_interpreter_interception_obj = CodeInterpreterInterceptionLogger.initialize_from_proxy_config( - litellm_settings=litellm_settings, - callback_specific_params=callback_specific_params, - ) - imported_list.append(code_interpreter_interception_obj) - continue + if config_parameterized_obj is not None: + imported_list.append(config_parameterized_obj) + continue # check if callback is a custom logger compatible callback if isinstance(callback, str): @@ -272,16 +358,6 @@ def initialize_callbacks_on_proxy( **azure_content_safety_params, ) imported_list.append(azure_content_safety_obj) - elif isinstance(callback, str) and callback == "websearch_interception": - from litellm.integrations.websearch_interception.handler import ( - WebSearchInterceptionLogger, - ) - - websearch_interception_obj = WebSearchInterceptionLogger.initialize_from_proxy_config( - litellm_settings=litellm_settings, - callback_specific_params=callback_specific_params, - ) - imported_list.append(websearch_interception_obj) elif isinstance(callback, str) and callback == "datadog_cost_management": from litellm.integrations.datadog.datadog_cost_management import ( DatadogCostManagementLogger, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 901ca39326b..a152513ca17 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -295,7 +295,10 @@ from litellm.proxy.common_request_processing import ( _should_return_raw_model_name, create_response, ) -from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.common_utils.callback_utils import ( + initialize_callbacks_on_proxy, + install_config_parameterized_callback, +) from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router @@ -5551,6 +5554,7 @@ class ProxyConfig: Adds callbacks from DB config to litellm """ litellm_settings = config_data.get("litellm_settings", {}) or {} + callback_specific_params = config_data.get("callback_settings") success_callbacks = litellm_settings.get("success_callback", None) failure_callbacks = litellm_settings.get("failure_callback", None) callbacks = litellm_settings.get("callbacks", None) @@ -5573,6 +5577,12 @@ class ProxyConfig: if callbacks is not None and isinstance(callbacks, list): for callback in callbacks: + if isinstance(callback, str) and install_config_parameterized_callback( + callback=callback, + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, + ): + continue self._add_callback_from_db_to_in_memory_litellm_callbacks( callback=callback, event_types=["success", "failure"], 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..e41766f0c86 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -15,8 +15,10 @@ from litellm.proxy.common_utils.callback_utils import ( encrypt_callback_vars, get_logging_caching_headers, initialize_callbacks_on_proxy, + install_config_parameterized_callback, get_remaining_tokens_and_requests_from_request_data, normalize_callback_names, + resolve_config_parameterized_callback, sanitize_openai_provider_metadata, strip_callback_config, ) @@ -220,6 +222,112 @@ def test_initialize_callbacks_on_proxy_instantiates_compression_interception( litellm.callbacks = original_callbacks +@pytest.mark.parametrize( + "callback_name, params_key, params, expected_attrs", + [ + ( + "websearch_interception", + "websearch_interception_params", + {"enabled_providers": ["bedrock", "openai"], "search_tool_name": "tavily-search"}, + {"enabled_providers": ["bedrock", "openai"], "search_tool_name": "tavily-search"}, + ), + ( + "compression_interception", + "compression_interception_params", + {"enabled": False, "compression_trigger": 4242}, + {"enabled": False, "compression_trigger": 4242}, + ), + ( + "code_interpreter_interception", + "code_interpreter_interception_params", + {"enabled": True, "sandbox_tool_name": "my-sandbox"}, + {"enabled": True, "sandbox_tool_name": "my-sandbox"}, + ), + ], +) +def test_resolve_config_parameterized_callback_forwards_litellm_settings_params( + callback_name, params_key, params, expected_attrs +): + resolved = resolve_config_parameterized_callback( + callback=callback_name, + litellm_settings={params_key: params}, + callback_specific_params={}, + ) + + assert resolved is not None + assert {attr: getattr(resolved, attr) for attr in expected_attrs} == expected_attrs + + +def test_resolve_config_parameterized_callback_falls_back_to_callback_specific_params(): + resolved = resolve_config_parameterized_callback( + callback="websearch_interception", + litellm_settings={}, + callback_specific_params={"websearch_interception": {"search_tool_name": "exa-search"}}, + ) + + assert resolved is not None + assert resolved.search_tool_name == "exa-search" + + +@pytest.mark.parametrize("callback_name", ["langfuse", "prometheus", "my_custom_module.my_logger"]) +def test_resolve_config_parameterized_callback_returns_none_for_other_callbacks(callback_name): + assert ( + resolve_config_parameterized_callback( + callback=callback_name, + litellm_settings={}, + callback_specific_params={}, + ) + is None + ) + + +def test_initialize_callbacks_on_proxy_instantiates_websearch_interception_with_params( + monkeypatch, +): + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(prisma_client=None), + ) + monkeypatch.setattr(litellm, "callbacks", [], raising=False) + + initialize_callbacks_on_proxy( + value=["websearch_interception"], + premium_user=False, + config_file_path=".", + litellm_settings={ + "websearch_interception_params": { + "enabled_providers": ["bedrock"], + "search_tool_name": "tavily-search", + } + }, + callback_specific_params={}, + ) + + installed = [cb for cb in litellm.callbacks if isinstance(cb, WebSearchInterceptionLogger)] + assert len(installed) == 1 + assert installed[0].search_tool_name == "tavily-search" + assert installed[0].enabled_providers == ["bedrock"] + assert "websearch_interception" not in litellm.callbacks + + +def test_install_config_parameterized_callback_ignores_unrelated_callbacks(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [], raising=False) + + handled = install_config_parameterized_callback( + callback="langfuse", + litellm_settings={}, + callback_specific_params={}, + ) + + assert handled is False + assert litellm.callbacks == [] + + # --------------------------------------------------------------------------- # encrypt_callback_vars / decrypt_callback_vars # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 9a3702d2355..03baa8e9afc 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -1852,6 +1852,140 @@ def test_ProxyConfig__add_callbacks_from_db_config_processes_lists(monkeypatch): } +def _websearch_db_config(search_tool_name: str, enabled_providers: list[str]) -> Dict[str, Any]: + return { + "litellm_settings": { + "callbacks": ["websearch_interception"], + "websearch_interception_params": { + "enabled_providers": enabled_providers, + "search_tool_name": search_tool_name, + }, + } + } + + +def _reset_callback_lists(monkeypatch) -> None: + for list_name in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, [], raising=False) + + +def test_ProxyConfig__add_callbacks_from_db_config_installs_websearch_interception_logger( + monkeypatch, +): + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + _reset_callback_lists(monkeypatch) + pc = ProxyConfig() + + pc._add_callbacks_from_db_config(_websearch_db_config("tavily-search", ["bedrock", "vertex_ai"])) + + installed = [cb for cb in litellm.callbacks if isinstance(cb, WebSearchInterceptionLogger)] + snapshot = { + "raw_string_registered": "websearch_interception" in litellm.callbacks, + "instance_count": len(installed), + "search_tool_name": installed[0].search_tool_name if installed else None, + "enabled_providers": installed[0].enabled_providers if installed else None, + } + assert snapshot == { + "raw_string_registered": False, + "instance_count": 1, + "search_tool_name": "tavily-search", + "enabled_providers": ["bedrock", "vertex_ai"], + } + + +def test_ProxyConfig__add_callbacks_from_db_config_installs_compression_interception_logger( + monkeypatch, +): + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + _reset_callback_lists(monkeypatch) + pc = ProxyConfig() + + pc._add_callbacks_from_db_config( + { + "litellm_settings": { + "callbacks": ["compression_interception"], + "compression_interception_params": { + "enabled": True, + "compression_trigger": 123456, + }, + } + } + ) + + installed = [cb for cb in litellm.callbacks if isinstance(cb, CompressionInterceptionLogger)] + snapshot = { + "instance_count": len(installed), + "compression_trigger": installed[0].compression_trigger if installed else None, + } + assert snapshot == {"instance_count": 1, "compression_trigger": 123456} + + +def test_ProxyConfig__add_callbacks_from_db_config_reuses_logger_across_polls(monkeypatch): + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + _reset_callback_lists(monkeypatch) + pc = ProxyConfig() + config = _websearch_db_config("tavily-search", ["bedrock"]) + + pc._add_callbacks_from_db_config(config) + first = next(cb for cb in litellm.callbacks if isinstance(cb, WebSearchInterceptionLogger)) + for _ in range(3): + pc._add_callbacks_from_db_config(_websearch_db_config("tavily-search", ["bedrock"])) + + installed = [cb for cb in litellm.callbacks if isinstance(cb, WebSearchInterceptionLogger)] + snapshot = {"instance_count": len(installed), "same_instance": installed[:1] == [first]} + assert snapshot == {"instance_count": 1, "same_instance": True} + + +@pytest.mark.parametrize( + "changed_config, expected", + [ + ( + ("perplexity-search", ["bedrock"]), + {"search_tool_name": "perplexity-search", "enabled_providers": ["bedrock"]}, + ), + ( + ("tavily-search", ["bedrock", "openai"]), + {"search_tool_name": "tavily-search", "enabled_providers": ["bedrock", "openai"]}, + ), + ], +) +def test_ProxyConfig__add_callbacks_from_db_config_replaces_logger_when_params_change( + monkeypatch, changed_config, expected +): + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + _reset_callback_lists(monkeypatch) + pc = ProxyConfig() + + pc._add_callbacks_from_db_config(_websearch_db_config("tavily-search", ["bedrock"])) + pc._add_callbacks_from_db_config(_websearch_db_config(*changed_config)) + + installed = [cb for cb in litellm.callbacks if isinstance(cb, WebSearchInterceptionLogger)] + snapshot = { + "instance_count": len(installed), + "search_tool_name": installed[0].search_tool_name if installed else None, + "enabled_providers": installed[0].enabled_providers if installed else None, + } + assert snapshot == {"instance_count": 1, **expected} + + def test_ProxyConfig__add_callbacks_from_db_config_bad_config_raises(): pc = ProxyConfig() with pytest.raises(AttributeError): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f071c381916..72ca648d65f 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23350 + "limit": 23344 }, "LIT002": { "limit": 27256