diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index dd94286b4ce..2a3b0520ac5 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45322 + "limit": 45321 }, "reportUnknownLambdaType": { "limit": 113 diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index b18994afc99..c8e26556804 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -44,10 +44,67 @@ reset_color_code = "\033[0m" TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY = "_pillar_response_headers_trusted" if TYPE_CHECKING: + from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + ) + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + ConfigParameterizedLoggerClass = ( + type[CompressionInterceptionLogger] + | type[CodeInterpreterInterceptionLogger] + | type[WebSearchInterceptionLogger] + ) + _NO_CALLBACK_SPECIFIC_PARAMS: Mapping[str, Any] = MappingProxyType({}) +CONFIG_PARAMETERIZED_CALLBACKS: frozenset[str] = frozenset( + { + "compression_interception", + "code_interpreter_interception", + "websearch_interception", + } +) + + +def config_parameterized_logger_class(callback: str) -> "ConfigParameterizedLoggerClass | None": + """ + The logger class 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 one of these + names 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. The names it answers to are + ``CONFIG_PARAMETERIZED_CALLBACKS``. + """ + match callback: + case "compression_interception": + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + return CompressionInterceptionLogger + case "code_interpreter_interception": + from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + ) + + return CodeInterpreterInterceptionLogger + case "websearch_interception": + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + return WebSearchInterceptionLogger + case _: + return None def resolve_config_parameterized_callback( @@ -56,46 +113,18 @@ def resolve_config_parameterized_callback( callback_specific_params: Mapping[str, Any] | None = None, ) -> CustomLogger | None: """ - Build the ``CustomLogger`` for callbacks whose construction needs config params. + Build the ``CustomLogger`` for a config-parameterized callback from its 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. + Returns ``None`` when ``callback`` is not one of them. """ - callback_params = callback_specific_params or _NO_CALLBACK_SPECIFIC_PARAMS - match callback: - case "compression_interception": - from litellm.integrations.compression_interception.handler import ( - CompressionInterceptionLogger, - ) + logger_class = config_parameterized_logger_class(callback) + if logger_class is None: + return None - 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 + return logger_class.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params or _NO_CALLBACK_SPECIFIC_PARAMS, + ) def _callback_config_state(logger: CustomLogger) -> tuple[tuple[str, object], ...]: @@ -146,6 +175,29 @@ def install_config_parameterized_callback( return True +def uninstall_deconfigured_parameterized_callbacks(configured_callbacks: Iterable[object]) -> None: + """ + Drop config-parameterized loggers whose callback name is no longer configured. + + The DB-config path re-runs on every poll and only visits the names the config + still lists, so without this an operator who removes ``websearch_interception`` + from ``litellm_settings.callbacks`` would keep paying for interception until the + proxy restarts. Call it only with a callbacks list the config actually produced; + a failed config load must not be read as "nothing is configured". + """ + configured = frozenset(name for name in configured_callbacks if isinstance(name, str)) + for callback in CONFIG_PARAMETERIZED_CALLBACKS - configured: + logger_class = config_parameterized_logger_class(callback) + if logger_class is None: + continue + for installed in tuple(existing for existing in litellm.callbacks if isinstance(existing, logger_class)): + verbose_proxy_logger.info( + "%s no longer configured; removing its callback from litellm.callbacks", + callback, + ) + litellm.logging_callback_manager.remove_callback_from_all_lists(installed) + + def initialize_callbacks_on_proxy( value: Any, premium_user: bool, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a152513ca17..59bab43fbe6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -298,6 +298,7 @@ from litellm.proxy.common_request_processing import ( from litellm.proxy.common_utils.callback_utils import ( initialize_callbacks_on_proxy, install_config_parameterized_callback, + uninstall_deconfigured_parameterized_callbacks, ) from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers @@ -5588,6 +5589,7 @@ class ProxyConfig: event_types=["success", "failure"], existing_callbacks=litellm.callbacks, ) + uninstall_deconfigured_parameterized_callbacks(callbacks) def _encrypt_env_variables(self, environment_variables: dict, new_encryption_key: str | None = None) -> dict: """ 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 e41766f0c86..22e2804ed8c 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -10,8 +10,11 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm.proxy.common_utils.callback_utils import ( + CONFIG_PARAMETERIZED_CALLBACKS, add_policy_to_applied_policies_header, + config_parameterized_logger_class, decrypt_callback_vars, + uninstall_deconfigured_parameterized_callbacks, encrypt_callback_vars, get_logging_caching_headers, initialize_callbacks_on_proxy, @@ -315,6 +318,32 @@ def test_initialize_callbacks_on_proxy_instantiates_websearch_interception_with_ assert "websearch_interception" not in litellm.callbacks +def test_every_config_parameterized_callback_name_resolves_to_a_logger_class(): + resolved = {name: config_parameterized_logger_class(name) for name in CONFIG_PARAMETERIZED_CALLBACKS} + + assert all(cls is not None for cls in resolved.values()), resolved + assert len({id(cls) for cls in resolved.values()}) == len(CONFIG_PARAMETERIZED_CALLBACKS) + + +def test_uninstall_deconfigured_parameterized_callbacks_keeps_still_configured_logger( + monkeypatch, +): + from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, + ) + + monkeypatch.setattr(litellm, "callbacks", [], raising=False) + install_config_parameterized_callback( + callback="websearch_interception", + litellm_settings={"websearch_interception_params": {"search_tool_name": "tavily-search"}}, + ) + + uninstall_deconfigured_parameterized_callbacks(["websearch_interception", "langfuse"]) + + installed = [cb for cb in litellm.callbacks if isinstance(cb, WebSearchInterceptionLogger)] + assert len(installed) == 1 + + def test_install_config_parameterized_callback_ignores_unrelated_callbacks(monkeypatch): monkeypatch.setattr(litellm, "callbacks", [], raising=False) 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 03baa8e9afc..4ffedda8323 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -1986,6 +1986,65 @@ def test_ProxyConfig__add_callbacks_from_db_config_replaces_logger_when_params_c assert snapshot == {"instance_count": 1, **expected} +def test_ProxyConfig__add_callbacks_from_db_config_uninstalls_logger_when_callback_removed( + 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"])) + installed_before = [cb for cb in litellm.callbacks if isinstance(cb, WebSearchInterceptionLogger)] + + pc._add_callbacks_from_db_config({"litellm_settings": {"callbacks": ["some_other_callback"]}}) + + remaining_anywhere = [ + cb + for callback_list in ( + litellm.callbacks, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ) + for cb in callback_list + if isinstance(cb, WebSearchInterceptionLogger) + ] + snapshot = { + "installed_before": len(installed_before), + "remaining_anywhere": len(remaining_anywhere), + } + assert snapshot == {"installed_before": 1, "remaining_anywhere": 0} + + +@pytest.mark.parametrize( + "config_without_callbacks", + [ + {}, + {"litellm_settings": {}}, + {"litellm_settings": {"success_callback": ["langfuse"]}}, + ], +) +def test_ProxyConfig__add_callbacks_from_db_config_keeps_logger_when_config_omits_callbacks( + monkeypatch, config_without_callbacks +): + 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(config_without_callbacks) + + installed = [cb for cb in litellm.callbacks if isinstance(cb, WebSearchInterceptionLogger)] + assert len(installed) == 1 + + 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 72ca648d65f..a2e9c2ed45b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23344 + "limit": 23338 }, "LIT002": { "limit": 27256