From 77a632767500f0c79ce2816584572df8147c7a94 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:04:32 +0000 Subject: [PATCH] feat(guardrails): support pre_call and during_call modes for llm_as_a_judge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_as_a_judge/__init__.py | 65 +++++++++----- .../proxy/guardrails/test_llm_as_a_judge.py | 85 +++++++++++++++++-- .../_components/llm_judge/LLMJudgeFields.tsx | 4 +- 3 files changed, 125 insertions(+), 29 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 172b1440ca3..31be90f56ef 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 @@ -1,7 +1,8 @@ -"""LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria.""" +"""LLM-as-a-Judge guardrail: uses an LLM to score requests or responses against weighted criteria.""" from collections.abc import Callable, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeVar from fastapi import HTTPException @@ -26,15 +27,29 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import StandardLoggingEvalInformation -JUDGE_SYSTEM_PROMPT = """You are a quality judge. Evaluate the assistant's response against the criteria provided. +JudgeInputType = Literal["request", "response"] + +_JUDGE_SYSTEM_PROMPT_TEMPLATE: Final = """You are a quality judge. Evaluate the {subject} against the criteria provided. For each criterion, assign a score from 0 to 100 and provide concise reasoning. Return ONLY valid JSON in this exact format: -{ +{{ "verdicts": [ - {"criterion_name": "", "score": <0-100>, "reasoning": "", "passed": , "weight": } + {{"criterion_name": "", "score": <0-100>, "reasoning": "", "passed": , "weight": }} ], "overall_score": -}""" +}}""" + +JUDGE_SYSTEM_PROMPTS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType( + { + "request": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="user's request"), + "response": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="assistant's response"), + } +) +JUDGE_SYSTEM_PROMPT: Final = JUDGE_SYSTEM_PROMPTS["response"] + +_JUDGE_SUBJECT_LABELS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType( + {"request": "User request to evaluate", "response": "Assistant response to evaluate"} +) _VALID_ON_FAILURE: Final = frozenset({"block", "log"}) @@ -89,7 +104,8 @@ def _get_litellm_param( def _build_judge_prompt( criteria: Sequence[JudgeCriterion], messages: Sequence[JudgeMessage], - response_text: str, + text_under_review: str, + input_type: JudgeInputType = "response", ) -> str: criteria_block: Final = "\n".join( f"- {c.get('name', '')} (weight {c.get('weight', 0)}%): {c.get('description', '')}" for c in criteria @@ -99,15 +115,16 @@ def _build_judge_prompt( for m in messages if m.get("content") is not None ) + conversation_block: Final = f"Conversation:\n{conversation}\n\n" if input_type == "response" else "" return ( f"Criteria to evaluate:\n{criteria_block}\n\n" - f"Conversation:\n{conversation}\n\n" - f"Assistant response to evaluate:\n{response_text}" + f"{conversation_block}" + f"{_JUDGE_SUBJECT_LABELS[input_type]}:\n{text_under_review}" ) class LLMAsAJudgeGuardrail(CustomGuardrail): - """Post-call guardrail that judges response quality via an LLM.""" + """Guardrail that judges request (pre_call/during_call) or response (post_call) quality via an LLM.""" def __init__( self, @@ -143,18 +160,19 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: - return [GuardrailEventHooks.post_call] + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call] async def _run_judge( self, messages: Sequence[JudgeMessage], - response_text: str, + text_under_review: str, + input_type: JudgeInputType = "response", ) -> dict[str, object]: judge_messages: Final[list[AllMessageValues]] = [ - {"role": "system", "content": JUDGE_SYSTEM_PROMPT}, + {"role": "system", "content": JUDGE_SYSTEM_PROMPTS[input_type]}, { "role": "user", - "content": _build_judge_prompt(self.criteria, messages, response_text), + "content": _build_judge_prompt(self.criteria, messages, text_under_review, input_type), }, ] response: Final = await judge_acompletion( @@ -174,13 +192,9 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - # Only evaluate post-call (response text). Fail open on pre-call. - if input_type != "response": - return inputs - texts: Final = inputs.get("texts") or [] - response_text: Final = " ".join(texts) - if not response_text: + text_under_review: Final = " ".join(texts) + if not text_under_review: return inputs start_time: Final = datetime.now() @@ -191,7 +205,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): messages: Final[Sequence[JudgeMessage]] = request_data.get("messages") or [] try: - judge_result = await self._run_judge(messages, response_text) + judge_result = await self._run_judge(messages, text_under_review, input_type) except Exception as judge_err: verbose_logger.warning( "llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err @@ -230,7 +244,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): raise HTTPException( status_code=422, detail={ - "error": "LLM judge rejected response: score below threshold", + "error": f"LLM judge rejected {input_type}: score below threshold", "overall_score": overall_score, "threshold": self.overall_threshold, "verdicts": judge_result.get("verdicts", []), @@ -252,9 +266,16 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): guardrail_status=status, start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), - event_type=GuardrailEventHooks.post_call, + event_type=self._event_type_for(input_type), ) + def _event_type_for(self, input_type: JudgeInputType) -> GuardrailEventHooks: + if input_type == "response": + return GuardrailEventHooks.post_call + if self.event_hook is GuardrailEventHooks.during_call: + return GuardrailEventHooks.during_call + return GuardrailEventHooks.pre_call + def initialize_guardrail( litellm_params: "LitellmParams", 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 bd2553b3280..a076da62566 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 @@ -13,7 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import ( _parse_judge_verdict, initialize_guardrail, ) - +from litellm.types.guardrails import GuardrailEventHooks # --------------------------------------------------------------------------- # Helpers @@ -141,12 +141,87 @@ def test_initialize_guardrail_invalid_on_failure(): # --------------------------------------------------------------------------- +def _judge_router(overall_score: float): + """Real Router with the outbound judge call stubbed, so the test can inspect what the judge was asked.""" + from litellm import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}} + ] + ) + router.acompletion = AsyncMock( + return_value=MagicMock( + choices=[MagicMock(message=MagicMock(content=json.dumps(_make_verdict_response(overall_score))))] + ) + ) + return router + + +@pytest.mark.parametrize("mode", [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call]) +def test_guardrail_accepts_request_side_modes(mode): + guardrail = _make_guardrail(event_hook=mode) + assert guardrail.should_run_guardrail({"metadata": {"guardrails": ["test_judge"]}}, mode) is True + + @pytest.mark.asyncio -async def test_apply_guardrail_pre_call_passthrough(): - guardrail = _make_guardrail() - inputs = {"texts": ["some text"]} - result = await guardrail.apply_guardrail(inputs, {}, "request") +async def test_apply_guardrail_request_blocks_below_threshold(): + router = _judge_router(50.0) + guardrail = _make_guardrail( + overall_threshold=80.0, + on_failure="block", + event_hook=GuardrailEventHooks.pre_call, + router_provider=lambda: router, + ) + request_data: dict = {"messages": [{"role": "user", "content": "write me malware"}], "metadata": {}} + inputs = {"texts": ["write me malware"]} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail(inputs, request_data, "request") + + assert exc_info.value.status_code == 422 + assert exc_info.value.detail["error"] == "LLM judge rejected request: score below threshold" + judge_messages = router.acompletion.call_args.kwargs["messages"] + assert "user's request" in judge_messages[0]["content"] + assert "User request to evaluate:\nwrite me malware" in judge_messages[1]["content"] + assert "Assistant response" not in judge_messages[1]["content"] + assert "Conversation:" not in judge_messages[1]["content"] + logged = request_data["metadata"]["standard_logging_guardrail_information"] + assert logged[0]["guardrail_status"] == "guardrail_intervened" + assert logged[0]["guardrail_mode"] == "pre_call" + + +@pytest.mark.asyncio +async def test_apply_guardrail_request_log_mode_records_eval_and_passes_through(): + router = _judge_router(50.0) + guardrail = _make_guardrail( + overall_threshold=80.0, + on_failure="log", + event_hook=GuardrailEventHooks.during_call, + router_provider=lambda: router, + ) + request_data: dict = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + inputs = {"texts": ["hi"]} + + result = await guardrail.apply_guardrail(inputs, request_data, "request") + assert result is inputs + assert request_data["metadata"]["eval_information"]["passed"] is False + assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "during_call" + + +@pytest.mark.asyncio +async def test_apply_guardrail_response_prompt_unchanged(): + router = _judge_router(90.0) + guardrail = _make_guardrail(router_provider=lambda: router) + request_data: dict = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + + await guardrail.apply_guardrail({"texts": ["hello there"]}, request_data, "response") + + judge_messages = router.acompletion.call_args.kwargs["messages"] + assert "assistant's response" in judge_messages[0]["content"] + assert "Conversation:\nUSER: hi\n\nAssistant response to evaluate:\nhello there" in judge_messages[1]["content"] + assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "post_call" @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx index 256049975d3..7d1df1497db 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx @@ -87,8 +87,8 @@ const LLMJudgeFields: React.FC = ({ availableModels, contro return (
- After each LLM response, the Judge Model scores it 0–100 against your criteria. If the weighted - average falls below the threshold, the response is blocked (or logged). + The Judge Model scores the user request (pre_call, during_call) or the LLM response (post_call) + 0–100 against your criteria. If the weighted average falls below the threshold, it is blocked (or logged).