From af66966aab52e8f561e6b9fe4eb36d01928fdf10 Mon Sep 17 00:00:00 2001 From: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com> Date: Wed, 6 May 2026 05:00:35 +0000 Subject: [PATCH] Fix LLM judge guardrail execution in playground Co-authored-by: ishaan-berri --- .../llm_as_a_judge/__init__.py | 14 +- .../unified_guardrail/unified_guardrail.py | 21 +++ .../test_llm_as_a_judge_guardrail_hook.py | 178 ++++++++++++++++++ 3 files changed, 212 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_as_a_judge_guardrail_hook.py 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 45bdd6dd09f..93f5262f6ec 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 @@ -13,6 +13,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus if TYPE_CHECKING: + from litellm import Router from litellm.types.guardrails import Guardrail, LitellmParams from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import StandardLoggingEvalInformation @@ -93,6 +94,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): criteria: List[Dict[str, Any]], overall_threshold: float = 80.0, on_failure: Literal["block", "log"] = "block", + llm_router: Optional["Router"] = None, event_hook: Optional[ Union[GuardrailEventHooks, List[GuardrailEventHooks]] ] = None, @@ -126,6 +128,9 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): self.criteria = criteria self.overall_threshold = overall_threshold self.on_failure = on_failure + self.llm_router = llm_router + # The judge needs the complete response before deciding whether to block. + self.buffer_streaming_response = True async def _run_judge( self, @@ -139,7 +144,12 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): "content": _build_judge_prompt(self.criteria, messages, response_text), }, ] - response = await litellm.acompletion( + completion_fn = ( + self.llm_router.acompletion + if self.llm_router is not None + else litellm.acompletion + ) + response = await completion_fn( model=self.judge_model, messages=judge_messages, response_format={"type": "json_object"}, @@ -244,6 +254,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): def initialize_guardrail( litellm_params: "LitellmParams", guardrail: "Guardrail", + llm_router: Optional["Router"] = None, ) -> LLMAsAJudgeGuardrail: guardrail_name = guardrail.get("guardrail_name") if not guardrail_name: @@ -286,6 +297,7 @@ def initialize_guardrail( criteria=criteria, overall_threshold=overall_threshold, on_failure=on_failure, + llm_router=llm_router, event_hook=event_hook, default_on=bool( _get_litellm_param(litellm_params, guardrail, "default_on", False) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index bc46beabc65..89362521c7f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -312,6 +312,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Get streaming configuration from guardrail or optional_params sampling_rate = 5 end_of_stream_only = False # If True, only apply guardrail at end of stream + buffer_streaming_response = False if guardrail_to_apply is not None: # Check direct attributes on guardrail first @@ -321,6 +322,11 @@ class UnifiedLLMGuardrails(CustomLogger): end_of_stream_only = getattr( guardrail_to_apply, "streaming_end_of_stream_only", end_of_stream_only ) + buffer_streaming_response = getattr( + guardrail_to_apply, + "buffer_streaming_response", + buffer_streaming_response, + ) # Also check guardrail_config dict if present guardrail_config = getattr(guardrail_to_apply, "guardrail_config", {}) @@ -331,6 +337,9 @@ class UnifiedLLMGuardrails(CustomLogger): end_of_stream_only = guardrail_config.get( "streaming_end_of_stream_only", end_of_stream_only ) + buffer_streaming_response = guardrail_config.get( + "buffer_streaming_response", buffer_streaming_response + ) # Also check optional_params as fallback sampling_rate = self.optional_params.get( @@ -339,6 +348,9 @@ class UnifiedLLMGuardrails(CustomLogger): end_of_stream_only = self.optional_params.get( "streaming_end_of_stream_only", end_of_stream_only ) + buffer_streaming_response = self.optional_params.get( + "buffer_streaming_response", buffer_streaming_response + ) if guardrail_to_apply is None: async for item in response: @@ -394,6 +406,11 @@ class UnifiedLLMGuardrails(CustomLogger): yield remaining_item return + # If this guardrail must make a final decision before any response is + # emitted, buffer chunks until the final post-call check completes. + if buffer_streaming_response: + continue + # If end_of_stream_only mode, yield chunks without processing if end_of_stream_only: yield item @@ -520,3 +537,7 @@ class UnifiedLLMGuardrails(CustomLogger): yield error_chunk else: raise + + if buffer_streaming_response: + for processed_item in responses_so_far: + yield processed_item diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_as_a_judge_guardrail_hook.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_as_a_judge_guardrail_hook.py new file mode 100644 index 00000000000..3bb63b01fef --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_as_a_judge_guardrail_hook.py @@ -0,0 +1,178 @@ +from types import SimpleNamespace +from typing import Any, AsyncGenerator + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import ( + LLMAsAJudgeGuardrail, +) +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, +) +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices + + +def _llm_judge_guardrail() -> LLMAsAJudgeGuardrail: + return LLMAsAJudgeGuardrail( + guardrail_name="judge-quality", + judge_model="gpt-4o-mini", + criteria=[ + { + "name": "quality", + "description": "Response must be high quality", + "weight": 100, + } + ], + overall_threshold=80, + on_failure="block", + ) + + +def test_llm_as_a_judge_selected_guardrail_matches_post_call_event() -> None: + guardrail = _llm_judge_guardrail() + + assert guardrail.should_run_guardrail( + data={"metadata": {"guardrails": ["judge-quality"]}}, + event_type=GuardrailEventHooks.post_call, + ) + + +@pytest.mark.asyncio +async def test_llm_as_a_judge_uses_proxy_router_for_judge_model_alias( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class MockRouter: + def __init__(self) -> None: + self.calls = 0 + + async def acompletion(self, **kwargs: Any) -> Any: + self.calls += 1 + assert kwargs["model"] == "proxy-judge-alias" + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace(content='{"overall_score": 10}') + ) + ] + ) + + async def direct_litellm_call_should_not_be_used(**kwargs: Any) -> Any: + raise RuntimeError("direct provider call cannot resolve proxy model alias") + + router = MockRouter() + guardrail = LLMAsAJudgeGuardrail( + guardrail_name="judge-quality", + judge_model="proxy-judge-alias", + criteria=[ + { + "name": "quality", + "description": "Response must be high quality", + "weight": 100, + } + ], + overall_threshold=80, + on_failure="block", + llm_router=router, + ) + monkeypatch.setattr(litellm, "acompletion", direct_litellm_call_should_not_be_used) + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["unsafe response"]}, + request_data={"messages": [{"role": "user", "content": "hello"}]}, + input_type="response", + ) + + assert exc_info.value.status_code == 422 + assert router.calls == 1 + + +@pytest.mark.asyncio +async def test_llm_as_a_judge_blocks_selected_streaming_chat_completion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail = _llm_judge_guardrail() + judge_calls = 0 + + async def mock_run_judge( + messages: list[dict[str, Any]], response_text: str + ) -> dict: + nonlocal judge_calls + judge_calls += 1 + assert messages == [{"role": "user", "content": "hello"}] + assert "unsafe response" in response_text + return { + "overall_score": 10, + "verdicts": [ + { + "criterion_name": "quality", + "score": 10, + "reasoning": "low quality", + "passed": False, + "weight": 100, + } + ], + } + + monkeypatch.setattr(guardrail, "_run_judge", mock_run_judge) + original_mappings = unified_module.endpoint_guardrail_translation_mappings + unified_module.endpoint_guardrail_translation_mappings = { + CallTypes.acompletion: unified_module.load_guardrail_translation_mappings()[ + CallTypes.acompletion + ], + } + + async def mock_stream() -> AsyncGenerator[ModelResponseStream, None]: + yield ModelResponseStream( + choices=[ + StreamingChoices( + delta=Delta(content="unsafe response", role="assistant"), + finish_reason=None, + ) + ], + model="gpt-4o-mini", + ) + yield ModelResponseStream( + choices=[ + StreamingChoices( + delta=Delta(content=None), + finish_reason="stop", + ) + ], + model="gpt-4o-mini", + ) + + request_data = { + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["judge-quality"]}, + "messages": [{"role": "user", "content": "hello"}], + "model": "gpt-4o-mini", + } + emitted_chunks: list[ModelResponseStream] = [] + + try: + with pytest.raises(HTTPException) as exc_info: + async for ( + chunk + ) in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key="test-key", + request_route="/v1/chat/completions", + ), + response=mock_stream(), + request_data=request_data, + ): + emitted_chunks.append(chunk) + finally: + unified_module.endpoint_guardrail_translation_mappings = original_mappings + + assert exc_info.value.status_code == 422 + assert judge_calls == 1 + assert emitted_chunks == []