diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d80a3f4cbd6..5fd0e9d1f36 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -759,20 +759,27 @@ 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 one when its verdict ends the request before they run.""" + return pipelines[index:] if result.terminal_action != "allow" else pipelines[index : index + 1] + + def _drop_response_steps_that_did_not_run( data: dict[str, object], # mutable-ok: the request's applied guardrails are trimmed in place - steps: Sequence["PipelineStep"], + pipelines: Sequence[tuple[str, "GuardrailPipeline"]], steps_that_ran: frozenset[str], ) -> None: """ - Drop from the applied guardrails the steps of a response pipeline that no pipeline ran, + Drop from the applied guardrails the steps of these response pipelines that no pipeline ran, 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 = frozenset(step.guardrail for step in steps) + post_call: Final = _pipeline_step_guardrail_names(pipelines) not_run: Final = frozenset( name for name in post_call - steps_that_ran if _runs_only_in_pipelines(name, data, pre_call, post_call) ) @@ -2796,7 +2803,7 @@ class ProxyLogging: 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 policy_name, pipeline in pipelines: + for index, (policy_name, pipeline) in enumerate(pipelines): if pipeline.mode != event_hook: continue @@ -2814,7 +2821,9 @@ class ProxyLogging: steps_that_ran = steps_that_ran | _steps_that_ran(result) if log_waits: - _drop_response_steps_that_did_not_run(request, pipeline.steps, steps_that_ran) + _drop_response_steps_that_did_not_run( + request, _pipelines_settled_by(pipelines, index, result), steps_that_ran + ) data = self._handle_pipeline_result( result=result, @@ -4734,7 +4743,7 @@ class ProxyLogging: call_type, endpoint_translation = translation steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran - for policy_name, pipeline in pipelines: + for index, (policy_name, pipeline) in enumerate(pipelines): result: PipelineExecutionResult = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode="post_call", @@ -4747,7 +4756,9 @@ class ProxyLogging: ) 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, pipeline.steps, steps_that_ran) + _drop_response_steps_that_did_not_run( + request_data, _pipelines_settled_by(pipelines, index, result), steps_that_ran + ) try: ProxyLogging._handle_pipeline_result( result, data=request_data, policy_name=policy_name, original_response=buffered 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 99a06f4e141..ac314667347 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 @@ -1496,6 +1497,86 @@ async def test_held_stream_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_ 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"] + + +@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,