From d4cbf197fb75a46a18f534799a8455dd32bcb786 Mon Sep 17 00:00:00 2001 From: mint Date: Wed, 4 Nov 2026 17:52:05 +0800 Subject: [PATCH] fix(guardrails): dispatch llm_as_a_judge judge calls in the caller's team scope --- .../llm_as_a_judge/__init__.py | 17 +++++- .../proxy/guardrails/test_llm_as_a_judge.py | 56 +++++++++++++++++++ 2 files changed, 70 insertions(+), 3 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 d319e34f9fd..8244ae72a52 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 @@ -14,6 +14,8 @@ 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 ( + attached_auth, + caller_team_for, default_router_provider, extract_text_from_content, judge_acompletion, @@ -29,6 +31,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.guardrails import Guardrail, LitellmParams from litellm.types.llms.openai import AllMessageValues + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.types.utils import StandardLoggingEvalInformation JudgeInputType = Literal["request", "response"] @@ -237,6 +240,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): messages: Sequence[JudgeMessage], text_under_review: str, input_type: JudgeInputType = "response", + auth: "UserAPIKeyAuth | None" = None, ) -> dict[str, object]: judge_messages: Final[list[AllMessageValues]] = [ {"role": "system", "content": JUDGE_SYSTEM_PROMPTS[input_type]}, @@ -245,13 +249,19 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): "content": _build_judge_prompt(self.criteria, messages, text_under_review, input_type), }, ] + router: Final = self._router_provider() + team_id: Final = caller_team_for(router, self.judge_model, auth) + judge_metadata: Final[Mapping[str, str]] = MappingProxyType( + {**_JUDGE_CALL_METADATA, **({} if team_id is None else {"user_api_key_team_id": team_id})} + ) response: Final = await judge_acompletion( - self._router_provider(), + router, self.judge_model, judge_messages, + team_id=team_id, response_format={"type": "json_object"}, temperature=0, - metadata=dict(_JUDGE_CALL_METADATA), + metadata=dict(judge_metadata), ) raw: Final = response.choices[0].message.content or "{}" return _parse_judge_verdict(raw) @@ -278,9 +288,10 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): messages: Final[Sequence[JudgeMessage]] = ( inputs.get("structured_messages") or request_data.get("messages") or [] ) + auth: Final = attached_auth(request_data) try: - judge_result = await self._run_judge(messages, text_under_review, input_type) + judge_result = await self._run_judge(messages, text_under_review, input_type, auth) except Exception as judge_err: verbose_logger.warning( "llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err diff --git a/tests/unit/proxy/guardrails/test_llm_as_a_judge.py b/tests/unit/proxy/guardrails/test_llm_as_a_judge.py index 6e00958eba4..a1fc3577667 100644 --- a/tests/unit/proxy/guardrails/test_llm_as_a_judge.py +++ b/tests/unit/proxy/guardrails/test_llm_as_a_judge.py @@ -9,6 +9,7 @@ from fastapi import HTTPException import litellm from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import ( LLMAsAJudgeGuardrail, _build_judge_prompt, @@ -675,6 +676,61 @@ async def test_judge_routes_through_router_for_router_served_model( mock_sdk_completion.assert_not_called() +_TEAM_JUDGE_MODEL_LIST = [ + { + "model_name": "model_name_team1_uuid", + "litellm_params": {"model": "hosted_vllm/Qwen/Qwen3.6-27B", "api_key": "vllm-key"}, + "model_info": {"team_id": "team1", "team_public_model_name": "team-judge"}, + } +] + + +@pytest.mark.asyncio +@patch("litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge.litellm.acompletion", new_callable=AsyncMock) +async def test_judge_team_public_alias_runs_under_caller_team(mock_sdk_completion): + """A judge model exposed only as a team's public name serves only that team. The + caller's attached auth must resolve the team and stamp it onto the dispatch so the + router arm can actually serve the alias; without the stamp the call fell through to + the SDK with the raw name and raised 'LLM Provider NOT provided' (#27767).""" + router = _real_router(_TEAM_JUDGE_MODEL_LIST) + guardrail = _make_guardrail(judge_model="team-judge", router_provider=lambda: router) + inputs = {"texts": ["good response"]} + request_data: dict = { + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"user_api_key_auth": UserAPIKeyAuth(team_id="team1")}, + } + + result = await guardrail.apply_guardrail(inputs, request_data, "response") + + assert result is inputs + router.acompletion.assert_awaited_once() + call_kwargs = router.acompletion.await_args.kwargs + assert call_kwargs["model"] == "team-judge" + assert call_kwargs["metadata"]["user_api_key_team_id"] == "team1" + mock_sdk_completion.assert_not_called() + + +@pytest.mark.asyncio +@patch("litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge.litellm.acompletion", new_callable=AsyncMock) +async def test_judge_team_public_alias_without_team_caller_falls_back_to_sdk(mock_sdk_completion): + """With no caller team, a team-public alias serves nobody through the router, so + the SDK receives the raw alias. LiteLLM only serves provider-qualified public + names, which is the exact shape of the reported 'LLM Provider NOT provided' + failure, so dispatch must stay on the SDK arm for SDK-resolvable names and never + consult the router with a team the caller does not have.""" + router = _real_router(_TEAM_JUDGE_MODEL_LIST) + guardrail = _make_guardrail(judge_model="team-judge", router_provider=lambda: router) + inputs = {"texts": ["good response"]} + request_data: dict = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + + result = await guardrail.apply_guardrail(inputs, request_data, "response") + + assert result is inputs + router.acompletion.assert_not_called() + mock_sdk_completion.assert_awaited_once() + assert mock_sdk_completion.await_args.kwargs["model"] == "team-judge" + + @pytest.mark.asyncio @patch("litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge.litellm.acompletion", new_callable=AsyncMock) async def test_judge_call_falls_back_to_sdk_when_model_not_in_router(mock_sdk_completion):