mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
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>
This commit is contained in:
parent
0b3e56448f
commit
77a6327675
3 changed files with 125 additions and 29 deletions
|
|
@ -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": "<name>", "score": <0-100>, "reasoning": "<one sentence>", "passed": <true|false>, "weight": <weight>}
|
||||
{{"criterion_name": "<name>", "score": <0-100>, "reasoning": "<one sentence>", "passed": <true|false>, "weight": <weight>}}
|
||||
],
|
||||
"overall_score": <weighted average 0-100>
|
||||
}"""
|
||||
}}"""
|
||||
|
||||
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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -87,8 +87,8 @@ const LLMJudgeFields: React.FC<LLMJudgeFieldsProps> = ({ availableModels, contro
|
|||
return (
|
||||
<FieldGroup>
|
||||
<div className="rounded-md border border-success/20 bg-success/10 px-3.5 py-2.5 text-[13px] text-success">
|
||||
After each LLM response, the <strong>Judge Model</strong> scores it 0–100 against your criteria. If the weighted
|
||||
average falls below the threshold, the response is blocked (or logged).
|
||||
The <strong>Judge Model</strong> 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).
|
||||
</div>
|
||||
|
||||
<GuardrailField
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue