mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(proxy): guard against None return in async_post_call_streaming_iterator_hook
Fixes "TypeError: 'async for' requires an object with __aiter__ method, got NoneType" error that occurs when a user-defined custom callback has a poorly implemented async_post_call_streaming_iterator_hook that returns None instead of an async generator. Changes: - Fix type(callback).__dict__ -> type(_callback).__dict__ to check the resolved instance - Add None check after calling hook to prevent current_response from becoming None - Add regression test for the NoneType crash scenario The fix ensures that if a callback returns None (e.g., a sync method instead of async generator), the streaming chain continues with the previous valid response iterator instead of crashing. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
ac6e1d9fd1
commit
7ea99461dd
2 changed files with 51 additions and 4 deletions
|
|
@ -2109,16 +2109,18 @@ class ProxyLogging:
|
|||
):
|
||||
if (
|
||||
"async_post_call_streaming_iterator_hook"
|
||||
in type(callback).__dict__
|
||||
in type(_callback).__dict__
|
||||
):
|
||||
current_response = (
|
||||
_new_response = (
|
||||
_callback.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=current_response,
|
||||
request_data=request_data,
|
||||
)
|
||||
)
|
||||
elif "apply_guardrail" in type(callback).__dict__:
|
||||
if _new_response is not None:
|
||||
current_response = _new_response
|
||||
elif "apply_guardrail" in type(_callback).__dict__:
|
||||
request_data["guardrail_to_apply"] = callback
|
||||
current_response = (
|
||||
unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
|
|
@ -2128,13 +2130,15 @@ class ProxyLogging:
|
|||
)
|
||||
)
|
||||
else:
|
||||
current_response = (
|
||||
_new_response = (
|
||||
_callback.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=current_response,
|
||||
request_data=request_data,
|
||||
)
|
||||
)
|
||||
if _new_response is not None:
|
||||
current_response = _new_response
|
||||
|
||||
# Actually iterate through the chained async generator and yield chunks
|
||||
async for chunk in current_response:
|
||||
|
|
|
|||
|
|
@ -192,3 +192,46 @@ async def test_streaming_hook_propagates_callback_errors():
|
|||
with pytest.raises(RuntimeError, match="Callback failed!"):
|
||||
async for _ in result:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_double_strip_on_second_call():
|
||||
"""Regression test: callback returning None should not break the streaming chain.
|
||||
|
||||
This tests the fix for: 'async for' requires an object with __aiter__ method, got NoneType
|
||||
Caused by a user-defined callback whose async_post_call_streaming_iterator_hook
|
||||
is a regular sync method (not an async generator) that returns None.
|
||||
"""
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=MagicMock())
|
||||
|
||||
class NoneReturningCallback(CustomLogger):
|
||||
"""Simulates a badly-implemented callback that returns None instead of an async generator."""
|
||||
|
||||
def async_post_call_streaming_iterator_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: Any,
|
||||
request_data: dict,
|
||||
):
|
||||
# Regular sync method, not an async generator — returns None implicitly
|
||||
pass
|
||||
|
||||
bad_callback = NoneReturningCallback()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test_key")
|
||||
request_data = {"model": "gpt-4", "messages": []}
|
||||
|
||||
with patch.object(litellm, "callbacks", [bad_callback]):
|
||||
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,
|
||||
)
|
||||
|
||||
# Should not raise TypeError about NoneType
|
||||
collected_chunks = []
|
||||
async for chunk in result:
|
||||
collected_chunks.append(chunk)
|
||||
|
||||
# All 4 original chunks should pass through unmodified
|
||||
assert len(collected_chunks) == 4
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue