From 345298f3c90871a7099952cb08e1e19d396883ee Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:04:19 -0700 Subject: [PATCH] fix(policy_engine): warn when a poll cannot re-match the submitted model name and keep default_on pre_call claims --- .../proxy/policy_engine/response_retrieval.py | 21 +++++++ litellm/proxy/utils.py | 14 ++++- .../policy_engine/test_response_retrieval.py | 56 ++++++++++++++++++- .../proxy_logging/test_guardrail_pipeline.py | 24 ++++++++ 4 files changed, 111 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 8fc75e40070..d284c44397e 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -18,6 +18,7 @@ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.types.proxy.policy_engine import PolicyMatchContext from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline @@ -46,9 +47,29 @@ def _model_group_for_response_id(response_id: object, llm_router: "Router | None deployment: Final = llm_router.get_deployment(model_id) if deployment is None: return UngovernedRetrieval("deployment no longer in the router") + hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias) + if hidden_by is not None: + verbose_proxy_logger.warning( + "Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), " + "so a policy attached to the model name it was submitted as does not run on it", + response_id, + deployment.model_name, + hidden_by, + ) return deployment.model_name +def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None: + if "*" in model_group: + return "a wildcard deployment" + aliases: Final = tuple( + alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group + ) + if not aliases: + return None + return f"the target of model_group_alias {', '.join(aliases)}" + + def _retrieval_context( data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str ) -> PolicyMatchContext: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e653ca6b62c..37ca52d6fc9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -543,6 +543,16 @@ def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline" return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps) +def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]: + return frozenset( + callback.guardrail_name + for callback in litellm.callbacks + if isinstance(callback, CustomGuardrail) + and callback.guardrail_name is not None + and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + ) + + def _without_names( bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write slot: str, @@ -567,7 +577,9 @@ def _withdraw_deferred_claims( outside_by_policy: Final = MappingProxyType( {policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred} ) - running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(*outside_by_policy.values()) + running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union( + _guardrails_run_standalone_pre_call(data), *outside_by_policy.values() + ) withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside) withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere _, bucket = get_or_create_metadata_bucket(data) diff --git a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py index b704f666646..9dcbef888ce 100644 --- a/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py +++ b/tests/test_litellm/proxy/policy_engine/test_response_retrieval.py @@ -14,11 +14,14 @@ GOVERNED_MODEL_GROUP = "gpt-5.4-mini" GOVERNED_MODEL_ID = "deployment-governed" UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini" UNGOVERNED_MODEL_ID = "deployment-ungoverned" +WILDCARD_MODEL_GROUP = "openai/*" +WILDCARD_MODEL_ID = "deployment-wildcard" class FakeRouter: - def __init__(self, deployments: dict[str, Deployment]): + def __init__(self, deployments: dict[str, Deployment], model_group_alias: dict[str, object] | None = None): self._deployments = deployments + self.model_group_alias = model_group_alias or {} def get_deployment(self, model_id: str) -> Deployment | None: return self._deployments.get(model_id) @@ -32,12 +35,14 @@ def _deployment(model_group: str, model_id: str) -> Deployment: ) -def _router() -> FakeRouter: +def _router(model_group_alias: dict[str, object] | None = None) -> FakeRouter: return FakeRouter( { GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID), UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID), - } + WILDCARD_MODEL_ID: _deployment(WILDCARD_MODEL_GROUP, WILDCARD_MODEL_ID), + }, + model_group_alias, ) @@ -133,6 +138,51 @@ def test_already_attached_policy_is_not_attached_twice(policy_engine): assert data["litellm_metadata"]["applied_policies"] == ["response-governance"] +def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "the model name it was submitted as" in record.getMessage() + ] + + +def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns(policy_engine, caplog): + data = _retrieval_data(WILDCARD_MODEL_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router()) + + assert data == _retrieval_data(WILDCARD_MODEL_ID) + assert [ + "as model group openai/* (a wildcard deployment)" in message for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_aliased_model_group_still_attaches_its_own_policies_and_warns(policy_engine, caplog): + data = _retrieval_data(GOVERNED_MODEL_ID) + router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}}) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router) + + assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),) + assert [ + "(the target of model_group_alias gpt-mini, gpt-hidden)" in message + for message in _hidden_submit_model_warnings(caplog) + ] == [True] + + +def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model(policy_engine, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + attach_post_call_pipelines_to_retrieval( + data=_retrieval_data(GOVERNED_MODEL_ID), + user_api_key_dict=UserAPIKeyAuth(), + llm_router=_router({"other-alias": UNGOVERNED_MODEL_GROUP}), + ) + + assert _hidden_submit_model_warnings(caplog) == [] + + def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: return [ record.getMessage() diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 8da07f39fd8..7e0736d4e67 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1561,6 +1561,30 @@ async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"} +@pytest.mark.asyncio +async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call( + proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry +): + class DualStageGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return response + + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-post", event_hook=["pre_call", "post_call"], default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _claimed_post_call_pipeline_data("response-governance") + + await proxy_logging.post_call_success_hook( + data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth() + ) + + assert "applied_policies" not in data["metadata"] + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + + @pytest.mark.asyncio async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran( proxy_logging, make_user_api_key_auth, monkeypatch, clear_policy_registry