diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 07bed1f88ad..1dc534fb8ff 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/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 diff --git a/tests/test_litellm/litellm_core_utils/test_agentic_loop_fake_stream.py b/tests/test_litellm/litellm_core_utils/test_agentic_loop_fake_stream.py new file mode 100644 index 00000000000..a4308b5a48a --- /dev/null +++ b/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