Merge remote-tracking branch 'origin/litellm_fix_post_call_policy_pipeline' into litellm_post_call_pipeline_background_responses

litellm/proxy/utils.py

@@ -630,6 +630,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 all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps)
+        )
+    )
+
+
 def _streamable_post_call_pipelines(
     request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
 ) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
@@ -646,13 +664,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_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 "
             "on their own: %s",
-            route,
+            user_api_key_dict.request_route,
             ", ".join(policy_name for policy_name, _pipeline in post_call_pipelines),
         )
         return ()
@@ -3360,15 +3377,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

tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py

@@ -2282,7 +2282,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(
@@ -2301,3 +2305,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 "