mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge pull request #41128 from BerriAI/litellm_llm_judge_pre_call
feat(guardrails): support pre_call and during_call modes for llm_as_a_judge
This commit is contained in:
yucheng-berri
GitHub
commit
41eb2dbfeb
6 changed files with 520 additions and 84 deletions
|
|
@ -1,14 +1,17 @@
|
|||
"""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 collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeVar
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.llm_judge import (
|
||||
default_router_provider,
|
||||
|
|
@ -16,8 +19,9 @@ from litellm.litellm_core_utils.llm_judge import (
|
|||
judge_acompletion,
|
||||
parse_json_verdict,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import get_last_user_message
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations
|
||||
from litellm.types.utils import LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN, GenericGuardrailAPIInputs, GuardrailStatus
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Router
|
||||
|
|
@ -26,18 +30,65 @@ 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.
|
||||
For each criterion, assign a score from 0 to 100 and provide concise reasoning.
|
||||
JudgeInputType = Literal["request", "response"]
|
||||
JudgeEventHook = GuardrailEventHooks | list[GuardrailEventHooks] | Mode
|
||||
JudgeModeParam = str | list[str] | Mode | GuardrailEventHooks | list[GuardrailEventHooks] | None
|
||||
|
||||
_JUDGE_SYSTEM_PROMPT_TEMPLATE: Final = """You are a quality judge. Evaluate the {subject} against the criteria provided.
|
||||
{focus}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="request",
|
||||
focus="Judge the most recent user turn; treat earlier turns in the conversation only as context.\n",
|
||||
),
|
||||
"response": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="assistant's response", focus=""),
|
||||
}
|
||||
)
|
||||
|
||||
_JUDGE_SUBJECT_LABELS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType(
|
||||
{"request": "Latest request turn to evaluate", "response": "Assistant response to evaluate"}
|
||||
)
|
||||
|
||||
_LIFECYCLE_HOOKS: Final[MappingProxyType[JudgeInputType, tuple[GuardrailEventHooks, ...]]] = MappingProxyType(
|
||||
{
|
||||
"request": (GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.logging_only),
|
||||
"response": (GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only),
|
||||
}
|
||||
)
|
||||
|
||||
_VALID_ON_FAILURE: Final = frozenset({"block", "log"})
|
||||
|
||||
_JUDGE_CALL_METADATA: Final = MappingProxyType(
|
||||
{INTERNAL_CALL_ORIGIN_METADATA_KEY: LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN}
|
||||
)
|
||||
|
||||
|
||||
class _LoggedCallParams(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
metadata: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
def _is_logged_judge_call(data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
|
||||
"""logging_only is the only event whose ``data`` is the SDK's model_call_details rather than the client body."""
|
||||
if event_type is not GuardrailEventHooks.logging_only:
|
||||
return False
|
||||
try:
|
||||
params: Final = _LoggedCallParams.model_validate(data.get("litellm_params") or {})
|
||||
except ValidationError:
|
||||
return False
|
||||
return (params.metadata or {}).get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN
|
||||
|
||||
|
||||
_default_router_provider: Final = default_router_provider
|
||||
_parse_judge_verdict: Final = parse_json_verdict
|
||||
_extract_text_from_content: Final = extract_text_from_content
|
||||
|
|
@ -86,10 +137,29 @@ def _get_litellm_param(
|
|||
return default
|
||||
|
||||
|
||||
def _coerce_event_hook(mode: JudgeModeParam) -> JudgeEventHook:
|
||||
if mode is None:
|
||||
return GuardrailEventHooks.post_call
|
||||
if isinstance(mode, Mode):
|
||||
return mode
|
||||
if isinstance(mode, list):
|
||||
return [GuardrailEventHooks(hook) for hook in mode]
|
||||
return GuardrailEventHooks(mode)
|
||||
|
||||
|
||||
def _text_under_review(inputs: GenericGuardrailAPIInputs, input_type: JudgeInputType) -> str:
|
||||
all_text: Final = "\n".join(inputs.get("texts") or [])
|
||||
if input_type == "response":
|
||||
return all_text
|
||||
latest_user_turn: Final = get_last_user_message(inputs.get("structured_messages") or [])
|
||||
return latest_user_turn if latest_user_turn is not None else all_text
|
||||
|
||||
|
||||
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 +169,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 conversation or 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,
|
||||
|
|
@ -116,22 +187,15 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
criteria: Sequence[JudgeCriterion],
|
||||
overall_threshold: float = 80.0,
|
||||
on_failure: Literal["block", "log"] = "block",
|
||||
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None,
|
||||
event_hook: JudgeModeParam = None,
|
||||
default_on: bool = False,
|
||||
router_provider: "Callable[[], Router | None] | None" = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
_event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None
|
||||
if event_hook is not None:
|
||||
if isinstance(event_hook, list):
|
||||
_event_hook = [GuardrailEventHooks(h) if isinstance(h, str) else h for h in event_hook]
|
||||
else:
|
||||
_event_hook = GuardrailEventHooks(event_hook) if isinstance(event_hook, str) else event_hook
|
||||
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
supported_event_hooks=list(self.get_supported_event_hooks()),
|
||||
event_hook=_event_hook or GuardrailEventHooks.post_call,
|
||||
event_hook=_coerce_event_hook(event_hook),
|
||||
default_on=default_on,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
@ -143,18 +207,24 @@ 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]
|
||||
|
||||
def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
|
||||
if _is_logged_judge_call(data, event_type):
|
||||
return False
|
||||
return super().should_run_guardrail(data, event_type)
|
||||
|
||||
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(
|
||||
|
|
@ -163,6 +233,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
judge_messages,
|
||||
response_format={"type": "json_object"},
|
||||
temperature=0,
|
||||
metadata=dict(_JUDGE_CALL_METADATA),
|
||||
)
|
||||
raw: Final = response.choices[0].message.content or "{}"
|
||||
return _parse_judge_verdict(raw)
|
||||
|
|
@ -174,13 +245,8 @@ 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 = _text_under_review(inputs, input_type)
|
||||
if not text_under_review:
|
||||
return inputs
|
||||
|
||||
start_time: Final = datetime.now()
|
||||
|
|
@ -188,10 +254,12 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
|
|||
judge_result: dict[str, object] = {}
|
||||
|
||||
try:
|
||||
messages: Final[Sequence[JudgeMessage]] = request_data.get("messages") or []
|
||||
messages: Final[Sequence[JudgeMessage]] = (
|
||||
inputs.get("structured_messages") or 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 +298,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 +320,13 @@ 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 | None:
|
||||
configured: Final = tuple(hook for hook in _LIFECYCLE_HOOKS[input_type] if self._event_hook_is_event_type(hook))
|
||||
return configured[0] if len(configured) == 1 else None
|
||||
|
||||
|
||||
def initialize_guardrail(
|
||||
litellm_params: "LitellmParams",
|
||||
|
|
@ -282,10 +354,7 @@ def initialize_guardrail(
|
|||
|
||||
overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0))
|
||||
|
||||
mode: Final[str | None] = _get_litellm_param(litellm_params, guardrail, "mode", None)
|
||||
event_hook: GuardrailEventHooks | None = None
|
||||
if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}:
|
||||
event_hook = GuardrailEventHooks(mode)
|
||||
mode: Final[JudgeModeParam] = _get_litellm_param(litellm_params, guardrail, "mode", None)
|
||||
|
||||
instance: Final = LLMAsAJudgeGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
|
|
@ -293,7 +362,7 @@ def initialize_guardrail(
|
|||
criteria=criteria,
|
||||
overall_threshold=overall_threshold,
|
||||
on_failure=on_failure,
|
||||
event_hook=event_hook,
|
||||
event_hook=mode,
|
||||
default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)),
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(instance)
|
||||
|
|
|
|||
|
|
@ -2669,34 +2669,15 @@ class ProxyLogging:
|
|||
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict)
|
||||
else:
|
||||
user_api_key_auth_dict = user_api_key_dict
|
||||
# Add task to list for parallel execution
|
||||
if (
|
||||
"apply_guardrail" in type(callback).__dict__
|
||||
and not callback.use_native_lifecycle_hooks
|
||||
and user_api_key_dict is not None
|
||||
and not getattr(callback, "use_native_during_call_hook", False)
|
||||
):
|
||||
data["guardrail_to_apply"] = callback
|
||||
guardrail_task = self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
unified_guardrail.async_moderation_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
guardrail_tasks.append(
|
||||
self._run_during_call_guardrail(
|
||||
callback=callback,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
user_api_key_auth_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
)
|
||||
else:
|
||||
guardrail_task = self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
callback.async_moderation_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
guardrail_tasks.append(guardrail_task)
|
||||
)
|
||||
|
||||
# Step 2: Run all guardrail tasks in parallel
|
||||
if guardrail_tasks:
|
||||
|
|
@ -2708,6 +2689,41 @@ class ProxyLogging:
|
|||
|
||||
return data
|
||||
|
||||
async def _run_during_call_guardrail(
|
||||
self,
|
||||
callback: CustomGuardrail,
|
||||
data: dict[str, object], # mutable-ok: request payload dict, guardrail_to_apply is written in place
|
||||
user_api_key_dict: UserAPIKeyAuth | None,
|
||||
user_api_key_auth_dict: UserAPIKeyAuth | dict[str, object] | None,
|
||||
call_type: CallTypesLiteral,
|
||||
) -> None:
|
||||
if (
|
||||
"apply_guardrail" in type(callback).__dict__
|
||||
and not callback.use_native_lifecycle_hooks
|
||||
and user_api_key_dict is not None
|
||||
and not callback.use_native_during_call_hook
|
||||
):
|
||||
data["guardrail_to_apply"] = callback
|
||||
await self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
unified_guardrail.async_moderation_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
return
|
||||
await self._run_guardrail_with_metrics(
|
||||
callback,
|
||||
callback.async_moderation_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_auth_dict,
|
||||
call_type=call_type,
|
||||
),
|
||||
"during_call",
|
||||
)
|
||||
|
||||
async def failed_tracking_alert(
|
||||
self,
|
||||
error_message: str,
|
||||
|
|
|
|||
|
|
@ -2957,6 +2957,7 @@ InternalCallOrigin = Literal[
|
|||
"autorouter_classifier",
|
||||
"shadow_eval_router",
|
||||
"shadow_eval_judge",
|
||||
"llm_as_a_judge_guardrail",
|
||||
"background_response_cost_poll",
|
||||
]
|
||||
"""Which internal litellm feature originated a billed sub-call, so a spend log row
|
||||
|
|
@ -2965,6 +2966,7 @@ records that it is not traffic the caller sent."""
|
|||
AUTOROUTER_CLASSIFIER_CALL_ORIGIN: Final[InternalCallOrigin] = "autorouter_classifier"
|
||||
SHADOW_EVAL_ROUTER_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_router"
|
||||
SHADOW_EVAL_JUDGE_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_judge"
|
||||
LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN: Final[InternalCallOrigin] = "llm_as_a_judge_guardrail"
|
||||
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_cost_poll"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,14 @@
|
|||
"""Unit tests for the LLM-as-a-Judge guardrail hook."""
|
||||
|
||||
import json
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
|
||||
LLMAsAJudgeGuardrail,
|
||||
_build_judge_prompt,
|
||||
|
|
@ -13,7 +16,8 @@ from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
|
|||
_parse_judge_verdict,
|
||||
initialize_guardrail,
|
||||
)
|
||||
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.utils import LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
|
|
@ -136,17 +140,314 @@ def test_initialize_guardrail_invalid_on_failure():
|
|||
initialize_guardrail(lp, g)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "runs_pre_call", "runs_post_call"),
|
||||
[
|
||||
("pre_call", True, False),
|
||||
(["pre_call", "post_call"], True, True),
|
||||
(Mode(tags={"judge": ["pre_call"]}, default="post_call"), True, False),
|
||||
(None, False, True),
|
||||
],
|
||||
ids=["scalar", "list", "tagged", "missing"],
|
||||
)
|
||||
def test_initialize_guardrail_preserves_every_mode_shape(
|
||||
mode: str | list[str] | Mode | None,
|
||||
runs_pre_call: bool,
|
||||
runs_post_call: bool,
|
||||
):
|
||||
lp: Final = _make_litellm_params(mode=mode)
|
||||
instance: Final = initialize_guardrail(lp, _make_guardrail_dict())
|
||||
request_data: Final[dict[str, object]] = {"metadata": {"guardrails": ["g"], "tags": ["judge"]}}
|
||||
premium: Final = patch("litellm.proxy.proxy_server.premium_user", True) # test-quality-ok: no seam for Mode tags
|
||||
try:
|
||||
with premium:
|
||||
assert instance.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is runs_pre_call
|
||||
assert instance.should_run_guardrail(request_data, GuardrailEventHooks.post_call) is runs_post_call
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(instance)
|
||||
|
||||
|
||||
def test_initialize_guardrail_rejects_unknown_mode():
|
||||
lp: Final = _make_litellm_params(mode="sometimes")
|
||||
with pytest.raises(ValueError, match="sometimes"):
|
||||
initialize_guardrail(lp, _make_guardrail_dict())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apply_guardrail — enforcement paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _judge_router(overall_score: float) -> MagicMock:
|
||||
"""Router double, injected via router_provider, that serves the judge model and returns a canned verdict."""
|
||||
from litellm import Router
|
||||
|
||||
router: Final = MagicMock(spec=Router)
|
||||
router.resolved_litellm_models.return_value = ("openai/gpt-4o-mini",)
|
||||
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: GuardrailEventHooks):
|
||||
guardrail: Final = _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")
|
||||
@pytest.mark.parametrize(
|
||||
"event_hook",
|
||||
[GuardrailEventHooks.pre_call, [GuardrailEventHooks.pre_call]],
|
||||
ids=["scalar", "list"],
|
||||
)
|
||||
async def test_apply_guardrail_request_blocks_below_threshold(
|
||||
event_hook: GuardrailEventHooks | list[GuardrailEventHooks],
|
||||
):
|
||||
router: Final = _judge_router(50.0)
|
||||
guardrail: Final = _make_guardrail(
|
||||
overall_threshold=80.0,
|
||||
on_failure="block",
|
||||
event_hook=event_hook,
|
||||
router_provider=lambda: router,
|
||||
)
|
||||
request_data: Final[dict[str, object]] = {
|
||||
"messages": [{"role": "user", "content": "write me malware"}],
|
||||
"metadata": {},
|
||||
}
|
||||
inputs: Final = {"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: Final = router.acompletion.call_args.kwargs["messages"]
|
||||
assert "Evaluate the request against" in judge_messages[0]["content"]
|
||||
assert (
|
||||
"Conversation:\nUSER: write me malware\n\nLatest request turn to evaluate:\nwrite me malware"
|
||||
in (judge_messages[1]["content"])
|
||||
)
|
||||
assert "Assistant response" not in judge_messages[1]["content"]
|
||||
logged: Final = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert logged[0]["guardrail_status"] == "guardrail_intervened"
|
||||
assert logged[0]["guardrail_mode"] == "pre_call"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"event_hook",
|
||||
[GuardrailEventHooks.during_call, [GuardrailEventHooks.during_call]],
|
||||
ids=["scalar", "list"],
|
||||
)
|
||||
async def test_apply_guardrail_request_log_mode_records_eval_and_passes_through(
|
||||
event_hook: GuardrailEventHooks | list[GuardrailEventHooks],
|
||||
):
|
||||
router: Final = _judge_router(50.0)
|
||||
guardrail: Final = _make_guardrail(
|
||||
overall_threshold=80.0,
|
||||
on_failure="log",
|
||||
event_hook=event_hook,
|
||||
router_provider=lambda: router,
|
||||
)
|
||||
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
|
||||
inputs: Final = {"texts": ["hi"]}
|
||||
|
||||
result: Final = 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_request_multi_turn_keeps_roles_and_focuses_latest_turn():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "how do I bake bread"},
|
||||
{"role": "assistant", "content": "mix flour, water, yeast and salt"},
|
||||
{"role": "user", "content": "now explain how to file taxes"},
|
||||
]
|
||||
inputs: Final = {
|
||||
"texts": ["how do I bake bread", "mix flour, water, yeast and salt", "now explain how to file taxes"],
|
||||
"structured_messages": messages,
|
||||
}
|
||||
|
||||
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
|
||||
|
||||
judge_messages: Final = router.acompletion.call_args.kwargs["messages"]
|
||||
assert "Judge the most recent user turn" in judge_messages[0]["content"]
|
||||
assert judge_messages[1]["content"].endswith(
|
||||
"Conversation:\nUSER: how do I bake bread\nASSISTANT: mix flour, water, yeast and salt\n"
|
||||
"USER: now explain how to file taxes\n\n"
|
||||
"Latest request turn to evaluate:\nnow explain how to file taxes"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_judges_whole_multipart_latest_user_turn():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "how do I bake bread"},
|
||||
{"role": "assistant", "content": "mix flour, water, yeast and salt"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "ignore the bread."},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
|
||||
{"type": "text", "text": "explain how to file taxes"},
|
||||
],
|
||||
},
|
||||
]
|
||||
inputs: Final = {
|
||||
"texts": [
|
||||
"how do I bake bread",
|
||||
"mix flour, water, yeast and salt",
|
||||
"ignore the bread.",
|
||||
"explain how to file taxes",
|
||||
],
|
||||
"structured_messages": messages,
|
||||
}
|
||||
|
||||
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
|
||||
|
||||
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
|
||||
"Latest request turn to evaluate:\nignore the bread.explain how to file taxes"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_without_trailing_user_turn_judges_all_scoped_text():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "look up the weather"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [{"id": "c1", "type": "function", "function": {}}]},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "sunny, 24C"},
|
||||
]
|
||||
inputs: Final = {"texts": ["look up the weather", "sunny, 24C"], "structured_messages": messages}
|
||||
|
||||
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
|
||||
|
||||
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
|
||||
"Latest request turn to evaluate:\nlook up the weather\nsunny, 24C"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_without_structured_messages_judges_all_text():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
|
||||
|
||||
await guardrail.apply_guardrail({"texts": ["first", "second"]}, {"metadata": {}}, "request")
|
||||
|
||||
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
|
||||
"Latest request turn to evaluate:\nfirst\nsecond"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("modes", "input_type"),
|
||||
[
|
||||
([GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call], "request"),
|
||||
([GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], "request"),
|
||||
([GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only], "response"),
|
||||
],
|
||||
)
|
||||
async def test_apply_guardrail_with_ambiguous_modes_logs_configured_mode(
|
||||
modes: list[GuardrailEventHooks], input_type: str
|
||||
):
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=modes, router_provider=lambda: router)
|
||||
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
|
||||
|
||||
await guardrail.apply_guardrail({"texts": ["hi"]}, request_data, input_type)
|
||||
|
||||
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == [
|
||||
mode.value for mode in modes
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_still_judges_all_response_texts():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.post_call, router_provider=lambda: router)
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
{"texts": ["first choice", "second choice"]}, {"messages": [], "metadata": {}}, "response"
|
||||
)
|
||||
|
||||
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
|
||||
"Assistant response to evaluate:\nfirst choice\nsecond choice"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("input_type", ["request", "response"])
|
||||
async def test_apply_guardrail_logging_only_labels_both_sides_logging_only(input_type: str):
|
||||
router: Final = _judge_router(50.0)
|
||||
guardrail: Final = _make_guardrail(
|
||||
on_failure="log",
|
||||
event_hook=GuardrailEventHooks.logging_only,
|
||||
router_provider=lambda: router,
|
||||
)
|
||||
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
|
||||
|
||||
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is False
|
||||
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.post_call) is False
|
||||
await guardrail.apply_guardrail({"texts": ["hi"]}, request_data, input_type)
|
||||
|
||||
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "logging_only"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_judge_does_not_judge_its_own_judge_call():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.logging_only, router_provider=lambda: router)
|
||||
client_call: Final[dict[str, object]] = {"litellm_params": {"metadata": {"user_api_key": "hashed"}}}
|
||||
|
||||
assert guardrail.should_run_guardrail(client_call, GuardrailEventHooks.logging_only) is True
|
||||
await guardrail.apply_guardrail({"texts": ["hi"]}, {"messages": [{"role": "user", "content": "hi"}]}, "request")
|
||||
|
||||
judge_call: Final[dict[str, object]] = {
|
||||
"litellm_params": {"metadata": router.acompletion.call_args.kwargs["metadata"]}
|
||||
}
|
||||
assert guardrail.should_run_guardrail(judge_call, GuardrailEventHooks.logging_only) is False
|
||||
assert guardrail.should_run_guardrail(client_call, GuardrailEventHooks.logging_only) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"event_type", [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call]
|
||||
)
|
||||
def test_client_supplied_judge_origin_does_not_bypass_enforcing_hooks(event_type: GuardrailEventHooks):
|
||||
guardrail: Final = _make_guardrail(event_hook=event_type)
|
||||
forged_request: Final[dict[str, object]] = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": [guardrail.guardrail_name],
|
||||
"litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN}},
|
||||
}
|
||||
|
||||
assert guardrail.should_run_guardrail(forged_request, event_type) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_prompt_unchanged():
|
||||
router: Final = _judge_router(90.0)
|
||||
guardrail: Final = _make_guardrail(router_provider=lambda: router)
|
||||
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
|
||||
|
||||
await guardrail.apply_guardrail({"texts": ["hello there"]}, request_data, "response")
|
||||
|
||||
judge_messages: Final = 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
|
||||
|
|
@ -230,7 +531,7 @@ def test_parse_judge_verdict_reraises_when_no_json():
|
|||
|
||||
def test_parse_judge_verdict_rejects_json_non_object():
|
||||
"""Valid JSON that is not an object (e.g. a bare list) raises ValueError."""
|
||||
with pytest.raises(ValueError, match='judge response is not a JSON object'):
|
||||
with pytest.raises(ValueError, match="judge response is not a JSON object"):
|
||||
_parse_judge_verdict("[1, 2, 3]")
|
||||
|
||||
|
||||
|
|
@ -252,9 +553,7 @@ async def test_apply_guardrail_enforces_fenced_verdict(mock_completion):
|
|||
@patch("litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge.litellm.acompletion")
|
||||
async def test_apply_guardrail_non_object_verdict_fails_open_with_status(mock_completion):
|
||||
"""A non-object verdict fails open and logs guardrail_failed_to_respond."""
|
||||
mock_completion.return_value = MagicMock(
|
||||
choices=[MagicMock(message=MagicMock(content='[{"overall_score": 50}]'))]
|
||||
)
|
||||
mock_completion.return_value = MagicMock(choices=[MagicMock(message=MagicMock(content='[{"overall_score": 50}]'))])
|
||||
guardrail = _make_guardrail(overall_threshold=80.0, on_failure="block", router_provider=lambda: None)
|
||||
inputs = {"texts": ["response"]}
|
||||
request_data: dict = {"messages": [], "metadata": {}}
|
||||
|
|
@ -314,7 +613,12 @@ def _real_router(model_list, **router_kwargs):
|
|||
"model_list, router_kwargs, judge_model",
|
||||
[
|
||||
(
|
||||
[{"model_name": "my-judge-alias", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
|
||||
[
|
||||
{
|
||||
"model_name": "my-judge-alias",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
|
||||
}
|
||||
],
|
||||
{},
|
||||
"my-judge-alias",
|
||||
),
|
||||
|
|
@ -324,12 +628,22 @@ def _real_router(model_list, **router_kwargs):
|
|||
"anthropic/claude-sonnet-4-6",
|
||||
),
|
||||
(
|
||||
[{"model_name": "backing-group", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
|
||||
[
|
||||
{
|
||||
"model_name": "backing-group",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
|
||||
}
|
||||
],
|
||||
{"model_group_alias": {"my-judge-alias": "backing-group"}},
|
||||
"my-judge-alias",
|
||||
),
|
||||
(
|
||||
[{"model_name": "backing-group", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
|
||||
[
|
||||
{
|
||||
"model_name": "backing-group",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
|
||||
}
|
||||
],
|
||||
{"model_group_alias": {"my-judge-alias": {"model": "backing-group", "hidden": True}}},
|
||||
"my-judge-alias",
|
||||
),
|
||||
|
|
@ -412,7 +726,12 @@ async def test_judge_resolves_router_lazily_per_call(mock_sdk_completion):
|
|||
mock_sdk_completion.assert_awaited_once()
|
||||
|
||||
holder["router"] = _real_router(
|
||||
[{"model_name": "my-judge-alias", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}]
|
||||
[
|
||||
{
|
||||
"model_name": "my-judge-alias",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
|
||||
}
|
||||
]
|
||||
)
|
||||
await guardrail.apply_guardrail({"texts": ["r"]}, {"messages": [], "metadata": {}}, "response")
|
||||
holder["router"].acompletion.assert_awaited_once()
|
||||
|
|
|
|||
|
|
@ -689,6 +689,36 @@ async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_
|
|||
assert recorded["status"] == "success"
|
||||
|
||||
|
||||
class _RecordingApplyGuardrail(CustomGuardrail):
|
||||
def __init__(self, guardrail_name: str, applied: list[str]) -> None:
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=GuardrailEventHooks.during_call,
|
||||
default_on=True,
|
||||
)
|
||||
self._applied = applied
|
||||
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
await asyncio.sleep(0)
|
||||
self._applied.append(self.guardrail_name or "")
|
||||
return inputs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_call_hook_runs_every_unified_guardrail(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
applied: list[str] = []
|
||||
guardrails = [_RecordingApplyGuardrail(f"judge-{i}", applied) for i in range(3)]
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
|
||||
await proxy_logging.during_call_hook(
|
||||
data={"model": "m", "messages": [{"role": "user", "content": "hi"}], "metadata": {}},
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert sorted(applied) == ["judge-0", "judge-1", "judge-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
cb = _moderation_guardrail()
|
||||
|
|
|
|||
|
|
@ -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