From dfcea2c1866630313ec3083794a922ff6971583a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:16:00 -0700 Subject: [PATCH] fix(policy_engine): execute post_call guardrail pipelines on responses --- litellm/proxy/utils.py | 47 +++++- .../proxy_logging/test_guardrail_pipeline.py | 152 +++++++++++++++++- 2 files changed, 196 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d880b529727..bd8ba6ae9f0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -455,6 +455,30 @@ def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[s ) +def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: + if data.get("stream") is not True: + return + post_call_policies: Final = tuple( + policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" + ) + if not post_call_policies: + return + raise HTTPException( + status_code=400, + detail={ + "error": { + "message": ( + "Policies with post_call guardrail pipelines cannot govern streaming responses yet: " + f"{', '.join(post_call_policies)}. Retry with stream=false, or move these policies' output " + "guardrails from pipeline steps to guardrails.add, which scans streamed output." + ), + "type": "guardrail_pipeline_error", + "policies": list(post_call_policies), + } + }, + ) + + def _prompt_block_text(block: object) -> str: if isinstance(block, str): return block @@ -1578,6 +1602,7 @@ class ProxyLogging: call_type: str, event_hook: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + response: LLMResponseTypes | None = None, ) -> dict: """ Execute guardrail pipelines if any are configured for this request. @@ -1596,6 +1621,8 @@ class ProxyLogging: if not pipelines: return data + step_input: Final = {**data, "response": response} if response is not None else data + for policy_name, pipeline in pipelines: if pipeline.mode != event_hook: continue @@ -1603,7 +1630,7 @@ class ProxyLogging: result: PipelineExecutionResult = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode=pipeline.mode, - data=data, + data=step_input, user_api_key_dict=user_api_key_dict, call_type=call_type, policy_name=policy_name, @@ -1614,6 +1641,7 @@ class ProxyLogging: result=result, data=data, policy_name=policy_name, + original_response=response, ) return data @@ -1623,14 +1651,18 @@ class ProxyLogging: result: PipelineExecutionResult, data: dict, policy_name: str, + original_response: LLMResponseTypes | None = None, ) -> dict: """ Handle a PipelineExecutionResult — allow, block, or modify_response. Returns data dict if allowed, raises on block/modify_response. + ``original_response`` is set on the post_call path, where allowed + modifications land on the response object in place, so the request + payload (already sent upstream) is left untouched. """ if result.terminal_action == "allow": - if result.modified_data is not None: + if result.modified_data is not None and original_response is None: data.update(result.modified_data) return data @@ -1671,6 +1703,7 @@ class ProxyLogging: request_data=data, guardrail_name=f"pipeline:{policy_name}", detection_info=None, + original_response=original_response, ) return data @@ -1786,6 +1819,8 @@ class ProxyLogging: ) try: + _raise_for_streaming_post_call_pipelines(data) + # Execute guardrail pipelines before the normal callback loop data = await self._maybe_execute_pipelines( data=data, @@ -2774,6 +2809,14 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks + await self._maybe_execute_pipelines( + data=data, + user_api_key_dict=user_api_key_dict, + call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", + event_hook="post_call", + response=response, + ) + guardrail_callbacks: Final[list[CustomGuardrail]] = [] other_callbacks: Final[list[CustomLogger]] = [] try: 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 e99e34d65d4..b7e86b68fe4 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 @@ -23,7 +23,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, @@ -865,3 +865,153 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p hook_kwargs = logging_obj.async_get_chat_completion_prompt.await_args.kwargs assert hook_kwargs["messages"] == [{"role": "user", "content": "Who are you?"}] assert hook_kwargs["prompt_spec"] is prompt_spec + + +# --------------------------------------------------------------------------- +# post_call pipeline execution (LIT-6410) +# --------------------------------------------------------------------------- + + +def _post_call_pipeline_data(guardrail: str = "gr-post", **extra: Any) -> Dict[str, Any]: + pipeline = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("response-governance", pipeline)], + "_pipeline_managed_guardrails": {guardrail}, + }, + **extra, + } + + +@pytest.mark.asyncio +async def test_post_call_success_hook_runs_post_call_pipeline_and_reraises_block( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class OutputBlockingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + monkeypatch.setattr( + litellm, + "callbacks", + [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = litellm.ModelResponse() + + with pytest.raises(HTTPException) as info: + await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert info.value.detail["error"] == "output blocked" + assert seen["response"] is response + + +@pytest.mark.asyncio +async def test_post_call_pipeline_pass_runs_once_and_leaves_request_data_untouched( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + seen["response"] = response + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [RecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = litellm.ModelResponse() + + out = await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert out is response + assert seen["response"] is response + assert seen["count"] == 1 + assert "response" not in data + assert "guardrails" not in data["metadata"] + + +def test_handle_pipeline_result_modify_response_carries_original_response(): + result = MagicMock() + result.terminal_action = "modify_response" + result.modify_response_message = "filtered" + response = litellm.ModelResponse() + + with pytest.raises(ModifyResponseException) as info: + ProxyLogging._handle_pipeline_result( + result=result, data={"model": "m"}, policy_name="p", original_response=response + ) + + assert info.value.original_response is response + + +def test_handle_pipeline_result_allow_discards_modifications_on_post_call(): + data = {"a": 1, "metadata": {"guardrails": ["other"]}} + result = MagicMock() + result.terminal_action = "allow" + result.modified_data = {"metadata": {"guardrails": ["gr-post"]}, "response": object()} + + out = ProxyLogging._handle_pipeline_result( + result=result, data=data, policy_name="p", original_response=litellm.ModelResponse() + ) + + assert out is data + assert data == {"a": 1, "metadata": {"guardrails": ["other"]}} + + +@pytest.mark.asyncio +async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", []) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["policies"] == ["response-governance"] + assert "stream=false" in info.value.detail["error"]["message"] + + +def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(): + post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) + + assert ( + _raise_for_streaming_post_call_pipelines( + {"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}} + ) + is None + ) + assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) is None + assert ( + _raise_for_streaming_post_call_pipelines( + {"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}} + ) + is None + ) + assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None