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 93f5262f6ec..6b96c780c1c 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 @@ -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) 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 89362521c7f..d3ee723e7f8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/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 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 index 3bb63b01fef..504a07fd304 100644 --- 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 @@ -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]