mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
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
This commit is contained in:
Raney Cain
GitHub
parent
2d112fc8b2
commit
eb689a1f07
2 changed files with 234 additions and 20 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue