diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d3fe1406abd..0f95cda3d46 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5565,7 +5565,8 @@ async def async_data_generator( request_data=request_data ) model_mismatch_logged = False - # Chunks for post-guardrail log: use exclude_none=True only so stream_chunk_builder gets required keys + # Chunks for post-guardrail log: use exclude_none=True only so stream_chunk_builder gets required keys. + # This list holds the full stream in memory until the background post-guardrail log task completes. _streaming_chunks_for_log: List[Dict[str, Any]] = [] # Use a running string instead of list + join to avoid O(n^2) overhead. # Previously "".join(str_so_far_parts) was called every chunk, re-joining diff --git a/tests/test_litellm/proxy/hooks/test_async_post_guardrail_log_success_event.py b/tests/test_litellm/proxy/hooks/test_async_post_guardrail_log_success_event.py index 88d94592b08..e66cbd36beb 100644 --- a/tests/test_litellm/proxy/hooks/test_async_post_guardrail_log_success_event.py +++ b/tests/test_litellm/proxy/hooks/test_async_post_guardrail_log_success_event.py @@ -85,18 +85,22 @@ async def test_post_guardrail_log_base_custom_logger_not_invoked(): """Base CustomLogger (no override) is not invoked; only overriders are called.""" base_logger = CustomLogger() overriding_logger = PostGuardrailLogger() - # Base first, then override: only overriding should be called - with patch("litellm.callbacks", [base_logger, overriding_logger]): - from litellm.proxy.utils import ProxyLogging - from litellm.caching.caching import DualCache - - proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) - await proxy_logging.async_post_guardrail_log_success_event( - data={"model": "gpt-4"}, - response=ModelResponse(id="r1", choices=[], model="gpt-4"), - user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), - ) + with patch.object( + base_logger, + "async_post_guardrail_log_success_event", + new_callable=AsyncMock, + ) as mock_base_method: + with patch("litellm.callbacks", [base_logger, overriding_logger]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + await proxy_logging.async_post_guardrail_log_success_event( + data={"model": "gpt-4"}, + response=ModelResponse(id="r1", choices=[], model="gpt-4"), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + mock_base_method.assert_not_called() assert overriding_logger.called is True