scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 69

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

Browse source 
When websearch_interception converts stream=True to stream=False internally,
the agentic loop returns a plain dict. Previously this dict was returned
directly to the client expecting SSE events, resulting in empty streams.

Added _maybe_wrap_in_fake_stream() which checks the
websearch_interception_converted_stream flag and wraps dict responses in
FakeAnthropicMessagesStreamIterator. Applied to all return paths in
_call_agentic_completion_hooks:
- async_run_agentic_loop (legacy path)
- _execute_anthropic_agentic_plan (plan-based path)
- plan.response_override
- plan.terminate

Includes unit tests for _maybe_wrap_in_fake_stream().
This commit is contained in:
Clawd 2026-05-08 15:20:51 +08:00
parent 98cd057f38
commit c3b5e07627
2 changed files with 130 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

106
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -4740,6 +4740,42 @@ class BaseLLMHTTPHandler:
**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,
@ -4821,7 +4857,7 @@ class BaseLLMHTTPHandler:
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,
@ -4832,6 +4868,7 @@ class BaseLLMHTTPHandler:
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,
@ -4846,29 +4883,34 @@ class BaseLLMHTTPHandler:
)
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")
@ -4885,37 +4927,9 @@ class BaseLLMHTTPHandler:
# 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

70
tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py Normal file
View file

@ -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