From f04734507a14aba076db795678d5fa598ca387d0 Mon Sep 17 00:00:00 2001 From: zayedu Date: Wed, 24 Jun 2026 21:06:39 -0700 Subject: [PATCH] fix(proxy): dispatch prompt injection llm api check --- litellm/proxy/utils.py | 41 +++++-- .../proxy_logging/test_during_call_hook.py | 102 +++++++++++++++++- 2 files changed, 131 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 755602cbcc0..36e3456f9cf 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -376,6 +376,7 @@ class _CallbackCapabilities: has_iterator_override: bool = False has_streaming_chunk_override: bool = False has_guardrail: bool = False + has_during_call_hook: bool = False has_pre_call_override: bool = False # Tuple[(resolved_callback, "override" | "apply_guardrail"), ...] # Ordered the same as ``litellm.callbacks``; used to build the streaming @@ -387,6 +388,14 @@ class _CallbackCapabilities: resolved_callbacks: Tuple[Any, ...] = field(default_factory=tuple) +def _has_native_during_call_hook(callback: CustomLogger) -> bool: + base_hook = CustomLogger.async_moderation_hook + callback_hook = getattr(type(callback), "async_moderation_hook", base_hook) + return getattr(callback_hook, "__func__", callback_hook) is not getattr( + base_hook, "__func__", base_hook + ) + + class ProxyLogging: """ Logging/Custom Handlers for proxy. @@ -1685,6 +1694,7 @@ class ProxyLogging: has_iterator_override = False has_streaming_chunk_override = False has_guardrail = False + has_during_call_hook = False has_pre_call_override = False iterator_overrides: List[Tuple[Any, str]] = [] # (callback, kind) resolved_callbacks: List[Any] = [] @@ -1706,6 +1716,9 @@ class ProxyLogging: continue if isinstance(resolved, CustomGuardrail): has_guardrail = True + has_during_call_hook = True + elif _has_native_during_call_hook(resolved): + has_during_call_hook = True # Use the same leaf-class ``__dict__`` check as the other hook # capabilities: only callbacks that actually override the hook # contribute to the flag. Setting this for every ``CustomLogger`` @@ -1746,6 +1759,7 @@ class ProxyLogging: or any(kind == "apply_guardrail" for _, kind in iterator_overrides), has_streaming_chunk_override=has_streaming_chunk_override, has_guardrail=has_guardrail, + has_during_call_hook=has_during_call_hook, has_pre_call_override=has_pre_call_override, iterator_overrides=tuple(iterator_overrides), resolved_callbacks=tuple(resolved_callbacks), @@ -1794,7 +1808,7 @@ class ProxyLogging: @staticmethod def has_during_call_guardrails() -> bool: - return ProxyLogging._callback_capabilities().has_guardrail + return ProxyLogging._callback_capabilities().has_during_call_hook async def during_call_hook( self, @@ -1803,18 +1817,19 @@ class ProxyLogging: call_type: CallTypesLiteral, ): """ - Runs the CustomGuardrail's async_moderation_hook() in parallel + Runs during-call async_moderation_hook() callbacks in parallel """ - # Fast path: skip the entire guardrail scan when no CustomGuardrail - # callbacks are registered. Saves per-request iteration over + # Fast path: skip the entire guardrail scan when no during-call + # moderation callbacks are registered. Saves per-request iteration over # ``litellm.callbacks`` plus an ``asyncio.gather([])`` round trip on - # deployments with no guardrails configured. - if not ProxyLogging._callback_capabilities().has_guardrail: + # deployments with no during-call checks configured. + caps = ProxyLogging._callback_capabilities() + if not caps.has_during_call_hook: return data # Step 1: Collect all guardrail tasks to run in parallel guardrail_tasks = [] - for callback in litellm.callbacks: + for callback in caps.resolved_callbacks: if isinstance(callback, CustomGuardrail): ################################################################ # Check if guardrail should be run for GuardrailEventHooks.during_call hook @@ -1871,6 +1886,18 @@ class ProxyLogging: ), ) guardrail_tasks.append(guardrail_task) + elif _has_native_during_call_hook(callback): + if call_type == CallTypes.call_mcp_tool.value: + continue + guardrail_task = self._run_guardrail_task_with_enrichment( + callback, + callback.async_moderation_hook( + data=data, + user_api_key_dict=user_api_key_dict, # type: ignore + call_type=call_type, # type: ignore + ), + ) + guardrail_tasks.append(guardrail_task) # Step 2: Run all guardrail tasks in parallel if guardrail_tasks: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py index 3c5d879c2dc..fa002aa8910 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py @@ -2,13 +2,17 @@ from __future__ import annotations -from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock import pytest import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import LiteLLMPromptInjectionParams +from litellm.proxy.hooks.prompt_injection_detection import ( + _OPTIONAL_PromptInjectionDetection, +) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks @@ -31,8 +35,27 @@ def _make_guardrail(name="g1", should_run=True, response=None): return cb +class _NativeDuringCallLogger(CustomLogger): + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + return data + + +def test_callback_capabilities_detect_native_during_call_hook(monkeypatch): + callback = _NativeDuringCallLogger() + monkeypatch.setattr(litellm, "callbacks", [callback]) + + caps = ProxyLogging._callback_capabilities() + + assert caps.has_during_call_hook is True + assert caps.has_guardrail is False + assert caps.resolved_callbacks == (callback,) + assert ProxyLogging.has_during_call_guardrails() is True + + @pytest.mark.asyncio -async def test_during_call_hook_no_guardrail_fast_path_returns_data(proxy_logging, make_user_api_key_auth, mock_callbacks_disabled): +async def test_during_call_hook_no_guardrail_fast_path_returns_data( + proxy_logging, make_user_api_key_auth, mock_callbacks_disabled +): data = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} out = await proxy_logging.during_call_hook( data=data, @@ -43,7 +66,9 @@ async def test_during_call_hook_no_guardrail_fast_path_returns_data(proxy_loggin @pytest.mark.asyncio -async def test_during_call_hook_runs_guardrails_in_parallel(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_during_call_hook_runs_guardrails_in_parallel( + proxy_logging, make_user_api_key_auth, monkeypatch +): g1 = _make_guardrail("a") g2 = _make_guardrail("b") monkeypatch.setattr(litellm, "callbacks", [g1, g2]) @@ -62,7 +87,72 @@ async def test_during_call_hook_runs_guardrails_in_parallel(proxy_logging, make_ @pytest.mark.asyncio -async def test_during_call_hook_guardrail_skipped_when_should_not_run(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_during_call_hook_runs_prompt_injection_llm_api_check( + proxy_logging, make_user_api_key_auth, monkeypatch +): + prompt_injection_detection = _OPTIONAL_PromptInjectionDetection( + prompt_injection_params=LiteLLMPromptInjectionParams( + llm_api_check=True, + llm_api_name="prompt-check", + llm_api_system_prompt="Return SAFE or UNSAFE", + llm_api_fail_call_string="UNSAFE", + ) + ) + router = MagicMock() + router.model_names = ["prompt-check"] + router.acompletion = AsyncMock(return_value=None) + prompt_injection_detection.update_environment(router=router) + monkeypatch.setattr(litellm, "callbacks", [prompt_injection_detection]) + + data = { + "model": "m", + "messages": [{"role": "user", "content": "hello"}], + "temperature": 0.1, + } + out = await proxy_logging.during_call_hook( + data=data, + user_api_key_dict=make_user_api_key_auth(), + call_type="completion", + ) + + assert out is data + router.acompletion.assert_awaited_once() + assert router.acompletion.await_args.kwargs["model"] == "prompt-check" + + +@pytest.mark.asyncio +async def test_during_call_hook_skips_prompt_injection_for_mcp_calls( + proxy_logging, make_user_api_key_auth, monkeypatch +): + prompt_injection_detection = _OPTIONAL_PromptInjectionDetection( + prompt_injection_params=LiteLLMPromptInjectionParams( + llm_api_check=True, + llm_api_name="prompt-check", + llm_api_system_prompt="Return SAFE or UNSAFE", + llm_api_fail_call_string="UNSAFE", + ) + ) + router = MagicMock() + router.model_names = ["prompt-check"] + router.acompletion = AsyncMock(return_value=None) + prompt_injection_detection.update_environment(router=router) + monkeypatch.setattr(litellm, "callbacks", [prompt_injection_detection]) + + data = {"model": "m", "name": "tool", "arguments": {"query": "hello"}} + out = await proxy_logging.during_call_hook( + data=data, + user_api_key_dict=make_user_api_key_auth(), + call_type="call_mcp_tool", + ) + + assert out is data + router.acompletion.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_during_call_hook_guardrail_skipped_when_should_not_run( + proxy_logging, make_user_api_key_auth, monkeypatch +): g = _make_guardrail("g", should_run=False) monkeypatch.setattr(litellm, "callbacks", [g]) await proxy_logging.during_call_hook( @@ -74,7 +164,9 @@ async def test_during_call_hook_guardrail_skipped_when_should_not_run(proxy_logg @pytest.mark.asyncio -async def test_during_call_hook_guardrail_error_raises(proxy_logging, make_user_api_key_auth, monkeypatch): +async def test_during_call_hook_guardrail_error_raises( + proxy_logging, make_user_api_key_auth, monkeypatch +): g = _make_guardrail("bad") g.async_moderation_hook = AsyncMock(side_effect=RuntimeError("blocked")) monkeypatch.setattr(litellm, "callbacks", [g])