diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index c92e32aeb7f..a38b8bed595 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -26,6 +26,7 @@ from litellm.integrations.websearch_interception.tools import ( get_litellm_web_search_tool_openai, get_litellm_web_search_tool_responses, is_anthropic_native_web_search_tool, + is_chat_completion_request_web_search_tool, is_web_search_tool, is_web_search_tool_chat_completion, is_web_search_tool_responses, @@ -457,7 +458,10 @@ class WebSearchInterceptionLogger(CustomLogger): if call_type in (CallTypes.responses, CallTypes.aresponses): return self._convert_responses_tools(kwargs=kwargs, tools=tools) - has_websearch: Final = any(is_web_search_tool(t) for t in tools) + is_search_tool: Final = ( + is_chat_completion_request_web_search_tool if call_type == CallTypes.acompletion else is_web_search_tool + ) + has_websearch: Final = any(is_search_tool(t) for t in tools) if not has_websearch: return None @@ -475,7 +479,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Convert native/custom web_search tools to LiteLLM standard converted_tools: Final = [] for tool in tools: - if is_web_search_tool(tool): + if is_search_tool(tool): # Convert to LiteLLM standard web search tool converted_tool = get_litellm_web_search_tool_openai() converted_tools.append(converted_tool) diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index 2e1ae07eb68..36b02c48dc5 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -160,6 +160,10 @@ def is_web_search_tool_responses(tool: Mapping[str, object]) -> bool: return tool_type == "web_search" or tool_type.startswith("web_search_") +def is_chat_completion_request_web_search_tool(tool: Mapping[str, object]) -> bool: + return tool.get("type") == "web_search" or is_web_search_tool(dict(tool)) + + def is_web_search_tool_chat_completion(tool: dict[str, Any]) -> bool: """ Check if a tool is a web search tool for Chat Completions API (strict check). diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py index 74fad7deaf1..b9d39b02afd 100644 --- a/tests/integration/streaming/test_interception_converted_stream_options.py +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -61,6 +61,13 @@ _HEADROOM_RETRIEVE: Final[Mapping[str, JsonValue]] = MappingProxyType( } ) _RESPONSES_SEARCH: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "web_search"}) +_BARE_WEB_SEARCH: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "web_search"}) +_APP_WEB_SEARCH: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "type": "function", + "function": {"name": "web_search", "description": "Search the app index", "parameters": dict(_PARAMETERS)}, + } +) _BURST_MODEL: Final = f"interception-burst-{uuid.uuid4().hex}" _BURST_DEPLOYMENT: Final = iv.deployment("search") _STARTED_WORKER: Final = "Started server process [" @@ -460,6 +467,70 @@ def test_openai_sdk_chat_without_streaming_through_web_search_answers_in_json(ri _assert_search_turn(rig, deployment, completion.id) +def _assert_answered_in_json(reply: _Stream, text: str) -> str: + assert reply.status == 200, reply.text + completion: Final = iv.JSON_OBJECT.validate_json(reply.text) + (choice,) = rv.ITEMS.validate_python(completion["choices"]) + assert object_value(choice["message"])["content"] == text, completion + return string_value(completion["id"]) + + +def _offered_tools(call: Mapping[str, JsonValue]) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(rv.ITEMS.validate_python(call["tools"])) + + +def _function_name(tool: Mapping[str, JsonValue]) -> str: + return string_value(object_value(tool["function"])["name"]) + + +@pytest.mark.parametrize("stream", (pytest.param(True, id="stream"), pytest.param(False, id="json"))) +def test_bare_web_search_type_on_chat_is_searched_and_the_apps_own_web_search_function_reaches_azure_untouched( + rig: _Rig, scenario: Scenario, stream: bool +) -> None: + deployment: Final = _deploy(rig, scenario, "search") + body: Final = { + **_chat_body(deployment, stream=stream), + "tools": [dict(_BARE_WEB_SEARCH), dict(_APP_WEB_SEARCH)], + } + reply: Final = _chat(rig, body) + identity: Final = ( + _assert_restreamed(reply, iv.answer(deployment.name), usage=None) + if stream + else _assert_answered_in_json(reply, iv.answer(deployment.name)) + ) + received: Final = _upstream(rig, deployment) + _assert_upstream(received, model_calls=2, searches=1) + assert [search["query"] for search in received.searches] == [iv.query(deployment.name)], received.searches + offered: Final = _offered_tools(received.model_calls[0]) + assert [tool.get("type") for tool in offered] == ["function", "function"], offered + assert [_function_name(tool) for tool in offered] == [iv.SEARCH_TOOL, "web_search"], offered + assert offered[1] == dict(_APP_WEB_SEARCH), offered + _assert_landed_once(deployment.model, frozenset({identity}), rows=1) + + +def test_bare_web_search_type_on_messages_through_the_azure_chat_adapter_is_never_searched( + rig: _Rig, scenario: Scenario +) -> None: + deployment: Final = _deploy(rig, scenario, "search") + reply: Final = _post( + rig, + "/v1/messages", + { + "model": deployment.model, + "max_tokens": 64, + "messages": [{"role": "user", "content": f"Search for {deployment.name}"}], + "tools": [dict(_BARE_WEB_SEARCH)], + **_NO_CACHE, + }, + ) + assert reply.status == 200, reply.text + blocks: Final = rv.ITEMS.validate_python(iv.JSON_OBJECT.validate_json(reply.text)["content"]) + assert [block.get("text") for block in blocks] == [iv.plain(deployment.name)], reply.text + received: Final = _upstream(rig, deployment) + _assert_upstream(received, model_calls=1, searches=0) + assert [call.get("tools") for call in received.model_calls] == [None], received.model_calls + + @pytest.mark.parametrize( ("tool", "mode"), ( diff --git a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py index 064b94942a3..91aa18d2440 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py @@ -669,6 +669,52 @@ async def test_async_pre_call_deployment_hook_returns_full_kwargs(): ) +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["anthropic", "openai"]) +async def test_deployment_hook_intercepts_bare_web_search_type_and_keeps_user_web_search_function(provider: str): + user_web_search_function: Final = { + "type": "function", + "function": { + "name": "web_search", + "description": "Search a private index", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + }, + } + logger: Final = WebSearchInterceptionLogger(enabled_providers=[provider]) + + result: Final = await logger.async_pre_call_deployment_hook( + kwargs={ + "model": "some-model", + "messages": [{"role": "user", "content": "Who won the latest Nobel Prize in Physics?"}], + "tools": [{"type": "web_search"}, user_web_search_function], + "custom_llm_provider": provider, + }, + call_type=CallTypes.acompletion, + ) + + assert result is not None + assert [t.get("function", {}).get("name") for t in result["tools"]] == [LITELLM_WEB_SEARCH_TOOL_NAME, "web_search"] + assert result["tools"][1] == user_web_search_function + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", [CallTypes.anthropic_messages, CallTypes.aanthropic_messages]) +async def test_deployment_hook_leaves_bare_web_search_type_alone_on_messages(call_type: CallTypes): + logger: Final = WebSearchInterceptionLogger(enabled_providers=["openai"]) + + result: Final = await logger.async_pre_call_deployment_hook( + kwargs={ + "model": "some-model", + "messages": [{"role": "user", "content": "Who won the latest Nobel Prize in Physics?"}], + "tools": [{"type": "web_search"}], + "custom_llm_provider": "openai", + }, + call_type=call_type, + ) + + assert result is None + + @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_skips_disabled_provider(): """Test that the hook returns None for providers not in enabled_providers.""" @@ -876,7 +922,6 @@ async def test_converted_stream_moves_stream_options_out_of_the_provider_request ) == (False, False, True, _STREAM_USAGE_OPTIONS) - def test_sync_forced_tool_choice_repoints_converted_web_search(): """Regression (tool_choice 400): a forced tool_choice naming the original web_search tool must be repointed to litellm_web_search after conversion.