diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 106fab7af8a..aa952a436a0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -560,6 +560,24 @@ def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> return False +def _route_supports_streaming_pipelines(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]: + if not _route_supports_streaming_pipelines(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 _streamable_post_call_pipelines( request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> tuple[tuple[str, "GuardrailPipeline"], ...]: @@ -577,7 +595,7 @@ def _streamable_post_call_pipelines( post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: return () - if not _route_has_endpoint_translation(user_api_key_dict): + if not _route_supports_streaming_pipelines(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 " @@ -593,30 +611,6 @@ 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 @@ -3300,15 +3294,15 @@ class ProxyLogging: # dict lookups + llm_router.get_deployment() per callback per chunk. _cached_guardrail_data: dict | None = None _guardrail_data_computed = False - pipeline_managed: Final = ( - pipeline_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset() + pipeline_gated: Final = ( + stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset() ) for callback in litellm.callbacks: try: _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): - if callback.guardrail_name in pipeline_managed: + if callback.guardrail_name in pipeline_gated: continue # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks 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 d01f881c446..46d6f2bb7a9 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 @@ -2234,7 +2234,11 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 return None - managed = RecordingGuardrail( + class UnifiedRecordingGuardrail(RecordingGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + managed = UnifiedRecordingGuardrail( guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True ) free = RecordingGuardrail( @@ -2253,3 +2257,34 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( assert result is not None assert seen.get("gr-post") is None assert seen["gr-free"] == 1 + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_stream( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class ChunkHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen["count"] = seen.get("count", 0) + 1 + seen["response"] = response + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [ChunkHookGuardrail(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) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + ) + + assert result is not None + assert seen["count"] == 1 + assert seen["response"] == "hello "