From 3deadd768255ab2c5202b57aea98a9e56bea2623 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Thu, 27 Aug 2026 10:05:06 -0700 Subject: [PATCH] fix(agentic loop): return a real stream when code-interpreter interception converts the request A client that sent stream=true and hit code-interpreter interception got back a bare ModelResponseStream, and `async for` over that raises "'async for' requires an object with __aiter__ method". Wrap the response in a CustomStreamWrapper so the streamed shape the caller asked for is what it receives. Rebased onto the current staging tip: the test file is rebuilt on top of the current one, since upstream removed its sys.path bootstrap in the meantime. Signed-off-by: Vineeth Sai --- .../chat_completion_agentic_loop.py | 51 +++++++++-- .../test_chat_completion_agentic_loop.py | 85 +++++++++++++++++++ 2 files changed, 128 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 07bed1f88ad..8b7f442c6b3 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -1,7 +1,7 @@ # this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api import json -from typing import Final, cast +from typing import TYPE_CHECKING, Final, cast from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger @@ -19,6 +19,9 @@ from litellm.types.integrations.custom_logger import ( from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + _FOLLOWUP_INTERNAL_PARAMS: Final = frozenset( ( "acompletion", @@ -87,16 +90,38 @@ def _check_agentic_loop_safety( return fingerprint -def _wrap_response_as_fake_stream(response: object) -> object: +def _wrap_response_as_fake_stream( + response: object, + *, + model: str, + custom_llm_provider: str, + logging_obj: object, +) -> object: + """Present a non-streamed response to a caller that asked for a stream. + + The caller is iterating the result, so this has to be a real async iterable. + Handing back a bare ModelResponseStream raised "'async for' requires an + object with __aiter__ method" and left a chunk-shaped object in the response + cache for the next request to trip over. + """ + if isinstance(response, CustomStreamWrapper): + return response if getattr(response, "object", None) == "chat.completion.chunk": return response if not hasattr(response, "choices"): return response - from litellm.llms.base_llm.base_model_iterator import ( - convert_model_response_to_streaming, - ) + if not hasattr(logging_obj, "model_call_details"): + # Without a logging object there is nothing to build a wrapper around; + # returning the response unchanged beats raising from the wrapper. + return response + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator - return convert_model_response_to_streaming(cast(ModelResponse, response)) + return CustomStreamWrapper( + completion_stream=MockResponseIterator(model_response=cast(ModelResponse, response)), + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=cast("LiteLLMLoggingObj", logging_obj), # cast-ok: duck-typed on model_call_details above + ) def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None: @@ -178,7 +203,12 @@ async def _execute_chat_completion_agentic_plan( str(e), ) if kwargs.get("_code_interpreter_interception_converted_stream") and not depth: - return _wrap_response_as_fake_stream(response_followup) + return _wrap_response_as_fake_stream( + response_followup, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) return response_followup finally: try: @@ -305,6 +335,11 @@ async def maybe_run_chat_completion_agentic_loop( if kwargs.get("_code_interpreter_interception_converted_stream") and not depth and hasattr(response, "choices"): return cast( "ModelResponse | CustomStreamWrapper", - _wrap_response_as_fake_stream(response), + _wrap_response_as_fake_stream( + response, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ), ) return None diff --git a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py index 434daab6ab5..0a0c15bac75 100644 --- a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -20,6 +20,7 @@ removed, so `test_internal_control_fields_never_leak_into_provider_body` proves they stay out of the body even without it. """ +import time from typing import Any, Dict, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch @@ -34,6 +35,8 @@ from litellm.integrations.code_interpreter_interception.handler import ( from litellm.litellm_core_utils.chat_completion_agentic_loop import ( maybe_run_chat_completion_agentic_loop, ) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.utils import CustomStreamWrapper from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, @@ -214,6 +217,19 @@ class _LoggingStub: dynamic_success_callbacks: List[Any] = [] +def _real_logging_obj() -> LiteLLMLoggingObj: + """A real logging object, which the streaming wrapper reads settings off.""" + return LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="call-test", + function_id="fn-test", + ) + + class _GateOnlyLogger(CustomLogger): """Overrides the gate to fire, but builds a plan from request_patch.""" @@ -409,3 +425,72 @@ async def test_dispatcher_raises_on_repeated_tool_call_fingerprint(restore_callb ) acompletion_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_converted_stream_result_is_async_iterable_after_the_loop_runs( + restore_callbacks, +): + """A client that sent stream=true gets something it can `async for` over. + + With code-interpreter interception the proxy converts the request to a + non-streaming call, so the dispatcher has to hand the streamed shape back. + It used to return a bare ModelResponseStream, and iterating that raised + "'async for' requires an object with __aiter__ method". + """ + followup = _plain_model_response("42") + plan = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=_patched_messages()), + ) + litellm.callbacks = [_GateOnlyLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]})] + + with patch.object(litellm, "acompletion", AsyncMock(return_value=followup)): + result = await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={}, + kwargs={ + "_code_interpreter_interception_active": True, + "_code_interpreter_interception_converted_stream": True, + }, + logging_obj=_real_logging_obj(), + custom_llm_provider="openai", + stream=True, + ) + + assert isinstance(result, CustomStreamWrapper) + chunks = [chunk async for chunk in result] + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "42" + + +@pytest.mark.asyncio +async def test_converted_stream_result_is_async_iterable_without_a_tool_call( + restore_callbacks, +): + """The same holds when the model never calls the tool. + + No callback gates, so no follow-up runs, and the dispatcher returns the + original response in streamed form. That path had the same defect, which is + why a plain assistant reply was enough to trigger the failure. + """ + litellm.callbacks = [] + + result = await maybe_run_chat_completion_agentic_loop( + response=_plain_model_response("no tool needed"), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + optional_params={}, + kwargs={ + "_code_interpreter_interception_active": True, + "_code_interpreter_interception_converted_stream": True, + }, + logging_obj=_real_logging_obj(), + custom_llm_provider="openai", + stream=True, + ) + + assert isinstance(result, CustomStreamWrapper) + chunks = [chunk async for chunk in result] + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "no tool needed"