diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d7bf832bca4..8b36061c6be 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -447,14 +447,15 @@ def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardrail ) +def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> frozenset[str]: + return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps) + + def _pipeline_managed_guardrail_names( data: Mapping[str, object], mode: Literal["pre_call", "post_call"] ) -> frozenset[str]: - return frozenset( - step.guardrail - for _policy_name, pipeline in _policy_pipelines(data) - if pipeline.mode == mode - for step in pipeline.steps + return _pipeline_step_guardrail_names( + tuple((policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == mode) ) @@ -547,7 +548,7 @@ def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> return True verbose_proxy_logger.warning( "Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, " - "which streaming pipelines need; the stream is released ungoverned by it: %s", + "which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s", policy_name, ", ".join(unsupported), ) @@ -563,9 +564,9 @@ def _streamable_post_call_pipelines( Streaming pipelines scan the buffered stream through the endpoint guardrail translation of the request route, so every step's guardrail needs the unified apply_guardrail interface and the route needs a translation. A - pipeline that cannot be run that way yet is left out and the stream is - released the way it was before pipelines ran on streams at all, with a - warning naming what went ungoverned. + pipeline that cannot be run that way yet is left out and its guardrails + run on the stream on their own, the way they did before pipelines ran on + streams at all, with a warning naming the pipeline. """ post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: @@ -574,7 +575,8 @@ def _streamable_post_call_pipelines( if route and resolve_endpoint_translation(user_api_key_dict, None) is None: verbose_proxy_logger.warning( "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " - "(no endpoint guardrail translation); the stream is released ungoverned by them: %s", + "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run " + "on their own: %s", route, ", ".join(policy_name for policy_name, _pipeline in post_call_pipelines), ) @@ -3361,10 +3363,10 @@ class ProxyLogging: current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) - pipeline_managed_names: Final = _pipeline_managed_guardrail_names(request_data, "post_call") + pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): - if resolved_callback.guardrail_name in pipeline_managed_names: + if resolved_callback.guardrail_name in pipeline_gated_names: continue if ( resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) 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 ad2b4a0efbb..73dd6746b29 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 @@ -1620,6 +1620,42 @@ async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_l assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeline_cannot_stream( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + + class IteratorHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["count"] = seen.get("count", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + monkeypatch.setattr( + litellm, + "callbacks", + [IteratorHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ) + ] + + assert seen["count"] == 1 + assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"] + assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + @pytest.mark.asyncio @pytest.mark.parametrize( "rewrite_attribute, value",