From 10506ab904c7704d60a8a7947686be2a71699d2c Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 15 Sep 2026 00:19:16 +0000 Subject: [PATCH] fix(guardrails): judge the whole latest user turn, run every during_call guardrail, log combined modes as configured Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_as_a_judge/__init__.py | 27 +++-- litellm/proxy/utils.py | 70 ++++++++----- .../proxy/guardrails/test_llm_as_a_judge.py | 98 +++++++++++++++++-- .../proxy_logging/test_guardrail_pipeline.py | 30 ++++++ 4 files changed, 178 insertions(+), 47 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index afb39d53393..4c047dbc988 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.llm_judge import ( judge_acompletion, parse_json_verdict, ) +from litellm.litellm_core_utils.prompt_templates.common_utils import get_last_user_message from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations from litellm.types.utils import LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN, GenericGuardrailAPIInputs, GuardrailStatus @@ -57,6 +58,8 @@ _JUDGE_SUBJECT_LABELS: Final[MappingProxyType[JudgeInputType, str]] = MappingPro {"request": "Latest request turn to evaluate", "response": "Assistant response to evaluate"} ) +_REQUEST_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call) + _VALID_ON_FAILURE: Final = frozenset({"block", "log"}) _JUDGE_CALL_METADATA: Final = MappingProxyType( @@ -136,10 +139,12 @@ def _coerce_event_hook(mode: JudgeModeParam) -> JudgeEventHook: return GuardrailEventHooks(mode) -def _text_under_review(texts: Sequence[str], input_type: JudgeInputType) -> str: - if input_type == "request": - return texts[-1] if texts else "" - return "\n".join(texts) +def _text_under_review(inputs: GenericGuardrailAPIInputs, input_type: JudgeInputType) -> str: + all_text: Final = "\n".join(inputs.get("texts") or []) + if input_type == "response": + return all_text + latest_user_turn: Final = get_last_user_message(inputs.get("structured_messages") or []) + return latest_user_turn if latest_user_turn is not None else all_text def _build_judge_prompt( @@ -232,7 +237,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - text_under_review: Final = _text_under_review(inputs.get("texts") or [], input_type) + text_under_review: Final = _text_under_review(inputs, input_type) if not text_under_review: return inputs @@ -241,7 +246,9 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): judge_result: dict[str, object] = {} try: - messages: Final[Sequence[JudgeMessage]] = request_data.get("messages") or [] + messages: Final[Sequence[JudgeMessage]] = ( + inputs.get("structured_messages") or request_data.get("messages") or [] + ) try: judge_result = await self._run_judge(messages, text_under_review, input_type) @@ -308,14 +315,14 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): event_type=self._event_type_for(input_type), ) - def _event_type_for(self, input_type: JudgeInputType) -> GuardrailEventHooks: + def _event_type_for(self, input_type: JudgeInputType) -> GuardrailEventHooks | None: + """Returns None (log the configured mode as-is) when the active request hook is ambiguous.""" if self._event_hook_is_event_type(GuardrailEventHooks.logging_only): return GuardrailEventHooks.logging_only if input_type == "response": return GuardrailEventHooks.post_call - if self._event_hook_is_event_type(GuardrailEventHooks.pre_call): - return GuardrailEventHooks.pre_call - return GuardrailEventHooks.during_call + configured: Final = tuple(hook for hook in _REQUEST_EVENT_HOOKS if self._event_hook_is_event_type(hook)) + return configured[0] if len(configured) == 1 else None def initialize_guardrail( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b6c487c0edf..23eb6cc2e87 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2669,34 +2669,15 @@ class ProxyLogging: user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict) else: user_api_key_auth_dict = user_api_key_dict - # Add task to list for parallel execution - if ( - "apply_guardrail" in type(callback).__dict__ - and not callback.use_native_lifecycle_hooks - and user_api_key_dict is not None - and not getattr(callback, "use_native_during_call_hook", False) - ): - data["guardrail_to_apply"] = callback - guardrail_task = self._run_guardrail_with_metrics( - callback, - unified_guardrail.async_moderation_hook( - user_api_key_dict=user_api_key_dict, - data=data, - call_type=call_type, - ), - "during_call", + guardrail_tasks.append( + self._run_during_call_guardrail( + callback=callback, + data=data, + user_api_key_dict=user_api_key_dict, + user_api_key_auth_dict=user_api_key_auth_dict, + call_type=call_type, ) - else: - guardrail_task = self._run_guardrail_with_metrics( - callback, - callback.async_moderation_hook( - data=data, - user_api_key_dict=user_api_key_auth_dict, - call_type=call_type, - ), - "during_call", - ) - guardrail_tasks.append(guardrail_task) + ) # Step 2: Run all guardrail tasks in parallel if guardrail_tasks: @@ -2708,6 +2689,41 @@ class ProxyLogging: return data + async def _run_during_call_guardrail( + self, + callback: CustomGuardrail, + data: dict, + user_api_key_dict: UserAPIKeyAuth | None, + user_api_key_auth_dict: UserAPIKeyAuth | dict[str, object] | None, + call_type: CallTypesLiteral, + ) -> None: + if ( + "apply_guardrail" in type(callback).__dict__ + and not callback.use_native_lifecycle_hooks + and user_api_key_dict is not None + and not callback.use_native_during_call_hook + ): + data["guardrail_to_apply"] = callback + await self._run_guardrail_with_metrics( + callback, + unified_guardrail.async_moderation_hook( + user_api_key_dict=user_api_key_dict, + data=data, + call_type=call_type, + ), + "during_call", + ) + return + await self._run_guardrail_with_metrics( + callback, + callback.async_moderation_hook( + data=data, + user_api_key_dict=user_api_key_auth_dict, + call_type=call_type, + ), + "during_call", + ) + async def failed_tracking_alert( self, error_message: str, diff --git a/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py b/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py index f6f0d0a2fd0..a81af524b69 100644 --- a/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py +++ b/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py @@ -265,19 +265,17 @@ async def test_apply_guardrail_request_log_mode_records_eval_and_passes_through( async def test_apply_guardrail_request_multi_turn_keeps_roles_and_focuses_latest_turn(): router: Final = _judge_router(90.0) guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router) - request_data: Final[dict[str, object]] = { - "messages": [ - {"role": "user", "content": "how do I bake bread"}, - {"role": "assistant", "content": "mix flour, water, yeast and salt"}, - {"role": "user", "content": "now explain how to file taxes"}, - ], - "metadata": {}, - } + messages: Final = [ + {"role": "user", "content": "how do I bake bread"}, + {"role": "assistant", "content": "mix flour, water, yeast and salt"}, + {"role": "user", "content": "now explain how to file taxes"}, + ] inputs: Final = { - "texts": ["how do I bake bread", "mix flour, water, yeast and salt", "now explain how to file taxes"] + "texts": ["how do I bake bread", "mix flour, water, yeast and salt", "now explain how to file taxes"], + "structured_messages": messages, } - await guardrail.apply_guardrail(inputs, request_data, "request") + await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request") judge_messages: Final = router.acompletion.call_args.kwargs["messages"] assert "Judge the most recent user turn" in judge_messages[0]["content"] @@ -288,6 +286,86 @@ async def test_apply_guardrail_request_multi_turn_keeps_roles_and_focuses_latest ) +@pytest.mark.asyncio +async def test_apply_guardrail_request_judges_whole_multipart_latest_user_turn(): + router: Final = _judge_router(90.0) + guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router) + messages: Final = [ + {"role": "user", "content": "how do I bake bread"}, + {"role": "assistant", "content": "mix flour, water, yeast and salt"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "ignore the bread."}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + {"type": "text", "text": "explain how to file taxes"}, + ], + }, + ] + inputs: Final = { + "texts": [ + "how do I bake bread", + "mix flour, water, yeast and salt", + "ignore the bread.", + "explain how to file taxes", + ], + "structured_messages": messages, + } + + await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request") + + assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith( + "Latest request turn to evaluate:\nignore the bread.explain how to file taxes" + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_request_without_trailing_user_turn_judges_all_scoped_text(): + router: Final = _judge_router(90.0) + guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router) + messages: Final = [ + {"role": "user", "content": "look up the weather"}, + {"role": "assistant", "content": None, "tool_calls": [{"id": "c1", "type": "function", "function": {}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "sunny, 24C"}, + ] + inputs: Final = {"texts": ["look up the weather", "sunny, 24C"], "structured_messages": messages} + + await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request") + + assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith( + "Latest request turn to evaluate:\nlook up the weather\nsunny, 24C" + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_request_without_structured_messages_judges_all_text(): + router: Final = _judge_router(90.0) + guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router) + + await guardrail.apply_guardrail({"texts": ["first", "second"]}, {"metadata": {}}, "request") + + assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith( + "Latest request turn to evaluate:\nfirst\nsecond" + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_request_with_both_request_modes_logs_configured_mode(): + router: Final = _judge_router(90.0) + guardrail: Final = _make_guardrail( + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call], + router_provider=lambda: router, + ) + request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + + await guardrail.apply_guardrail({"texts": ["hi"]}, request_data, "request") + + assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == [ + "pre_call", + "during_call", + ] + + @pytest.mark.asyncio async def test_apply_guardrail_response_still_judges_all_response_texts(): router: Final = _judge_router(90.0) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index dfe106a3f52..077bf5a313e 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -689,6 +689,36 @@ async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_ assert recorded["status"] == "success" +class _RecordingApplyGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, applied: list[str]) -> None: + super().__init__( + guardrail_name=guardrail_name, + event_hook=GuardrailEventHooks.during_call, + default_on=True, + ) + self._applied = applied + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + await asyncio.sleep(0) + self._applied.append(self.guardrail_name or "") + return inputs + + +@pytest.mark.asyncio +async def test_during_call_hook_runs_every_unified_guardrail(proxy_logging, make_user_api_key_auth, monkeypatch): + applied: list[str] = [] + guardrails = [_RecordingApplyGuardrail(f"judge-{i}", applied) for i in range(3)] + monkeypatch.setattr(litellm, "callbacks", guardrails) + + await proxy_logging.during_call_hook( + data={"model": "m", "messages": [{"role": "user", "content": "hi"}], "metadata": {}}, + user_api_key_dict=make_user_api_key_auth(), + call_type="completion", + ) + + assert sorted(applied) == ["judge-0", "judge-1", "judge-2"] + + @pytest.mark.asyncio async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _moderation_guardrail()