diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index fa1253d9005..4f844e2be09 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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 diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py new file mode 100644 index 00000000000..6b38a703427 --- /dev/null +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_streaming_wrap.py @@ -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