From eb689a1f07afb3677d3e34834f819de27222817b Mon Sep 17 00:00:00 2001 From: Raney Cain <36416768+rcII@users.noreply.github.com> Date: Mon, 8 Dec 2025 07:29:42 +0000 Subject: [PATCH] fix(proxy): async_post_call_streaming_iterator_hook now properly iterates async generators (#17626) The async_post_call_streaming_iterator_hook function was broken: 1. Was a sync function (def) not async generator 2. Returned AsyncGenerator without iterating it 3. Callback generators were chained but never consumed This fix: 1. Makes the function an async generator (async def + yield) 2. Actually iterates through the chained callbacks with 'async for' 3. Properly yields chunks to the caller Fixes #9639 --- litellm/proxy/utils.py | 60 ++++-- ...async_post_call_streaming_iterator_hook.py | 194 ++++++++++++++++++ 2 files changed, 234 insertions(+), 20 deletions(-) create mode 100644 tests/test_litellm/proxy/hooks/test_async_post_call_streaming_iterator_hook.py diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 81d709c332c..ec9daebbf70 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -825,7 +825,12 @@ class ProxyLogging: return data def _process_prompt_template( - self, data: dict, litellm_logging_obj: Any, prompt_id: Any, prompt_version: Any, call_type: CallTypesLiteral + self, + data: dict, + litellm_logging_obj: Any, + prompt_id: Any, + prompt_version: Any, + call_type: CallTypesLiteral, ) -> None: """Process prompt template if applicable.""" from litellm.utils import get_non_default_completion_params @@ -878,27 +883,37 @@ class ProxyLogging: from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) + metadata_standard = data.get("metadata") or {} metadata_litellm = data.get("litellm_metadata") or {} - + guardrails_in_metadata = [] if isinstance(metadata_standard, dict) and "guardrails" in metadata_standard: guardrails_in_metadata = metadata_standard.get("guardrails", []) elif isinstance(metadata_litellm, dict) and "guardrails" in metadata_litellm: guardrails_in_metadata = metadata_litellm.get("guardrails", []) - + if guardrails_in_metadata and isinstance(guardrails_in_metadata, list): applied_guardrails = [] - if isinstance(metadata_standard, dict) and "applied_guardrails" in metadata_standard: + if ( + isinstance(metadata_standard, dict) + and "applied_guardrails" in metadata_standard + ): applied_guardrails = metadata_standard.get("applied_guardrails", []) - elif isinstance(metadata_litellm, dict) and "applied_guardrails" in metadata_litellm: + elif ( + isinstance(metadata_litellm, dict) + and "applied_guardrails" in metadata_litellm + ): applied_guardrails = metadata_litellm.get("applied_guardrails", []) - + if not isinstance(applied_guardrails, list): applied_guardrails = [] - + for guardrail_name in guardrails_in_metadata: - if isinstance(guardrail_name, str) and guardrail_name not in applied_guardrails: + if ( + isinstance(guardrail_name, str) + and guardrail_name not in applied_guardrails + ): add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=guardrail_name ) @@ -1022,10 +1037,10 @@ class ProxyLogging: start_time=start_time, end_time=end_time, ) - + if data is not None: self._process_guardrail_metadata(data) - + return data except Exception as e: raise e @@ -1602,7 +1617,7 @@ class ProxyLogging: raise e return response - def async_post_call_streaming_iterator_hook( + async def async_post_call_streaming_iterator_hook( self, response, user_api_key_dict: UserAPIKeyAuth, @@ -1615,6 +1630,7 @@ class ProxyLogging: Covers: 1. /chat/completions """ + current_response = response for callback in litellm.callbacks: @@ -1631,23 +1647,27 @@ class ProxyLogging: ) or _callback.should_run_guardrail( data=request_data, event_type=GuardrailEventHooks.post_call ): - if "apply_guardrail" in type(callback).__dict__: request_data["guardrail_to_apply"] = callback - response = ( + current_response = ( unified_guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, request_data=request_data, - response=response, + response=current_response, ) ) else: - response = _callback.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=response, - request_data=request_data, + current_response = ( + _callback.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=current_response, + request_data=request_data, + ) ) - return response + + # Actually iterate through the chained async generator and yield chunks + async for chunk in current_response: + yield chunk def _init_response_taking_too_long_task(self, data: Optional[dict] = None): """ @@ -3143,7 +3163,7 @@ class PrismaClient: key = (check.model_id, check.model_name) else: key = (None, check.model_name) - + # Only add if we haven't seen this key yet (since checks are ordered by checked_at desc) if key not in latest_checks: latest_checks[key] = check diff --git a/tests/test_litellm/proxy/hooks/test_async_post_call_streaming_iterator_hook.py b/tests/test_litellm/proxy/hooks/test_async_post_call_streaming_iterator_hook.py new file mode 100644 index 00000000000..50c6a580f91 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_async_post_call_streaming_iterator_hook.py @@ -0,0 +1,194 @@ +""" +Tests for async_post_call_streaming_iterator_hook fix. + +Verifies that the hook: +1. Is an async generator (not a sync function) +2. Properly iterates through callback chain +3. Actually yields chunks from async generators +""" + +import os +import sys +from typing import AsyncGenerator, Any +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging + + +class MockStreamingCallback(CustomLogger): + """Test callback that tracks chunk processing.""" + + def __init__(self, prefix: str = ""): + super().__init__() + self.prefix = prefix + self.chunks_processed = 0 + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: AsyncGenerator[Any, None], + request_data: dict, + ) -> AsyncGenerator[Any, None]: + """Transform chunks by tracking and optionally prefixing.""" + async for chunk in response: + self.chunks_processed += 1 + # Optionally modify chunk content for testing + if self.prefix and isinstance(chunk, dict): + if "choices" in chunk: + for choice in chunk["choices"]: + if "delta" in choice and "content" in choice["delta"]: + choice["delta"]["content"] = ( + f"[{self.prefix}]" + choice["delta"]["content"] + ) + yield chunk + + +async def mock_streaming_response() -> AsyncGenerator[dict, None]: + """Simulate an LLM streaming response.""" + chunks = [ + {"choices": [{"delta": {"content": "Hello"}}]}, + {"choices": [{"delta": {"content": " "}}]}, + {"choices": [{"delta": {"content": "World"}}]}, + {"choices": [{"delta": {"content": "!"}}]}, + ] + for chunk in chunks: + yield chunk + + +@pytest.mark.asyncio +async def test_streaming_hook_is_async_generator(): + """Verify that the hook is an async generator that yields chunks.""" + # Arrange + proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + callback = MockStreamingCallback() + + user_api_key_dict = UserAPIKeyAuth(api_key="test_key") + request_data = {"model": "gpt-4", "messages": []} + + with patch.object(litellm, "callbacks", [callback]): + # Act + result = proxy_logging.async_post_call_streaming_iterator_hook( + response=mock_streaming_response(), + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + + # Assert - result should be an async generator + assert hasattr(result, "__anext__"), "Result should be an async iterator" + + # Collect chunks + collected_chunks = [] + async for chunk in result: + collected_chunks.append(chunk) + + # Verify all chunks were yielded + assert ( + len(collected_chunks) == 4 + ), f"Expected 4 chunks, got {len(collected_chunks)}" + assert ( + callback.chunks_processed == 4 + ), "Callback should have processed 4 chunks" + + +@pytest.mark.asyncio +async def test_streaming_hook_chains_multiple_callbacks(): + """Verify that multiple callbacks are properly chained.""" + # Arrange + proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + callback1 = MockStreamingCallback(prefix="CB1") + callback2 = MockStreamingCallback(prefix="CB2") + + user_api_key_dict = UserAPIKeyAuth(api_key="test_key") + request_data = {"model": "gpt-4", "messages": []} + + with patch.object(litellm, "callbacks", [callback1, callback2]): + # Act + result = proxy_logging.async_post_call_streaming_iterator_hook( + response=mock_streaming_response(), + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + + # Collect chunks + collected_chunks = [] + async for chunk in result: + collected_chunks.append(chunk) + + # Assert - both callbacks should have processed all chunks + assert callback1.chunks_processed == 4 + assert callback2.chunks_processed == 4 + + # Verify chaining worked (CB2 wraps CB1's output) + first_content = collected_chunks[0]["choices"][0]["delta"]["content"] + assert "[CB2]" in first_content, "CB2 prefix should be present" + assert "[CB1]" in first_content, "CB1 prefix should be present (wrapped by CB2)" + + +@pytest.mark.asyncio +async def test_streaming_hook_handles_empty_callbacks(): + """Verify that the hook works with no callbacks registered.""" + # Arrange + proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + + user_api_key_dict = UserAPIKeyAuth(api_key="test_key") + request_data = {"model": "gpt-4", "messages": []} + + with patch.object(litellm, "callbacks", []): + # Act + result = proxy_logging.async_post_call_streaming_iterator_hook( + response=mock_streaming_response(), + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + + # Collect chunks + collected_chunks = [] + async for chunk in result: + collected_chunks.append(chunk) + + # Assert - all chunks should pass through unchanged + assert len(collected_chunks) == 4 + + +@pytest.mark.asyncio +async def test_streaming_hook_propagates_callback_errors(): + """Verify that callback errors during iteration are properly propagated.""" + # Arrange + proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + + class FailingCallback(CustomLogger): + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: AsyncGenerator[Any, None], + request_data: dict, + ) -> AsyncGenerator[Any, None]: + raise RuntimeError("Callback failed!") + yield # Make it a generator + + failing_callback = FailingCallback() + + user_api_key_dict = UserAPIKeyAuth(api_key="test_key") + request_data = {"model": "gpt-4", "messages": []} + + with patch.object(litellm, "callbacks", [failing_callback]): + # Act + result = proxy_logging.async_post_call_streaming_iterator_hook( + response=mock_streaming_response(), + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + + # Assert - error should propagate when iterating + with pytest.raises(RuntimeError, match="Callback failed!"): + async for _ in result: + pass