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 4c047dbc988..46fccc126e7 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 @@ -58,7 +58,12 @@ _JUDGE_SUBJECT_LABELS: Final[MappingProxyType[JudgeInputType, str]] = MappingPro {"request": "Latest request turn to evaluate", "response": "Assistant response to evaluate"} ) -_REQUEST_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call) +_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"}) @@ -316,12 +321,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): ) def _event_type_for(self, input_type: JudgeInputType) -> GuardrailEventHooks | None: - """Returns None (log the configured mode as-is) when the active request hook is ambiguous.""" - if self._event_hook_is_event_type(GuardrailEventHooks.logging_only): - return GuardrailEventHooks.logging_only - if input_type == "response": - return GuardrailEventHooks.post_call - configured: Final = tuple(hook for hook in _REQUEST_EVENT_HOOKS if self._event_hook_is_event_type(hook)) + 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 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 a81af524b69..1646f730984 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 @@ -350,19 +350,25 @@ async def test_apply_guardrail_request_without_structured_messages_judges_all_te @pytest.mark.asyncio -async def test_apply_guardrail_request_with_both_request_modes_logs_configured_mode(): +@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=[GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call], - router_provider=lambda: router, - ) + 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, "request") + await guardrail.apply_guardrail({"texts": ["hi"]}, request_data, input_type) assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == [ - "pre_call", - "during_call", + mode.value for mode in modes ]