fix response not being redacted for custom callbacks with streaming enabled

litellm/litellm_core_utils/redact_messages.py

@@ -99,6 +99,23 @@ def perform_redaction(model_call_details: dict, result):
             if hasattr(_streaming_response, "reasoning") and _streaming_response.reasoning is not None:
                 _streaming_response.reasoning = None
 
+    if (
+        model_call_details.get("stream", False) is True
+        and "async_complete_streaming_response" in model_call_details
+    ):
+        _streaming_response = model_call_details["async_complete_streaming_response"]
+        if hasattr(_streaming_response, "choices"):
+            for choice in _streaming_response.choices:
+                _redact_choice_content(choice)
+        elif hasattr(_streaming_response, "output"):
+            _redact_responses_api_output(_streaming_response.output)
+            # Redact reasoning field in ResponsesAPIResponse
+            if (
+                hasattr(_streaming_response, "reasoning")
+                and _streaming_response.reasoning is not None
+            ):
+                _streaming_response.reasoning = None            
+
     # Redact result
     if result is not None:
         # Check if result is a coroutine, async generator, or other async object - these cannot be deepcopied

tests/logging_callback_tests/test_logging_redaction_e2e_test.py

@@ -2,7 +2,7 @@ import io
 import os
 import sys
 
-from typing import Optional
+from typing import Optional, Union
 
 sys.path.insert(0, os.path.abspath("../.."))
 
@@ -12,23 +12,26 @@ import json
 import logging
 import time
 from unittest.mock import AsyncMock, patch
+from datetime import datetime
 
 import pytest
 
 import litellm
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.types.utils import StandardLoggingPayload
+from litellm.types.utils import StandardLoggingPayload, ModelResponse, TextCompletionResponse, ResponsesAPIResponse
 
 
 class TestCustomLogger(CustomLogger):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None
+        self.response_obj: Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]] = None
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         standard_logging_payload = kwargs.get("standard_logging_object", None)
         self.logged_standard_logging_payload = standard_logging_payload
+        self.response_obj = response_obj
 
 
 @pytest.mark.asyncio
@@ -120,6 +123,47 @@ async def test_global_redaction_off_with_dynamic_params(turn_off_message_logging
         assert standard_logging_payload["messages"][0]["content"] == "hi"
 
 
+@pytest.mark.asyncio
+async def test_redaction_with_custom_logger_streaming():
+    """Test redaction of responses for custom logger callbacks"""
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    class LoggingWithoutSyncSuccessHandler(Logging):
+        def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs):
+            pass
+
+    litellm.turn_off_message_logging = True
+    test_custom_logger = TestCustomLogger()
+
+    litellm_logging_obj = LoggingWithoutSyncSuccessHandler(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=True,
+        call_type="acompletion",
+        litellm_call_id="1234",
+        start_time=datetime.now(),
+        function_id="1234",
+        dynamic_async_success_callbacks=[test_custom_logger],
+    )
+
+    response = await litellm.acompletion(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        mock_response="hello",
+        stream=True,
+        litellm_logging_obj=litellm_logging_obj,
+    )
+    
+    # Consume the stream to trigger logging
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+    
+    await asyncio.sleep(1)
+    async_complete_streaming_response = test_custom_logger.response_obj
+    assert async_complete_streaming_response is not None    
+    assert (async_complete_streaming_response.choices[0].message.content == "redacted-by-litellm")
+
 @pytest.mark.asyncio
 async def test_redaction_responses_api():
     """Test redaction with ResponsesAPIResponse format"""