fix(websearch): preserve forced search and Messages streaming

litellm/integrations/websearch_interception/handler.py

@@ -37,6 +37,7 @@ from litellm.integrations.websearch_interception.transformation import (
 from litellm.litellm_core_utils.agentic_loop_settings import (
     validated_max_agentic_loops,
 )
+from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.llms.base_llm.search.transformation import SearchResponse
 from litellm.types.integrations.custom_logger import (
     CHAT_COMPLETION_AGENTIC_SURFACE,
@@ -514,7 +515,13 @@ class WebSearchInterceptionLogger(CustomLogger):
             get_litellm_web_search_tool_responses() if is_web_search_tool_responses(tool) else tool for tool in tools
         ]
 
-        converted_kwargs: Final = {**kwargs, "tools": converted_tools}
+        tool_choice: Final = as_str_mapping(kwargs.get("tool_choice"))
+        choice_patch: Final = (
+            {"tool_choice": {"type": "function", "name": LITELLM_WEB_SEARCH_TOOL_NAME}}
+            if tool_choice is not None and is_web_search_tool_responses(tool_choice)
+            else {}
+        )
+        converted_kwargs: Final = {**kwargs, "tools": converted_tools, **choice_patch}
 
         if kwargs.get("stream"):
             verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")
@@ -1266,7 +1273,7 @@ class WebSearchInterceptionLogger(CustomLogger):
             model=full_model_name,
             messages=input_list,
             tools=tools_param if isinstance(tools_param, list) else None,
-            optional_params=optional_params_clean,
+            optional_params={**optional_params_clean, "tool_choice": "auto"},
             kwargs=kwargs_for_followup,
         )
 

litellm/llms/anthropic/messages/native/fake_stream_iterator.py

@@ -93,7 +93,7 @@ class FakeAnthropicMessagesStreamIterator:
             content_block_start = {
                 "type": "content_block_start",
                 "index": index,
-                "content_block": {"type": "redacted_thinking"},
+                "content_block": block_dict,
             }
             chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
 
@@ -102,10 +102,8 @@ class FakeAnthropicMessagesStreamIterator:
                 "type": "content_block_start",
                 "index": index,
                 "content_block": {
-                    "type": "tool_use",
-                    "id": block_dict.get("id"),
-                    "name": block_dict.get("name"),
-                    "input": {},
+                    **block_dict,
+                    "input": dict[str, object](),
                 },
             }
             chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())

litellm/utils.py

@@ -1014,6 +1014,33 @@ def _is_converted_stream_result(result: object) -> bool:
     return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator))
 
 
+_MESSAGES_REPLAY_REQUEST: Final = TypeAdapter(Mapping[str, object])
+_MESSAGES_REPLAY_RESPONSE: Final = TypeAdapter(object)
+
+
+def _is_buffered_messages_stream(result: object, request_data: Mapping[str, object], call_type: str) -> bool:
+    from litellm.types.integrations.custom_logger import converted_stream_requested
+
+    return (
+        call_type == CallTypes.anthropic_messages.value
+        and isinstance(result, dict)
+        and (bool(request_data.get("stream")) or converted_stream_requested(request_data))
+    )
+
+
+def _restore_buffered_messages_stream(result: object, request_data: Mapping[str, object], call_type: str) -> object:
+    if not _is_buffered_messages_stream(result, request_data, call_type):
+        return result
+    from litellm.llms.anthropic.messages.native.fake_stream_iterator import FakeAnthropicMessagesStreamIterator
+    from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
+
+    return FakeAnthropicMessagesStreamIterator(
+        response=cast(  # cast-ok: adapter-validated Messages body; preserve provider fields
+            AnthropicMessagesResponse, result
+        )
+    )
+
+
 async def _run_success_deployment_hook_on_converted_chat_stream(
     result: object, request_data: dict[str, object], call_type: str
 ) -> None:
@@ -2205,6 +2232,12 @@ def client(original_function):
             if modified_kwargs is not None:
                 kwargs = modified_kwargs
 
+            messages_replay_request: Final[Mapping[str, object]] = (
+                _MESSAGES_REPLAY_REQUEST.validate_python(kwargs)
+                if call_type == CallTypes.anthropic_messages.value
+                else {}
+            )
+
             # Sync logging_obj.stream after deployment hooks (they may convert it).
             _hook_stream: Final = kwargs.get("stream")
             if (
@@ -2252,7 +2285,9 @@ def client(original_function):
                     if _is_converted_stream_result(_caching_handler_response.cached_result):
                         logging_obj.stream = True
                         logging_obj.model_call_details["stream"] = True
-                    return _caching_handler_response.cached_result
+                    return _restore_buffered_messages_stream(
+                        _caching_handler_response.cached_result, messages_replay_request, call_type
+                    )
 
                 elif _caching_handler_response.embedding_all_elements_cache_hit is True:
                     return _caching_handler_response.final_embedding_cached_response
@@ -2316,7 +2351,10 @@ def client(original_function):
             end_time = datetime.datetime.now()
 
             streaming_requested: Final = _is_streaming_request(kwargs=kwargs, call_type=call_type)
-            if streaming_requested or _is_converted_stream_result(result):
+            buffered_messages_stream: Final = _is_buffered_messages_stream(
+                _MESSAGES_REPLAY_RESPONSE.validate_python(result), messages_replay_request, call_type
+            )
+            if (streaming_requested and not buffered_messages_stream) or _is_converted_stream_result(result):
                 logging_obj.stream = True
                 logging_obj.model_call_details["stream"] = True
                 if not streaming_requested:
@@ -2406,7 +2444,9 @@ def client(original_function):
                 is_litellm_internal_call=_is_litellm_internal_call,
             )
 
-            return result
+            return _restore_buffered_messages_stream(
+                _MESSAGES_REPLAY_RESPONSE.validate_python(result), messages_replay_request, call_type
+            )
         except Exception as e:
             traceback_exception: Final = traceback.format_exc()
             # Reuse the timestamp taken right when the deployment call itself failed, before

tests/unit/integrations/websearch_interception/test_websearch_responses.py

@@ -6,11 +6,16 @@ tool calls returned by /v1/responses, executes the search server-side, and
 builds a Responses-format follow-up request.
 """
 
+import json
 from types import SimpleNamespace
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
+import respx
 
+import litellm
 from litellm.integrations.websearch_interception.handler import (
     WebSearchInterceptionLogger,
 )
@@ -247,7 +252,7 @@ async def test_build_responses_plan_produces_responses_input():
     }
 
     assert patch_obj.tools == [{"type": "function", "name": "litellm_web_search"}]
-    assert "tool_choice" not in patch_obj.optional_params
+    assert patch_obj.optional_params["tool_choice"] == "auto"
     assert "_agentic_loop_api_surface" not in patch_obj.kwargs
     assert patch_obj.model == "openai/gpt-4o"
 
@@ -316,3 +321,169 @@ async def test_deployment_hook_responses_converts_stream_to_non_stream():
     assert result is not None
     assert result["stream"] is False
     assert result["_websearch_interception_converted_stream"] is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("search_type", ["web_search", "web_search_preview", "web_search_preview_2025_03_11"])
+async def test_forced_responses_search_names_the_converted_function(search_type: str, stream: bool) -> None:
+    logger: Final = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.BEDROCK])
+    choice: Final = {"type": search_type}
+    result: Final = await logger.async_pre_call_deployment_hook(
+        kwargs={
+            "model": "us.openai.gpt-6.1-sol",
+            "custom_llm_provider": "bedrock",
+            "tools": [{"type": search_type}],
+            "tool_choice": choice,
+            "stream": stream,
+        },
+        call_type=CallTypes.aresponses,
+    )
+
+    assert result is not None
+    assert result["tool_choice"] == {"type": "function", "name": result["tools"][0]["name"]}
+    assert result["stream"] is False
+    assert choice == {"type": search_type}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "choice",
+    [
+        None,
+        "auto",
+        "none",
+        "required",
+        {"type": "function", "name": "get_weather"},
+        {"type": "function", "name": "litellm_web_search"},
+    ],
+)
+async def test_responses_search_conversion_preserves_other_tool_choices(choice: object) -> None:
+    logger: Final = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.BEDROCK])
+    weather: Final = {"type": "function", "name": "get_weather", "parameters": {"type": "object"}}
+    result: Final = await logger.async_pre_call_deployment_hook(
+        kwargs={
+            "model": "us.openai.gpt-6.1-sol",
+            "custom_llm_provider": "bedrock",
+            "tools": [{"type": "web_search"}, weather],
+            "tool_choice": choice,
+        },
+        call_type=CallTypes.aresponses,
+    )
+
+    assert result is not None
+    assert result["tool_choice"] == choice
+    assert result["tools"][1] is weather
+
+
+@pytest.mark.asyncio
+async def test_forced_responses_search_executes_once_then_allows_the_answer(monkeypatch: pytest.MonkeyPatch) -> None:
+    from litellm.proxy import proxy_server
+
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setenv("TAVILY_API_BASE", "https://search.example.test")
+    monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["bedrock"])])
+    router: Final = litellm.Router(
+        model_list=[
+            {
+                "model_name": "search-canary",
+                "litellm_params": {
+                    "model": "bedrock/us.openai.gpt-6.1-sol",
+                    "api_key": "synthetic-bedrock-bearer",
+                    "aws_region_name": "us-east-1",
+                },
+            }
+        ],
+        search_tools=[
+            {
+                "search_tool_name": "fixture-search",
+                "litellm_params": {
+                    "search_provider": "tavily",
+                    "api_key": "synthetic-search-key",
+                    "api_base": "https://search.example.test",
+                },
+            }
+        ],
+        num_retries=0,
+    )
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+    common: Final = {
+        "object": "response",
+        "created_at": 1,
+        "status": "completed",
+        "model": "us.openai.gpt-6.1-sol",
+        "usage": {"input_tokens": 11, "output_tokens": 5, "total_tokens": 16},
+    }
+    with respx.mock as upstream:
+        search: Final = upstream.post("https://search.example.test/search").respond(
+            200,
+            json={
+                "results": [
+                    {
+                        "title": "Synthetic search",
+                        "url": "https://example.test/search",
+                        "content": "Fixture evidence",
+                    }
+                ]
+            },
+        )
+        responses_route: Final = upstream.post(
+            "https://bedrock-runtime.us-east-1.amazonaws.com/openai/v1/responses"
+        ).mock(
+            side_effect=[
+                httpx.Response(
+                    200,
+                    json={
+                        **common,
+                        "id": "resp_search",
+                        "output": [
+                            {
+                                "type": "function_call",
+                                "id": "fc_search",
+                                "call_id": "call_search",
+                                "name": "litellm_web_search",
+                                "arguments": '{"query":"synthetic query"}',
+                                "status": "completed",
+                            }
+                        ],
+                    },
+                ),
+                httpx.Response(
+                    200,
+                    json={
+                        **common,
+                        "id": "resp_answer",
+                        "output": [
+                            {
+                                "type": "message",
+                                "id": "msg_answer",
+                                "role": "assistant",
+                                "status": "completed",
+                                "content": [{"type": "output_text", "text": "Answer from search", "annotations": []}],
+                            }
+                        ],
+                    },
+                ),
+            ]
+        )
+        result: Final = await router.aresponses(
+            model="search-canary",
+            input="Search for synthetic evidence",
+            max_output_tokens=100,
+            tools=[{"type": "web_search"}],
+            tool_choice={"type": "web_search"},
+        )
+
+    assert result.output[0].content[0].text == "Answer from search"
+    assert len(responses_route.calls) == 2
+    first: Final = json.loads(responses_route.calls[0].request.content)
+    followup: Final = json.loads(responses_route.calls[1].request.content)
+    assert first["tool_choice"] == {"type": "function", "name": first["tools"][0]["name"]}
+    assert followup["tool_choice"] == "auto"
+    assert followup["input"][-1] == {
+        "type": "function_call_output",
+        "call_id": "call_search",
+        "output": "Title: Synthetic search\nURL: https://example.test/search\nSnippet: Fixture evidence",
+    }
+    assert len(search.calls) == 1
+    assert json.loads(search.calls[0].request.content)["query"] == "synthetic query"

tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py

@@ -2,34 +2,38 @@ import asyncio
 import copy
 import json
 import os
+import unittest.mock
 import uuid
+from itertools import chain
 from typing import Any, Dict, Final, List
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
+import pytest
 import responses
 import respx
-import pytest
 from fastapi.testclient import TestClient
 from pydantic import ValidationError
 
-
-from unittest.mock import AsyncMock, MagicMock, patch
-
 import litellm
 from litellm.anthropic_interface import messages
+from litellm.caching.caching import Cache
+from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.router import Router
 from litellm.types.utils import (
+    CallTypes,
     Delta,
+    LLMResponseTypes,
     ModelResponse,
     StandardLoggingPayloadErrorInformation,
     StreamingChoices,
 )
 from litellm.utils import _invalidate_model_cost_lowercase_map
-import unittest.mock
 
 
 @pytest.mark.asyncio
@@ -2396,3 +2400,145 @@ def test_native_messages_tool_config_point_marks_last_tool_and_drops_the_key(mon
     assert "cache_control" not in captured["params"]["tools"][0]
     assert captured["messages"][0]["content"][-1]["cache_control"] == {"type": "ephemeral"}
     assert tools == tools_before
+
+
+class _MessagesAnswerRewrite(CustomLogger):
+    async def async_post_call_success_deployment_hook(
+        self, request_data: dict, response: LLMResponseTypes, call_type: CallTypes | None
+    ) -> LLMResponseTypes | None:
+        if call_type == CallTypes.anthropic_messages and isinstance(response, dict):
+            return {**response, "content": [{"type": "text", "text": "Reviewed answer"}]}
+        return None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("tool_call", [False, True], ids=["text", "ordinary-tool"])
+@pytest.mark.parametrize("rewrite", [False, True], ids=["original", "success-hook"])
+@pytest.mark.parametrize("cache_enabled", [False, True], ids=["uncached", "response-cache"])
+async def test_bedrock_messages_mixed_search_tools_preserve_the_client_stream(
+    monkeypatch: pytest.MonkeyPatch, stream: bool, tool_call: bool, rewrite: bool, cache_enabled: bool
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [WebSearchInterceptionLogger(enabled_providers=["bedrock"])] + ([_MessagesAnswerRewrite()] if rewrite else []),
+    )
+    if cache_enabled:
+        monkeypatch.setattr(litellm, "cache", Cache(type="local"))
+    content: Final = (
+        [{"toolUse": {"toolUseId": "call_weather", "name": "get_weather", "input": {"city": "Paris"}}}]
+        if tool_call
+        else [{"text": "No search needed"}]
+    )
+    stop_reason: Final = "tool_use" if tool_call else "end_turn"
+    with respx.mock as upstream:
+        converse: Final = upstream.post(
+            "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.openai.gpt-6.1-sol/converse"
+        ).respond(
+            200,
+            json={
+                "output": {"message": {"role": "assistant", "content": content}},
+                "stopReason": stop_reason,
+                "usage": {"inputTokens": 11, "outputTokens": 5, "totalTokens": 16},
+                "metrics": {"latencyMs": 1},
+            },
+        )
+        result: Final = await litellm.anthropic.messages.acreate(
+            model="bedrock/us.openai.gpt-6.1-sol",
+            api_key="synthetic-bedrock-bearer",
+            aws_region_name="us-east-1",
+            max_tokens=100,
+            messages=[{"role": "user", "content": "Answer or call get_weather for Paris without searching."}],
+            tools=[
+                {"type": "web_search_20250305", "name": "web_search"},
+                {
+                    "name": "get_weather",
+                    "description": "Get weather",
+                    "input_schema": {
+                        "type": "object",
+                        "properties": {"city": {"type": "string"}},
+                        "required": ["city"],
+                    },
+                },
+            ],
+            stream=stream,
+        )
+        assert len(converse.calls) == 1
+        assert "stream" not in json.loads(converse.calls.last.request.content)
+        if cache_enabled:
+            await asyncio.gather(*tuple(_PENDING_CACHE_WRITES))
+            replay: Final = await litellm.anthropic.messages.acreate(
+                model="bedrock/us.openai.gpt-6.1-sol",
+                api_key="synthetic-bedrock-bearer",
+                aws_region_name="us-east-1",
+                max_tokens=100,
+                messages=[{"role": "user", "content": "Answer or call get_weather for Paris without searching."}],
+                tools=[
+                    {"type": "web_search_20250305", "name": "web_search"},
+                    {
+                        "name": "get_weather",
+                        "description": "Get weather",
+                        "input_schema": {
+                            "type": "object",
+                            "properties": {"city": {"type": "string"}},
+                            "required": ["city"],
+                        },
+                    },
+                ],
+                stream=stream,
+            )
+            if stream:
+                original_chunks: Final = tuple([chunk async for chunk in result])
+                replay_chunks: Final = tuple([chunk async for chunk in replay])
+                assert replay_chunks == original_chunks
+            else:
+                assert replay["content"] == result["content"]
+            assert len(converse.calls) == 1
+
+    if not stream:
+        assert result["stop_reason"] == stop_reason
+        assert result["content"] == (
+            [{"type": "text", "text": "Reviewed answer"}]
+            if rewrite
+            else (
+                [
+                    {"type": "text", "text": ""},
+                    {"type": "tool_use", "id": "call_weather", "name": "get_weather", "input": {"city": "Paris"}},
+                ]
+                if tool_call
+                else [{"type": "text", "text": "No search needed"}]
+            )
+        )
+        return
+
+    assert hasattr(result, "__aiter__"), result
+    chunks: Final = original_chunks if cache_enabled else tuple([chunk async for chunk in result])
+    lines: Final = chain.from_iterable(chunk.decode().splitlines() for chunk in chunks)
+    events: Final = [json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: ")]
+    assert events[0]["type"] == "message_start"
+    assert events[-1] == {"type": "message_stop", "usage": {"input_tokens": 11, "output_tokens": 5}}
+    assert events[0]["message"]["usage"]["input_tokens"] == 11
+    assert events[-2]["usage"]["output_tokens"] == 5
+    assert events[-2]["delta"]["stop_reason"] == stop_reason
+    starts: Final = [event["content_block"] for event in events if event["type"] == "content_block_start"]
+    deltas: Final = [event["delta"] for event in events if event["type"] == "content_block_delta"]
+    assert starts[-1] == (
+        {"type": "text", "text": ""}
+        if rewrite
+        else (
+            {"type": "tool_use", "id": "call_weather", "name": "get_weather", "input": {}}
+            if tool_call
+            else {"type": "text", "text": ""}
+        )
+    )
+    assert deltas[-1] == (
+        {"type": "text_delta", "text": "Reviewed answer"}
+        if rewrite
+        else (
+            {"type": "input_json_delta", "partial_json": json.dumps({"city": "Paris"})}
+            if tool_call
+            else {"type": "text_delta", "text": "No search needed"}
+        )
+    )

tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py

@@ -91,3 +91,25 @@ def test_message_delta_omits_absent_optional_fields():
         "usage": {"input_tokens": 10, "output_tokens": 2},
     }
     assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None}
+
+
+def test_replayed_stream_preserves_redacted_reasoning_and_tool_signatures():
+    redacted: Final = {"type": "redacted_thinking", "data": "opaque-continuation-data"}
+    tool: Final = {
+        "type": "tool_use",
+        "id": "call_signed",
+        "name": "get_weather",
+        "input": {"city": "Paris"},
+        "provider_specific_fields": {"thought_signature": "opaque-tool-signature"},
+    }
+    response: Final = {
+        "id": "msg_signed",
+        "content": [redacted, tool],
+        "stop_reason": "tool_use",
+        "usage": {"input_tokens": 11, "output_tokens": 5},
+    }
+    events: Final = _decoded_events(FakeAnthropicMessagesStreamIterator(response=response))
+    starts: Final = tuple(payload["content_block"] for kind, payload in events if kind == "content_block_start")
+    assert starts == (redacted, {**tool, "input": {}})
+    deltas: Final = tuple(payload["delta"] for kind, payload in events if kind == "content_block_delta")
+    assert deltas == ({"type": "input_json_delta", "partial_json": json.dumps(tool["input"])},)