From 0274d8496b3448e068debe399cd8111d71b4c95e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 16:33:33 -0800 Subject: [PATCH] fix(responses): replay a BM25 tool search under its own name --- .../tool_search_results.py | 16 +++++++++- .../test_anthropic_responses_bridge.py | 30 +++++++++++++++++++ .../test_litellm_completion_responses.py | 30 +++++++++++++++++++ 3 files changed, 75 insertions(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/tool_search_results.py b/litellm/responses/litellm_completion_transformation/tool_search_results.py index 4aa366854fe..566e8cb6414 100644 --- a/litellm/responses/litellm_completion_transformation/tool_search_results.py +++ b/litellm/responses/litellm_completion_transformation/tool_search_results.py @@ -32,6 +32,7 @@ _MessageT: Final = TypeVar("_MessageT") _MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _TOOL_SEARCH_ITEM_TYPES: Final = frozenset({"tool_search_call", "tool_search_output"}) +_BM25_TOOL_SEARCH_TOOL_NAME: Final = "tool_search_tool_bm25" class _ProviderToolResults(LiteLLMBaseModel): @@ -213,13 +214,26 @@ def _anthropic_tool_search_result(output: _ToolSearchOutputItem) -> Mapping[str, return result.model_dump(mode="json") +def _search_arguments(arguments: object) -> Mapping[str, object]: + try: + if isinstance(arguments, str): + return _MAPPING_ADAPTER.validate_json(arguments) + return _MAPPING_ADAPTER.validate_python(arguments) + except ValidationError: + return MappingProxyType({}) + + +def _search_tool_name(arguments: object) -> str: + return _BM25_TOOL_SEARCH_TOOL_NAME if "query" in _search_arguments(arguments) else ANTHROPIC_TOOL_SEARCH_TOOL_NAME + + def _replayed_item(item: ResponseInputItemParam, replayed_ids: frozenset[str]) -> ResponseInputItemParam | None: call: Final = _tool_search_call(item) if call is not None and call.call_id in replayed_ids: return ResponseFunctionToolCallParam( type="function_call", call_id=call.call_id, - name=ANTHROPIC_TOOL_SEARCH_TOOL_NAME, + name=_search_tool_name(call.arguments), arguments=serialize_tool_call_arguments(call.arguments, "{}"), ) output: Final = _tool_search_output(item) diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index c7dbabb7f0b..3202b0d7d37 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -1651,6 +1651,36 @@ def test_responses_over_anthropic_return_the_tool_search_output_and_send_the_sea assert _sent_assistant_content(route) == [_SEARCH_USE, search_result, weather_use] + +@pytest.mark.parametrize( + ("search_tools", "search_use"), + [ + pytest.param((), _SEARCH_USE, id="regex"), + pytest.param( + ({"type": "tool_search_tool_bm25_20251119", "name": "tool_search_tool_bm25"},), + {**_SEARCH_USE, "name": "tool_search_tool_bm25", "input": {"query": "weather"}}, + id="bm25", + ), + ], +) +def test_responses_over_anthropic_send_a_search_back_under_the_anthropic_search_tool_that_ran_it( + respx_mock: respx.MockRouter, search_tools: tuple[Mapping[str, object], ...], search_use: Mapping[str, object] +) -> None: + tools: Final = [{"type": "tool_search"}, *search_tools, _WEATHER_TOOL] + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message(search_use, _SEARCH_RESULT, _WEATHER_USE) + ) + + first: Final = litellm.responses(model=_MODEL, input=list(_WEATHER_QUESTION), tools=tools, api_key="mock_api_key") + litellm.responses( + model=_MODEL, + input=[*_WEATHER_QUESTION, *(item.model_dump(exclude_none=True) for item in first.output), _WEATHER_ANSWER], + tools=tools, + api_key="mock_api_key", + ) + + assert _sent_assistant_content(route) == [search_use, _SEARCH_RESULT, _WEATHER_USE] + def _tool_search_stream(*events: Mapping[str, object]) -> httpx.Response: return httpx.Response( 200, diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 9ce9d68f48c..ea1e3123c20 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2460,6 +2460,36 @@ class TestToolTransformation: assert [call["id"] for call in request["messages"][1]["tool_calls"]] == replayed_call_ids + @pytest.mark.parametrize( + ("arguments", "search_tool_name"), + [ + pytest.param({"pattern": "weather"}, "tool_search_tool_regex", id="regex_object"), + pytest.param({"query": "weather"}, "tool_search_tool_bm25", id="bm25_object"), + pytest.param('{"query": "weather"}', "tool_search_tool_bm25", id="bm25_json"), + pytest.param("weather", "tool_search_tool_regex", id="not_json"), + ], + ) + def test_stateless_claude_search_is_replayed_under_the_search_tool_its_arguments_belong_to( + self, arguments: object, search_tool_name: str + ) -> None: + items: Final = ( + {**item, "arguments": arguments} if item["type"] == "tool_search_call" else item + for item in _stateless_tool_search_input("srvtoolu_01Search", "completed") + ) + request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( # cast-ok: tool search items are not in the SDK's input TypedDicts + ResponseInputParam, list(items) + ), + responses_api_request={"tools": _deferred_tools_with_server_tool_search()}, + custom_llm_provider="anthropic", + ) + + assert [(call["id"], call["function"]["name"]) for call in request["messages"][1]["tool_calls"]] == [ + ("srvtoolu_01Search", search_tool_name), + ("toolu_01Weather", "get_weather"), + ] + def test_transform_local_shell_tools_dropped(self) -> None: ( result_tools,