mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
test(agentic loop): pin that a converted stream is async iterable
The fix this PR carried has since landed upstream ina14daf18d6andfa5a4f06ed, so only the regression tests are left. `_wrap_response_as_fake_stream` now builds a real CustomStreamWrapper around a MockResponseIterator. Nothing in the suite pins that, and the defect it replaced was silent: returning a bare ModelResponseStream to a caller that sent stream=true raised "'async for' requires an object with __aiter__ method" only once the client iterated. Two cases, both driving maybe_run_chat_completion_agentic_loop with a real Logging object: the tool-call path that runs a follow-up, and the plain reply that gates nothing. Each asserts the result is a CustomStreamWrapper and then iterates it to the original content.
This commit is contained in:
parent
ee1a6407cb
commit
38dac9e31c
1 changed files with 89 additions and 0 deletions
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue