diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7e85b66682c..2a6eb0a16ff 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3343,15 +3343,15 @@ class ProxyBaseLLMRequestProcessing: from litellm.proxy.proxy_server import llm_router as _global_llm_router from litellm.proxy.utils import ( _check_and_merge_model_level_guardrails, - pipeline_managed_guardrail_names, + stream_gated_guardrail_names, ) guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router) - pipeline_managed: Final = pipeline_managed_guardrail_names(captured_data, "post_call") + stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict) for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue - if cb.guardrail_name in pipeline_managed: + if cb.guardrail_name in stream_gated: continue if not cb.should_run_guardrail( data=guardrail_data, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bbec1b842c4..106fab7af8a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -538,12 +538,16 @@ def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> No ) -def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: - unsupported: Final = tuple( +def _pipeline_unsupported_streaming_guardrails(pipeline: "GuardrailPipeline") -> tuple[str, ...]: + return tuple( dict.fromkeys( step.guardrail for step in pipeline.steps if not _pipeline_step_supports_streaming(step.guardrail) ) ) + + +def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: + unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline) if not unsupported: return True verbose_proxy_logger.warning( @@ -573,13 +577,12 @@ def _streamable_post_call_pipelines( post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: return () - route: Final = user_api_key_dict.request_route - if route and resolve_endpoint_translation(user_api_key_dict, None) is None: + if not _route_has_endpoint_translation(user_api_key_dict): verbose_proxy_logger.warning( "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run " "on their own: %s", - route, + user_api_key_dict.request_route, ", ".join(policy_name for policy_name, _pipeline in post_call_pipelines), ) return () @@ -590,6 +593,30 @@ def _streamable_post_call_pipelines( ) +def _route_has_endpoint_translation(user_api_key_dict: UserAPIKeyAuth) -> bool: + return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None + + +def stream_gated_guardrail_names( + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> frozenset[str]: + """ + The guardrails whose post_call pipelines gate a streaming response on this + route: the selection ``_streamable_post_call_pipelines`` makes, without its + warnings, so the post-call pass deferred to the end of the stream skips + exactly the guardrails the pipelines already ran and no others. + """ + if not _route_has_endpoint_translation(user_api_key_dict): + return frozenset() + return _pipeline_step_guardrail_names( + tuple( + (policy_name, pipeline) + for policy_name, pipeline in _post_call_pipelines(request_data) + if not _pipeline_unsupported_streaming_guardrails(pipeline) + ) + ) + + def _prompt_block_text(block: object) -> str: if isinstance(block, str): return block diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index f4165f98ba5..10de7e7fc3e 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -697,6 +697,69 @@ async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monk assert pipeline_managed.native_hooks_ran == [] +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_whose_pipeline_could_not_stream(monkeypatch): + """A pipeline step with neither streaming interface keeps the whole pipeline off the + stream, so the deferred audit is the only place the opted-out guardrail's own hook + still runs, the way it did before pipelines ran on streams.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + class NeitherHookGuardrail(CustomGuardrail): + pass + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + neither = NeitherHookGuardrail(guardrail_name="gr-neither", event_hook=GuardrailEventHooks.post_call) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed, neither]) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="keeps_native", on_fail="next"), + PipelineStep(guardrail="gr-neither", on_fail="block"), + ], + ) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_on_route_without_translation(monkeypatch): + """A route with no endpoint guardrail translation cannot gate the stream through its + pipelines, so the deferred audit still owes the opted-out guardrail its own hook.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/custom/stream"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + @pytest.mark.asyncio async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch): """The realtime path calls apply_guardrail directly, so the opt-out has to be