From 969ded2b54199b6ec94b688de66a0cc7afd0b2ae Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 30 Oct 2026 17:00:44 -0700 Subject: [PATCH] fix(guardrails): drop the response pipeline steps that never ran from x-litellm-applied-guardrails (#49264) * refactor(guardrails): record which policy pipeline steps and deployment-level hooks ran * fix(guardrails): name only the request pipeline steps that ran in x-litellm-applied-guardrails * fix(guardrails): drop the response pipeline steps that never ran from x-litellm-applied-guardrails * test(guardrails): integration cells for request pipeline steps in x-litellm-applied-guardrails * fix(guardrails): keep the agent cells on the shared proxy and shorten a long suppression line * fix(guardrails): keep request steps when a response pipeline blocks before a later request pipeline --- litellm/proxy/utils.py | 105 +++- .../test_guardrail_pipeline_applied_header.py | 37 ++ .../proxy_logging/test_guardrail_pipeline.py | 487 ++++++++++++++++++ 3 files changed, 627 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bd91e56cb01..af95a70261a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -746,6 +746,88 @@ def _request_steps_that_have_not_run(data: Mapping[str, object]) -> frozenset[st return frozenset(name for name in pre_call if _runs_only_in_pipelines(name, data, pre_call, frozenset())) +def _steps_that_ran(result: PipelineExecutionResult) -> frozenset[str]: + return frozenset(step.guardrail_name for step in result.step_results if step.outcome != "skip") + + +def _whole_response_log_waits_for_its_pipelines(data: Mapping[str, object], call_type: str) -> bool: + """ + The success log is parked until the response hooks finish, and the headers are built + after them. The LLM pass-through route builds its headers first, and the generateContent + adapter logs its nested completion before the response hooks run. + """ + logging_obj: Final = data.get("litellm_logging_obj") + return ( + call_type not in (CallTypes.allm_passthrough_route.value, CallTypes.agenerate_content.value) + and isinstance(logging_obj, Logging) + and logging_obj.enqueue_deferred_logging is not None + ) + + +def _stream_log_waits_for_its_pipelines(data: Mapping[str, object]) -> bool: + return getattr(data.get("litellm_logging_obj"), "_on_deferred_stream_complete", None) is not None + + +def _pipelines_settled_by( + pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult +) -> Sequence[tuple[str, "GuardrailPipeline"]]: + """This pipeline, plus every later response pipeline when its verdict ends the request before they run.""" + if result.terminal_action == "allow": + return pipelines[index : index + 1] + return ( + pipelines[index], + *((policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call"), + ) + + +def _steps_later_pipelines_still_run( + pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult +) -> frozenset[str]: + """The steps of the later response pipelines, which still run when this pipeline's verdict lets the request on.""" + if result.terminal_action != "allow": + return frozenset() + later: Final = tuple( + (policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call" + ) + return _pipeline_step_guardrail_names(later) + + +def _drop_response_steps_that_did_not_run( + data: dict[str, object], # mutable-ok: the request's applied guardrails are trimmed in place + pipelines: Sequence[tuple[str, "GuardrailPipeline"]], + steps_to_keep: frozenset[str], +) -> None: + """ + Drop from the applied guardrails the steps of these response pipelines outside ``steps_to_keep``, + unless their guardrail also runs outside pipelines. Every step is named up front, so a + log built before the pipeline finished, such as a cache hit's, still lists the steps + that ran. The list is trimmed in place, since a route's logging object may hold a + shallow copy of the metadata. + """ + pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") + post_call: Final = _pipeline_step_guardrail_names(pipelines) + not_run: Final = frozenset( + name for name in post_call - steps_to_keep if _runs_only_in_pipelines(name, data, pre_call, post_call) + ) + if not not_run: + return + for bucket_key in ("metadata", "litellm_metadata"): + _drop_applied_guardrails_in_place(data.get(bucket_key), not_run) + + +def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None: + if not isinstance(bucket, dict): + return + metadata: Final = cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed + applied: Final = metadata.get("applied_guardrails") + if not isinstance(applied, list): + return + entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names + entries[:] = [name for name in entries if name not in names] + if not entries: + metadata.pop("applied_guardrails") + + def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: resolved: Final = tuple( litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -2814,7 +2896,10 @@ class ProxyLogging: return data, None current_response = response # rebind-ok: chains each pipeline's replacement response into the next - for policy_name, pipeline in pipelines: + steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran + request: Final = cast(dict[str, object], data) # cast-ok: request payload, updated in place by every pipeline + log_waits: Final = event_hook == "post_call" and _whole_response_log_waits_for_its_pipelines(request, call_type) + for index, (policy_name, pipeline) in enumerate(pipelines): if pipeline.mode != event_hook: continue @@ -2830,6 +2915,14 @@ class ProxyLogging: raw_request_snapshot=raw_request_snapshot, ) + steps_that_ran = steps_that_ran | _steps_that_ran(result) + if log_waits: + _drop_response_steps_that_did_not_run( + request, + _pipelines_settled_by(pipelines, index, result), + steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result), + ) + data = self._handle_pipeline_result( result=result, data=data, @@ -4781,7 +4874,8 @@ class ProxyLogging: call_type, endpoint_translation = translation - for policy_name, pipeline in pipelines: + steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran + for index, (policy_name, pipeline) in enumerate(pipelines): result: PipelineExecutionResult = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode="post_call", @@ -4792,6 +4886,13 @@ class ProxyLogging: streaming_chunks=buffered, endpoint_translation=endpoint_translation, ) + steps_that_ran = steps_that_ran | _steps_that_ran(result) + if _stream_log_waits_for_its_pipelines(request_data): + _drop_response_steps_that_did_not_run( + request_data, + _pipelines_settled_by(pipelines, index, result), + steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result), + ) try: ProxyLogging._handle_pipeline_result( result, data=request_data, policy_name=policy_name, original_response=buffered diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index 8b40df163bb..ae8fd1c350b 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -546,6 +546,21 @@ def test_a_request_pipeline_that_blocks_at_its_first_step_names_only_that_step(r assert response.headers.get(_HEADER) == _REQUEST_BLOCKER, dict(response.headers) +def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_streamed_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + def test_a_request_pipeline_names_a_step_that_flagged_the_request_before_moving_on(rig: _Rig) -> None: response: Final = _ask_as_team_tagged(rig, _MOVES_PAST_FIRST) assert response.status_code == 200, response.text @@ -596,6 +611,17 @@ def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step_o assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step_on_every_endpoint( + rig: _Rig, path: str, stream: bool +) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=stream, path=path) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + def test_a_request_pipeline_names_a_step_whose_service_errored_before_moving_on(rig: _Rig) -> None: response: Final = _ask_as_team_tagged(rig, _ERRORS_PAST_FIRST) assert response.status_code == 200, response.text @@ -644,6 +670,15 @@ def _spans_naming_guardrails(rig: _Rig, since: int, team: str) -> tuple[Span, .. pytest.param(_STOPS_AFTER_FIRST, "/v1/chat/completions", False, [_REQUEST_FIRST], id="request-pipeline"), pytest.param(_STOPS_AFTER_FIRST, "/v1/chat/completions", True, [_REQUEST_FIRST], id="request-pipeline-stream"), pytest.param(_STOPS_AFTER_FIRST, "/v1/responses", True, [_REQUEST_FIRST], id="request-pipeline-responses"), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/chat/completions", False, [_RESPONSE_FIRST], id="response-pipeline" + ), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/chat/completions", True, [_RESPONSE_FIRST], id="response-pipeline-stream" + ), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/messages", True, [_RESPONSE_FIRST], id="response-pipeline-messages" + ), ], ) def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran( @@ -1009,6 +1044,7 @@ def agent(rig: _Rig) -> Iterator[str]: pytest.param( (_RESPONSE_MOVES_PAST_FIRST,), [_RESPONSE_BLOCKER, _RESPONSE_SECOND], id="first-step-flags-and-moves-on" ), + pytest.param((_RESPONSE_STOPS_AFTER_FIRST,), [_RESPONSE_FIRST], id="first-step-allows"), pytest.param( (_RESPONSE_STOPS_AFTER_FIRST, _RESPONSE_STOPS_AFTER_SECOND), [_RESPONSE_FIRST, _RESPONSE_SECOND], @@ -1174,6 +1210,7 @@ def test_a_concurrent_burst_names_only_each_requests_own_steps(rig: _Rig) -> Non { _STOPS_AFTER_FIRST: _REQUEST_FIRST, _RUNS_EVERY_STEP: f"{_REQUEST_FIRST},{_REQUEST_SECOND}", + _RESPONSE_STOPS_AFTER_FIRST: _RESPONSE_FIRST, } ) cases: Final = tuple(itertools.product(expected, _ENDPOINTS, (False, True))) * 2 diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 26f8f9e04b2..5fd40696224 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -12,6 +12,7 @@ from __future__ import annotations import asyncio import json from collections.abc import AsyncGenerator, Iterator +from contextlib import AbstractContextManager, nullcontext from copy import deepcopy import logging from typing import Any, Callable, Dict, Final, List @@ -30,6 +31,7 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( @@ -1206,6 +1208,21 @@ def _two_step_pipeline_data( } +def _logging_obj(*, log_parked: bool, call_type: str = "acompletion") -> LiteLLMLoggingObj: + logging_obj: Final = LiteLLMLoggingObj( + model="m", + messages=[], + stream=False, + call_type=call_type, + start_time=None, + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.defer_async_logging = True + logging_obj.enqueue_deferred_logging = (lambda: None) if log_parked else None + return logging_obj + + @pytest.mark.asyncio @pytest.mark.parametrize( ("first_on_pass", "expected_calls", "expected_header"), @@ -1343,6 +1360,403 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat assert set(get_logging_caching_headers(result)["x-litellm-applied-guardrails"].split(",")) == {"first", "second"} +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("log_parked", "call_type", "bucket", "expected_after_response"), + [ + pytest.param(True, "acompletion", "metadata", ["first"], id="success-log-waits-for-the-response-hooks"), + pytest.param(True, "acompletion", "litellm_metadata", ["first"], id="success-log-waits-litellm-metadata"), + pytest.param(False, "acompletion", "metadata", ["first", "second"], id="success-log-already-built"), + pytest.param( + True, "allm_passthrough_route", "metadata", ["first", "second"], id="pass-through-headers-already-out" + ), + pytest.param( + True, "agenerate_content", "metadata", ["first", "second"], id="generate-content-adapter-logged-its-call" + ), + ], +) +async def test_post_call_pipeline_drops_the_steps_that_did_not_run_only_while_nothing_has_read_them( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + log_parked: bool, + call_type: str, + bucket: str, + expected_after_response: list[str], +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + request: Final = _two_step_pipeline_data("post_call", "allow") + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + "model": request["model"], + "messages": request["messages"], + bucket: request["metadata"], + "litellm_logging_obj": _logging_obj(log_parked=log_parked, call_type=call_type), + }, + call_type="completion", + ) + assert data[bucket]["applied_guardrails"] == ["first", "second"] + logged_metadata: Final = data[bucket].copy() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 0) + assert data[bucket]["applied_guardrails"] == expected_after_response + assert logged_metadata["applied_guardrails"] == expected_after_response + + +def _three_step_two_pipeline_data() -> dict[str, object]: + runs_second: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="first", on_pass="next", on_fail="block"), + PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + ], + ) + stops_before_second: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="third", on_pass="allow", on_fail="block"), + PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + ], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second", "third"], + "_guardrail_pipelines": [("runs-second", runs_second), ("stops-before-second", stops_before_second)], + "_pipeline_managed_guardrails": {"first", "second", "third"}, + }, + } + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")] + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + **_three_step_two_pipeline_data(), + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1] + assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"] + + +@pytest.mark.asyncio +async def test_held_stream_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")] + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + logging_obj: Final = _logging_obj(log_parked=False) + + async def stream_log(*args: object) -> None: + return None + + logging_obj._on_deferred_stream_complete = stream_log + auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions") + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, + data={ + **_three_step_two_pipeline_data(), + "stream": True, + "litellm_logging_obj": logging_obj, + }, + call_type="completion", + ) + + delivered: Final[list[object]] = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data + ) + ] + + assert len(delivered) == 2 + assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1] + assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"] + + +def _ends_before_a_later_pipeline_data(ending: str) -> dict[str, object]: + ends: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass=ending)]) + never_reached: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="second")]) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second"], + "_guardrail_pipelines": [("ends", ends), ("never-reached", never_reached)], + "_pipeline_managed_guardrails": {"first", "second"}, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("ending", "raised"), [("block", HTTPException), ("modify_response", ModifyResponseException)]) +async def test_post_call_pipeline_that_ends_the_request_drops_the_steps_of_a_later_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + ending: str, + raised: type[Exception], +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_ends_before_a_later_pipeline_data(ending), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + + with pytest.raises(raised): + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 0) + assert data["metadata"]["applied_guardrails"] == ["first"] + + +def _blocks_before_a_later_request_pipeline_data() -> dict[str, object]: + blocks: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass="block")]) + request: Final = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="gate", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "gate", "second"], + "_guardrail_pipelines": [("blocks", blocks), ("request", request)], + "_pipeline_managed_guardrails": {"first", "gate", "second"}, + }, + } + + +@pytest.mark.asyncio +async def test_post_call_pipeline_that_blocks_keeps_a_later_request_step_whose_guardrail_runs_outside_pipelines( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + gate: Final = _PipelineStepGuardrail("gate", "pre_call") + second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + monkeypatch.setattr(litellm, "callbacks", [first, gate, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_blocks_before_a_later_request_pipeline_data(), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert sorted(data["metadata"]["applied_guardrails"]) == ["first", "gate", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("ending", "stream_end"), [("block", pytest.raises(HTTPException)), ("modify_response", nullcontext())] +) +async def test_held_stream_pipeline_that_ends_the_stream_drops_the_steps_of_a_later_pipeline( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + ending: str, + stream_end: AbstractContextManager[object], +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + logging_obj: Final = _logging_obj(log_parked=False) + + async def stream_log(*args: object) -> None: + return None + + logging_obj._on_deferred_stream_complete = stream_log + auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions") + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, + data={**_ends_before_a_later_pipeline_data(ending), "stream": True, "litellm_logging_obj": logging_obj}, + call_type="completion", + ) + + with stream_end: + async for _item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data + ): + pass + + assert (first.calls, second.calls) == (1, 0) + assert data["metadata"]["applied_guardrails"] == ["first"] + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_that_never_ran_whose_guardrail_ran_outside_the_pipeline( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("steps", "expected"), + [(("first", "second"), ["second"]), (("first",), None)], + ids=["skipped-then-ran", "only-step-skipped"], +) +async def test_post_call_pipeline_drops_a_step_its_guardrail_skips_on_a_whole_response( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + steps: tuple[str, ...], + expected: list[str] | None, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in steps] + guardrails[0].apply_stream_scope("streaming") + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline: Final = GuardrailPipeline( + mode="post_call", steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="block") for name in steps] + ) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": list(steps), + "_guardrail_pipelines": [("governance", pipeline)], + "_pipeline_managed_guardrails": set(steps), + }, + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert [guardrail.calls for guardrail in guardrails] == [0, 1][: len(steps)] + assert data["metadata"].get("applied_guardrails") == expected + + +class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: + self.calls += 1 + data["litellm_metadata"] = {"user_api_key_hash": "hashed"} + + +@pytest.mark.asyncio +async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_step_opens_litellm_metadata( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_two_step_pipeline_data("post_call", "next"), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_a_later_pipeline_runs_after_an_earlier_one_opens_litellm_metadata( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + stops_before_second: Final = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="first", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")], + ) + stops_before_first: Final = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="second", on_pass="allow"), PipelineStep(guardrail="first", on_pass="next")], + ) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second"], + "_guardrail_pipelines": [ + ("stops-before-second", stops_before_second), + ("stops-before-first", stops_before_first), + ], + "_pipeline_managed_guardrails": {"first", "second"}, + }, + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + @pytest.mark.asyncio async def test_pre_call_hook_names_every_pipeline_step_when_the_pipelines_run_on_another_request( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, @@ -1409,6 +1823,33 @@ async def test_deployment_hook_names_a_skipped_pre_call_step_that_raised( assert data["metadata"]["applied_guardrails"] == ["first", "second"] +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_skipped_step_a_deployment_hook_ran( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + + await second.async_post_call_success_deployment_hook( + request_data={**data, "guardrails": ["second"]}, + response=litellm.ModelResponse(), + call_type=CallTypes.acompletion, + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + @pytest.mark.asyncio async def test_deployment_hook_leaves_a_request_without_policy_pipelines_untouched() -> None: guardrail: Final = _PipelineStepGuardrail("second", "pre_call") @@ -1509,6 +1950,51 @@ async def test_deployment_hook_records_a_pipeline_step_it_ran_without_naming_one assert "applied_guardrails" not in kwargs["metadata"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("log_waits", "expected_after_stream"), + [(True, ["first"]), (False, ["first", "second"])], + ids=["stream-log-waits-for-the-pipeline", "stream-log-not-deferred"], +) +async def test_held_stream_pipeline_drops_the_steps_that_did_not_run_only_while_its_log_waits( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + log_waits: bool, + expected_after_stream: list[str], +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + logging_obj: Final = _logging_obj(log_parked=False) + + async def stream_log(*args: object) -> None: + return None + + logging_obj._on_deferred_stream_complete = stream_log if log_waits else None + auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions") + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, + data={**_two_step_pipeline_data("post_call", "allow"), "stream": True, "litellm_logging_obj": logging_obj}, + call_type="completion", + ) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + logged_metadata: Final = data["metadata"].copy() + + delivered: Final[list[object]] = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data + ) + ] + + assert len(delivered) == 2 + assert (first.calls, second.calls) == (1, 0) + assert data["metadata"]["applied_guardrails"] == expected_after_stream + assert logged_metadata["applied_guardrails"] == expected_after_stream + + @pytest.mark.asyncio async def test_post_call_pipeline_replacement_response_reaches_caller( proxy_logging, make_user_api_key_auth, monkeypatch @@ -3769,3 +4255,4 @@ async def test_live_pipeline_disconnect_scan_that_cannot_run_records_every_step_ assert seen.get("count") is None assert _post_call_statuses(data) == ["guardrail_failed_to_respond"] assert any("cut short" in message and "RuntimeError" in message for message in _warnings(caplog)) +