Skip to content

fix(websearch): wrap agentic loop response in fake stream for streaming requests #27449

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
vokako wants to merge 1 commit into from
Closed

fix(websearch): wrap agentic loop response in fake stream for streaming requests #27449

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 60 additions & 46 deletions litellm/llms/custom_httpx/llm_http_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4740,6 +4740,42 @@ async def _execute_chat_completion_agentic_plan(
**kwargs_for_followup,
)

def _maybe_wrap_in_fake_stream(
self,
response: Any,
logging_obj: "LiteLLMLoggingObj",
) -> Any:
"""
If the original request was streaming but converted to non-streaming for
WebSearch interception, wrap the dict response in a FakeAnthropicMessagesStreamIterator.
"""
websearch_converted_stream = (
logging_obj.model_call_details.get(
"websearch_interception_converted_stream", False
)
if logging_obj is not None
else False
)
if websearch_converted_stream and isinstance(response, dict):
from typing import cast

from litellm._logging import verbose_logger
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)

verbose_logger.debug(
"WebSearchInterception: Agentic loop completed, "
"converting non-streaming response to fake stream"
)
return FakeAnthropicMessagesStreamIterator(
response=cast(AnthropicMessagesResponse, response)
)
return response

async def _call_agentic_completion_hooks(
self,
response: Any,
Expand Down Expand Up @@ -4821,7 +4857,7 @@ async def _call_agentic_completion_hooks(
is not CustomLogger.async_build_agentic_loop_plan
)
if not build_plan_overridden:
return await callback.async_run_agentic_loop(
agentic_result = await callback.async_run_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
Expand All @@ -4832,6 +4868,7 @@ async def _call_agentic_completion_hooks(
stream=stream,
kwargs=kwargs_with_provider,
)
return self._maybe_wrap_in_fake_stream(agentic_result, logging_obj)

plan = await callback.async_build_agentic_loop_plan(
tools=tool_calls,
Expand All @@ -4846,29 +4883,34 @@ async def _call_agentic_completion_hooks(
)

if plan.response_override is not None:
return plan.response_override
return self._maybe_wrap_in_fake_stream(
plan.response_override, logging_obj
)
if plan.terminate:
verbose_logger.debug(
"Agentic loop terminated by callback=%s reason=%s",
callback.__class__.__name__,
plan.stop_reason,
)
return response
return self._maybe_wrap_in_fake_stream(response, logging_obj)
if not plan.run_agentic_loop:
continue

return await self._execute_anthropic_agentic_plan(
plan=plan,
model=model,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs_with_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
stream=stream,
return self._maybe_wrap_in_fake_stream(
await self._execute_anthropic_agentic_plan(
plan=plan,
model=model,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs_with_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
stream=stream,
),
logging_obj,
)
except Exception as e:
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
Expand All @@ -4885,37 +4927,9 @@ async def _call_agentic_completion_hooks(
# 1. Stream was originally True but converted to False for WebSearch interception
# 2. No agentic loop ran (LLM didn't use the tool)
# 3. We have a non-streaming response that needs to be converted to streaming
websearch_converted_stream = (
logging_obj.model_call_details.get(
"websearch_interception_converted_stream", False
)
if logging_obj is not None
else False
)

if websearch_converted_stream:
from typing import cast

from litellm._logging import verbose_logger
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)

verbose_logger.debug(
"WebSearchInterception: No tool call made, converting non-streaming response to fake stream"
)

# Convert the non-streaming response to a fake stream
# The response should be an AnthropicMessagesResponse (dict)
if isinstance(response, dict):
# Create a fake streaming iterator
fake_stream = FakeAnthropicMessagesStreamIterator(
response=cast(AnthropicMessagesResponse, response)
)
return fake_stream
result = self._maybe_wrap_in_fake_stream(response, logging_obj)
if result is not response:
return result

return None

Expand Down
70 changes: 70 additions & 0 deletions tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
"""
Unit tests for _maybe_wrap_in_fake_stream in BaseLLMHTTPHandler.

Tests that agentic loop responses are correctly wrapped in
FakeAnthropicMessagesStreamIterator when the original request was streaming.
"""

from unittest.mock import MagicMock


from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)


class TestMaybeWrapInFakeStream:
def setup_method(self):
self.handler = BaseLLMHTTPHandler()

def test_wraps_dict_when_converted_stream_flag_is_true(self):
"""When websearch_interception_converted_stream is True and response is dict, wrap it."""
logging_obj = MagicMock()
logging_obj.model_call_details = {
"websearch_interception_converted_stream": True
}
response = {
"id": "msg_123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "hello"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 5},
}

result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)

assert isinstance(result, FakeAnthropicMessagesStreamIterator)

def test_returns_response_unchanged_when_flag_is_false(self):
"""When flag is False, return response as-is."""
logging_obj = MagicMock()
logging_obj.model_call_details = {
"websearch_interception_converted_stream": False
}
response = {"id": "msg_123", "content": []}

result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)

assert result is response

def test_returns_response_unchanged_when_not_dict(self):
"""When response is not a dict (e.g., already a stream), return as-is."""
logging_obj = MagicMock()
logging_obj.model_call_details = {
"websearch_interception_converted_stream": True
}
response = MagicMock() # Not a dict

result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)

assert result is response

def test_returns_response_when_logging_obj_is_none(self):
"""When logging_obj is None, return response as-is."""
response = {"id": "msg_123", "content": []}

result = self.handler._maybe_wrap_in_fake_stream(response, None)

assert result is response
Loading