From 447016817cfc5f5af57eb7e5e4550d3523b143de Mon Sep 17 00:00:00 2001
From: Yuta Saito <uc4w6c@bma.biglobe.ne.jp>
Date: Thu, 4 Sep 2025 07:11:44 +0900
Subject: [PATCH] fix: Call guardrail during stream processing

---
 litellm/proxy/common_request_processing.py | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
index f4d794d94bc..d41da69b6dc 100644
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -108,7 +108,6 @@ async def create_streaming_response(
     final_status_code = default_status_code
 
     try:
-
         # Handle coroutine that returns a generator
         if asyncio.iscoroutine(generator):
             generator = await generator
@@ -117,7 +116,6 @@ async def create_streaming_response(
         first_chunk_value = await generator.__anext__()
 
         if first_chunk_value is not None:
-
             try:
                 error_code_from_chunk = await _parse_event_data_for_error(
                     first_chunk_value
@@ -131,7 +129,6 @@ async def create_streaming_response(
                 verbose_proxy_logger.debug(f"Error parsing first chunk value: {e}")
 
     except StopAsyncIteration:
-
         # Generator was empty. Default status
         async def empty_gen() -> AsyncGenerator[str, None]:
             if False:
@@ -144,7 +141,6 @@ async def create_streaming_response(
             status_code=default_status_code,
         )
     except Exception as e:
-
         # Unexpected error consuming first chunk.
         verbose_proxy_logger.exception(
             f"Error consuming first chunk from generator: {e}"
@@ -167,7 +163,6 @@ async def create_streaming_response(
             with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE):
                 yield first_chunk_value
         async for chunk in generator:
-
             with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE):
                 yield chunk
 
@@ -460,7 +455,6 @@ class ProxyBaseLLMRequestProcessing:
         ) or self._is_streaming_response(
             response
         ):  # use generate_responses to stream responses
-
             custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
                 user_api_key_dict=user_api_key_dict,
                 call_id=logging_obj.litellm_call_id,
@@ -478,7 +472,6 @@ class ProxyBaseLLMRequestProcessing:
             if route_type == "allm_passthrough_route":
                 # Check if response is an async generator
                 if self._is_streaming_response(response):
-
                     if asyncio.iscoroutine(response):
                         generator = await response
                     else:
@@ -499,7 +492,6 @@ class ProxyBaseLLMRequestProcessing:
                         headers=custom_headers,
                     )
             else:
-
                 selected_data_generator = select_data_generator(
                     response=response,
                     user_api_key_dict=user_api_key_dict,
@@ -738,7 +730,11 @@ class ProxyBaseLLMRequestProcessing:
         verbose_proxy_logger.debug("inside generator")
         try:
             str_so_far = ""
-            async for chunk in response:
+            async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook(
+                user_api_key_dict=user_api_key_dict,
+                response=response,
+                request_data=request_data,
+            ):
                 verbose_proxy_logger.debug(
                     "async_data_generator: received streaming chunk - {}".format(chunk)
                 )