diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 26071cd878b..c616d9e8723 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -11,7 +11,7 @@ import sys import threading import time import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Mapping, Sequence from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart @@ -186,6 +186,7 @@ if TYPE_CHECKING: from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline Span = _Span | object else: @@ -408,6 +409,46 @@ def _exception_changes_request_flow(exc: BaseException) -> bool: return isinstance(exc, (SensitiveDataRouteException, ModifyResponseException)) +def _policy_state_metadata(data: Mapping[str, object]) -> Mapping[str, object]: + """ + Return the metadata bucket the policy engine wrote its pipeline state into. + + The route decides the bucket (``litellm_metadata`` for ``/v1/messages``, + responses, batches, files and bedrock, ``metadata`` everywhere else), and both + buckets can be present at once because callers send their own provider-facing + ``metadata`` (Claude Code sends ``metadata.user_id``) or their own + ``litellm_metadata``. Pipeline slots are stripped from caller input before the + policy engine runs, so whichever bucket carries them is the proxy's own write. + """ + return next( + ( + bucket + for bucket in (data.get("metadata"), data.get("litellm_metadata")) + if isinstance(bucket, dict) + and ("_guardrail_pipelines" in bucket or "_pipeline_managed_guardrails" in bucket) + ), + {}, + ) + + +def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]: + pipelines: Final = _policy_state_metadata(data).get("_guardrail_pipelines") + return ( + tuple(cast("Sequence[tuple[str, GuardrailPipeline]]", pipelines)) # cast-ok: the policy engine wrote the slot + if pipelines + else () + ) + + +def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[str]: + managed: Final = _policy_state_metadata(data).get("_pipeline_managed_guardrails") + return ( + frozenset(cast("Collection[str]", managed)) # cast-ok: the policy engine wrote these guardrail names + if managed + else frozenset() + ) + + def _prompt_block_text(block: object) -> str: if isinstance(block, str): return block @@ -1446,8 +1487,7 @@ class ProxyLogging: Returns the (possibly modified) data dict. """ - metadata: Final = data.get("metadata", data.get("litellm_metadata", {})) or {} - pipelines: Final = metadata.get("_guardrail_pipelines") + pipelines: Final = _policy_pipelines(data) if not pipelines: return data @@ -1631,8 +1671,7 @@ class ProxyLogging: ) # Get pipeline-managed guardrails to skip in normal loop - metadata: Final = data.get("metadata", data.get("litellm_metadata", {})) or {} - pipeline_managed: Final[set] = metadata.get("_pipeline_managed_guardrails", set()) + pipeline_managed: Final = _pipeline_managed_guardrail_names(data) caps: Final = ProxyLogging._callback_capabilities() # Skip the per-request callback walk entirely when nothing in 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 5c711fc6c34..7df39b0ef82 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 @@ -25,6 +25,10 @@ from litellm.integrations.custom_guardrail import ( from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.policy_engine.pipeline_types import ( + GuardrailPipeline, + PipelineStep, +) @pytest.fixture(autouse=True) @@ -350,6 +354,45 @@ async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_log assert out is data +@pytest.mark.parametrize( + ("policy_state_key", "caller_metadata_key", "call_type"), + [ + ("litellm_metadata", "metadata", "anthropic_messages"), + ("metadata", "litellm_metadata", "completion"), + ], +) +@pytest.mark.asyncio +async def test_maybe_execute_pipelines_finds_policy_state_when_caller_sends_own_metadata( + proxy_logging, make_user_api_key_auth, monkeypatch, policy_state_key, caller_metadata_key, call_type +): + """The route picks the bucket the policy engine writes to (``litellm_metadata`` on + /v1/messages, ``metadata`` on chat completions), and the caller can populate the other + one, e.g. Claude Code sending ``metadata.user_id``. The pipeline must still run and block.""" + + class BlockingGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise HTTPException(status_code=400, detail={"error": "blocked by pipeline"}) + + monkeypatch.setattr(litellm, "callbacks", [BlockingGuardrail(guardrail_name="gr-1")]) + pipeline = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-1", on_fail="block")]) + data = { + caller_metadata_key: {"user_id": "user_abc"}, + policy_state_key: {"_guardrail_pipelines": [("policy-1", pipeline)]}, + "messages": [], + "model": "m", + } + + with pytest.raises(HTTPException) as exc_info: + await proxy_logging._maybe_execute_pipelines( + data=data, + user_api_key_dict=make_user_api_key_auth(), + call_type=call_type, + event_hook="pre_call", + ) + assert exc_info.value.detail["error"] == "blocked by pipeline" + assert exc_info.value.detail["guardrail_name"] == "gr-1" + + @pytest.mark.asyncio async def test_maybe_execute_pipelines_blocks_on_block_terminal_action_raises( proxy_logging, make_user_api_key_auth, monkeypatch