Merge f68284f75f into 02dcc4d347

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

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,76 @@ 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(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """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()),
+    )
+    # monkeypatch rather than a raw module-global write or patch.object: both are
+    # process-wide on the SDK, and the fixture undoes them at teardown.
+    monkeypatch.setattr(
+        litellm, "callbacks", [_GateOnlyLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]})]
+    )
+    monkeypatch.setattr(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(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """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.
+    """
+    monkeypatch.setattr(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"