Merge 77e928180e into 1df25e26cf

litellm/litellm_core_utils/chat_completion_agentic_loop.py

@@ -87,16 +87,36 @@ def _check_agentic_loop_safety(
     return fingerprint
 
 
-def _wrap_response_as_fake_stream(response: object) -> object:
-    if getattr(response, "object", None) == "chat.completion.chunk":
+def _wrap_response_as_fake_stream(
+    response: object,
+    *,
+    model: str,
+    custom_llm_provider: str,
+    logging_obj: object,
+) -> object:
+    """
+    Present a non-streamed response as something the streaming path can consume.
+
+    Interception downgrades `stream: true` to a single non-streamed call so the
+    agentic loop can run, then has to hand a stream back. Returning the bare
+    converted chunk made callers fail with `TypeError: 'async for' requires an
+    object with __aiter__ method, got ModelResponseStream`, and put a
+    `chat.completion.chunk`-shaped object where a full `chat.completion` was
+    expected — which is how a malformed entry reaches the response cache and
+    later raises `KeyError: 'message'`.
+    """
+    if isinstance(response, CustomStreamWrapper):
         return response
     if not hasattr(response, "choices"):
         return response
-    from litellm.llms.base_llm.base_model_iterator import (
-        convert_model_response_to_streaming,
-    )
+    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=logging_obj,
+    )
 
 
 def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None:
@@ -178,7 +198,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 +330,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_agentic_loop_fake_stream.py

@@ -0,0 +1,72 @@
+"""A downgraded stream must come back as something `async for` can consume."""
+
+import pytest
+
+from litellm.litellm_core_utils.chat_completion_agentic_loop import (
+    _wrap_response_as_fake_stream,
+)
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
+from litellm.types.utils import Choices, Message, ModelResponse
+from litellm.utils import CustomStreamWrapper
+
+
+def _response() -> ModelResponse:
+    return ModelResponse(
+        id="chatcmpl-1",
+        choices=[Choices(index=0, finish_reason="stop", message=Message(role="assistant", content="hello"))],
+        model="gpt-4o",
+        object="chat.completion",
+    )
+
+
+def _logging_obj() -> LiteLLMLogging:
+    return LiteLLMLogging(
+        model="gpt-4o",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=True,
+        call_type="completion",
+        start_time=None,
+        litellm_call_id="test-call-id",
+        function_id="1",
+    )
+
+
+def _wrap(response):
+    return _wrap_response_as_fake_stream(
+        response, model="gpt-4o", custom_llm_provider="openai", logging_obj=_logging_obj()
+    )
+
+
+class TestWrapResponseAsFakeStream:
+    def test_result_is_async_iterable(self):
+        """The reported crash: `async for` got a bare ModelResponseStream."""
+        wrapped = _wrap(_response())
+
+        assert hasattr(wrapped, "__aiter__"), "result must support `async for`"
+        assert isinstance(wrapped, CustomStreamWrapper)
+
+    def test_result_is_sync_iterable_too(self):
+        wrapped = _wrap(_response())
+
+        assert hasattr(wrapped, "__iter__")
+
+    @pytest.mark.asyncio
+    async def test_yields_the_original_content_as_chunks(self):
+        wrapped = _wrap(_response())
+
+        chunks = [chunk async for chunk in wrapped]
+
+        assert chunks, "expected at least one chunk"
+        assert all(getattr(c, "object", None) == "chat.completion.chunk" for c in chunks)
+        text = "".join((c.choices[0].delta.content or "") for c in chunks if getattr(c, "choices", None))
+        assert "hello" in text
+
+    def test_an_already_wrapped_stream_is_passed_through(self):
+        wrapped = _wrap(_response())
+
+        assert _wrap(wrapped) is wrapped
+
+    def test_object_without_choices_is_returned_unchanged(self):
+        sentinel = object()
+
+        assert _wrap(sentinel) is sentinel