Fix LLM judge guardrail execution in playground

Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>

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)

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

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 == []