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

# Conflicts:
#	litellm/proxy/utils.py

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

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 "