diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f3138f10dac..438da037527 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1604,6 +1604,12 @@ class ProxyBaseLLMRequestProcessing: # here would duplicate the guardrail API call # (e.g. double OpenAI Moderation charges). continue + if "async_post_call_streaming_iterator_hook" in type(cb).__dict__: + # Skip — the guardrail already scanned the assembled + # response via its own streaming iterator hook in the + # streaming pipeline. re running this function async_post_call_success_hook + # here would duplicate the scan and can spuriously block the guardrail that already passed / failed. + continue else: guardrail_result = await cb.async_post_call_success_hook( user_api_key_dict=captured_user_api_key_dict, diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 59b6e24f430..e10258c0829 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -758,6 +758,60 @@ class TestDeferredStreamingClosure: apply_guardrail_called is False ), "apply_guardrail guardrails must be SKIPPED in deferred path" + @pytest.mark.asyncio + async def test_streaming_iterator_hook_skipped_in_deferred_path(self): + """regression test: guardrails that define async_post_call_streaming_iterator_hook must be SKIPPED in _run_deferred_stream_guardrails. + The iterator hook already scanned the assembled response in the streaming + pipeline""" + success_hook_called = False + + class IteratorHookGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="iterator-hook", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict, response, request_data + ): + async for chunk in response: + yield chunk + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + nonlocal success_hook_called + success_hook_called = True + return response + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = IteratorHookGuardrail() + + with patch("litellm.callbacks", [guardrail]): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + + assert success_hook_called is False, ( + "Guardrails that implement async_post_call_streaming_iterator_hook " + "must be SKIPPED in deferred path — the iterator hook already ran" + ) + @pytest.mark.asyncio async def test_hooks_receive_merged_guardrail_data(self): """Hooks must receive guardrail_data (the merged dict from