Make LLM judge streaming guardrails safe

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

litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py

@@ -154,6 +154,10 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
             messages=judge_messages,
             response_format={"type": "json_object"},
             temperature=0,
+            metadata={
+                "user_api_key_metadata": {"disable_global_guardrails": True},
+                "user_api_key_team_metadata": {"disable_global_guardrails": True},
+            },
         )
         raw = response.choices[0].message.content or "{}"  # type: ignore[union-attr]
         return json.loads(raw)

litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py

@@ -29,6 +29,32 @@ A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message)
 GUARDRAIL_NAME = "unified_llm_guardrails"
 
 
+def _get_http_exception_detail(e: HTTPException) -> dict:
+    detail = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)}
+    return detail
+
+
+def _get_openai_sse_error_chunk(e: HTTPException) -> str:
+    detail = _get_http_exception_detail(e)
+    message = detail.get("error", detail.get("message", str(e.detail)))
+    return (
+        "data: "
+        + json.dumps(
+            {
+                "error": {
+                    "message": message,
+                    "type": "guardrail_violation",
+                    "code": e.status_code,
+                    **{
+                        k: v for k, v in detail.items() if k not in ("error", "message")
+                    },
+                }
+            }
+        )
+        + "\n\n"
+    )
+
+
 def _get_a2a_request_id(
     responses_so_far: List[Any], request_data: dict
 ) -> Optional[str]:
@@ -509,11 +535,7 @@ class UnifiedLLMGuardrails(CustomLogger):
             except HTTPException as e:
                 if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
                     request_id = _get_a2a_request_id(responses_so_far, request_data)
-                    detail = (
-                        e.detail
-                        if isinstance(e.detail, dict)
-                        else {"message": str(e.detail)}
-                    )
+                    detail = _get_http_exception_detail(e)
                     error_chunk = (
                         json.dumps(
                             {
@@ -535,6 +557,9 @@ class UnifiedLLMGuardrails(CustomLogger):
                         + "\n"
                     )
                     yield error_chunk
+                elif buffer_streaming_response:
+                    yield _get_openai_sse_error_chunk(e)
+                    return
                 else:
                     raise
 

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_llm_as_a_judge_guardrail_hook.py

@@ -1,6 +1,7 @@
 from types import SimpleNamespace
 from typing import Any, AsyncGenerator
 
+import json
 import pytest
 from fastapi import HTTPException
 
@@ -51,9 +52,11 @@ async def test_llm_as_a_judge_uses_proxy_router_for_judge_model_alias(
     class MockRouter:
         def __init__(self) -> None:
             self.calls = 0
+            self.kwargs: dict[str, Any] = {}
 
         async def acompletion(self, **kwargs: Any) -> Any:
             self.calls += 1
+            self.kwargs = kwargs
             assert kwargs["model"] == "proxy-judge-alias"
             return SimpleNamespace(
                 choices=[
@@ -92,6 +95,10 @@ async def test_llm_as_a_judge_uses_proxy_router_for_judge_model_alias(
 
     assert exc_info.value.status_code == 422
     assert router.calls == 1
+    assert router.kwargs["metadata"] == {
+        "user_api_key_metadata": {"disable_global_guardrails": True},
+        "user_api_key_team_metadata": {"disable_global_guardrails": True},
+    }
 
 
 @pytest.mark.asyncio
@@ -155,24 +162,30 @@ async def test_llm_as_a_judge_blocks_selected_streaming_chat_completion(
         "messages": [{"role": "user", "content": "hello"}],
         "model": "gpt-4o-mini",
     }
-    emitted_chunks: list[ModelResponseStream] = []
+    emitted_chunks: list[Any] = []
 
     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)
+        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 == []
+    assert len(emitted_chunks) == 1
+    assert isinstance(emitted_chunks[0], str)
+    assert emitted_chunks[0].startswith("data: ")
+    error_payload = json.loads(emitted_chunks[0].removeprefix("data: ").strip())
+    assert (
+        error_payload["error"]["message"]
+        == "LLM judge rejected response: score below threshold"
+    )
+    assert "unsafe response" not in emitted_chunks[0]