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 1693cf75882..863cff6827f 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 @@ -17,7 +17,7 @@ from litellm.litellm_core_utils.llm_judge import ( judge_acompletion, parse_json_verdict, ) -from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations +from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus if TYPE_CHECKING: @@ -28,6 +28,8 @@ if TYPE_CHECKING: from litellm.types.utils import StandardLoggingEvalInformation 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. For each criterion, assign a score from 0 to 100 and provide concise reasoning. @@ -41,13 +43,13 @@ Return ONLY valid JSON in this exact format: JUDGE_SYSTEM_PROMPTS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType( { - "request": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="user's request"), + "request": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="request"), "response": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="assistant's response"), } ) _JUDGE_SUBJECT_LABELS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType( - {"request": "User request to evaluate", "response": "Assistant response to evaluate"} + {"request": "Request text to evaluate", "response": "Assistant response to evaluate"} ) _VALID_ON_FAILURE: Final = frozenset({"block", "log"}) @@ -100,6 +102,16 @@ 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 _build_judge_prompt( criteria: Sequence[JudgeCriterion], messages: Sequence[JudgeMessage], @@ -132,22 +144,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, ) @@ -302,10 +307,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, @@ -313,7 +315,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) 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 dd42e137ba0..cab24514844 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 @@ -14,7 +14,7 @@ from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import ( _parse_judge_verdict, initialize_guardrail, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, Mode # --------------------------------------------------------------------------- # Helpers @@ -137,6 +137,37 @@ 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"], +) +@patch("litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge.litellm.logging_callback_manager") +def test_initialize_guardrail_preserves_every_mode_shape( + _mock_mgr: MagicMock, + 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"]}} + + 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 + + +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 # --------------------------------------------------------------------------- @@ -190,8 +221,8 @@ async def test_apply_guardrail_request_blocks_below_threshold( 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 "user's request" in judge_messages[0]["content"] - assert "User request to evaluate:\nwrite me malware" in judge_messages[1]["content"] + assert "Evaluate the request against" in judge_messages[0]["content"] + assert "Request text 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: Final = request_data["metadata"]["standard_logging_guardrail_information"]