diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 13a16947fb4..d025590def9 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -329,15 +329,25 @@ class WebSearchInterceptionLogger(CustomLogger): if not tools: return None - if call_type in (CallTypes.responses, CallTypes.aresponses): - return self._convert_responses_tools(kwargs=kwargs, tools=tools) - - # Check if any tool is a web search tool (native or already LiteLLM standard) - has_websearch: Final = any(is_web_search_tool(t) for t in tools) - + is_responses_call: Final = call_type in (CallTypes.responses, CallTypes.aresponses) + has_websearch: Final = ( + any(is_web_search_tool_responses(tool) for tool in tools) + if is_responses_call + else any(is_web_search_tool(tool) for tool in tools) + ) if not has_websearch: return None + if self.search_tool_name: + try: + from litellm.proxy.proxy_server import llm_router + except ImportError: + llm_router = None + self._select_search_tool_from_router(llm_router=llm_router) + + if is_responses_call: + return self._convert_responses_tools(kwargs=kwargs, tools=tools) + verbose_logger.debug("WebSearchInterception: Converting native web_search tools to LiteLLM standard") # If the client sent an Anthropic-native web_search_* tool, mark the @@ -1539,9 +1549,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None": - if llm_router is None or not hasattr(llm_router, "search_tools"): - return None - search_tools: Final = list(getattr(llm_router, "search_tools") or []) + search_tools: Final = list(getattr(llm_router, "search_tools", []) or []) return self._select_search_tool_from_list(search_tools=search_tools, source="router") def _select_search_tool_from_list( @@ -1551,20 +1559,26 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> "_SearchToolConfig | None": if self.search_tool_name: matching_tools = [tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name] - if matching_tools: - search_provider = (matching_tools[0].get("litellm_params", {}) or {}).get("search_provider") - verbose_logger.debug( - "WebSearchInterception: Found search tool '%s' from %s with provider '%s'", - self.search_tool_name, - source, - search_provider, + if not matching_tools: + raise ValueError(f"Configured search tool '{self.search_tool_name}' was not found") + + selected_tool: Final = matching_tools[0] + litellm_params: Final = selected_tool.get("litellm_params") + selected_search_provider: Final = ( + litellm_params.get("search_provider") if isinstance(litellm_params, Mapping) else None + ) + if not isinstance(selected_search_provider, str) or not selected_search_provider.strip(): + raise ValueError( + f"Configured search tool '{self.search_tool_name}' does not define a valid search provider" ) - return matching_tools[0] + verbose_logger.debug( - "WebSearchInterception: Search tool '%s' not found in %s, falling back to first available or perplexity", + "WebSearchInterception: Found search tool '%s' from %s with provider '%s'", self.search_tool_name, source, + selected_search_provider, ) + return selected_tool if search_tools: first_tool: Final = search_tools[0] diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index f39f41a6d12..ec4bc1f49eb 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -14,7 +14,7 @@ from litellm.integrations.websearch_interception.handler import ( ) from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, ProxyException, UserAPIKeyAuth -from litellm.types.utils import LlmProviders +from litellm.types.utils import CallTypes, LlmProviders def test_initialize_from_proxy_config(): @@ -230,6 +230,124 @@ async def test_execute_search_passes_selected_search_tool_litellm_params(monkeyp assert forwarded_kwargs["max_retries"] == 2 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search_tools", "error"), + [ + pytest.param(None, "was not found", id="router-not-configured"), + pytest.param( + [{"search_tool_name": "other-search", "litellm_params": {"search_provider": "tavily"}}], + "was not found", + id="requested-tool-not-configured", + ), + pytest.param( + [{"search_tool_name": "parallel-search", "litellm_params": "not-a-mapping"}], + "does not define a valid search provider", + id="invalid-parameters", + ), + pytest.param( + [{"search_tool_name": "parallel-search", "litellm_params": {}}], + "does not define a valid search provider", + id="missing-provider", + ), + pytest.param( + [{"search_tool_name": "parallel-search", "litellm_params": {"search_provider": " "}}], + "does not define a valid search provider", + id="whitespace-provider", + ), + pytest.param( + [{"search_tool_name": "parallel-search", "litellm_params": {"search_provider": 123}}], + "does not define a valid search provider", + id="invalid-provider", + ), + ], +) +async def test_execute_search_rejects_invalid_explicit_search_tool(monkeypatch, search_tools, error): + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger(search_tool_name="parallel-search") + router = None if search_tools is None else MagicMock(search_tools=search_tools) + mock_asearch = AsyncMock() + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + with pytest.raises(ValueError, match=f"Configured search tool 'parallel-search' {error}"): + await logger._execute_search("what is litellm") + + mock_asearch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_execute_search_honors_explicit_parallel_search_tool(monkeypatch): + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger(search_tool_name="parallel-search") + router = MagicMock( + search_tools=[ + { + "search_tool_name": "other-search", + "litellm_params": {"search_provider": "tavily", "api_key": "other-key"}, + }, + { + "search_tool_name": "parallel-search", + "litellm_params": {"search_provider": "parallel_ai", "api_key": "parallel-key"}, + }, + ], + ) + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("what is litellm") + + mock_asearch.assert_awaited_once_with( + query="what is litellm", + search_provider="parallel_ai", + api_key="parallel-key", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search_tools", "expected_search_kwargs"), + [ + pytest.param(None, {"search_provider": "perplexity"}, id="router-not-configured"), + pytest.param( + [ + { + "search_tool_name": "first-search", + "litellm_params": {"search_provider": "tavily", "api_key": "first-key"}, + }, + { + "search_tool_name": "parallel-search", + "litellm_params": {"search_provider": "parallel_ai", "api_key": "parallel-key"}, + }, + ], + {"search_provider": "tavily", "api_key": "first-key"}, + id="first-configured-tool", + ), + ], +) +async def test_execute_search_preserves_implicit_provider_selection(monkeypatch, search_tools, expected_search_kwargs): + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger() + router = None if search_tools is None else MagicMock(search_tools=search_tools) + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("what is litellm") + + mock_asearch.assert_awaited_once_with(query="what is litellm", **expected_search_kwargs) + + @pytest.mark.asyncio async def test_execute_search_attributes_spend_to_the_calling_key(monkeypatch): """An intercepted search is billed and logged against the key that made the LLM request. @@ -397,6 +515,72 @@ async def test_execute_search_enforces_team_search_tool_permission(monkeypatch): mock_asearch.assert_not_awaited() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "web_search_tool"), + [ + pytest.param( + CallTypes.acompletion, + {"type": "web_search_20250305", "name": "web_search"}, + id="chat-completion", + ), + pytest.param(CallTypes.responses, {"type": "web_search"}, id="responses"), + pytest.param(CallTypes.aresponses, {"type": "web_search"}, id="async-responses"), + pytest.param( + CallTypes.anthropic_messages, + {"type": "web_search_20250305", "name": "web_search"}, + id="anthropic-messages", + ), + ], +) +async def test_deployment_hook_dispatcher_propagates_missing_explicit_search_tool( + monkeypatch, call_type, web_search_tool +): + import litellm + from litellm.proxy import proxy_server + from litellm.utils import async_pre_call_deployment_hook + + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"], search_tool_name="parallel-search") + mock_asearch = AsyncMock() + kwargs = { + "model": "bedrock/claude-sonnet-4", + "tools": [web_search_tool], + "custom_llm_provider": "bedrock", + } + + monkeypatch.setattr( + proxy_server, + "llm_router", + MagicMock(search_tools=[{"search_tool_name": "other-search", "litellm_params": {"search_provider": "tavily"}}]), + ) + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + with pytest.raises(ValueError, match="Configured search tool 'parallel-search' was not found"): + await async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type.value) + + assert kwargs["tools"] == [web_search_tool] + mock_asearch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_deployment_hook_skips_explicit_tool_validation_for_non_search_responses(monkeypatch): + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"], search_tool_name="parallel-search") + monkeypatch.setattr(proxy_server, "llm_router", MagicMock(search_tools=[])) + + result = await logger.async_pre_call_deployment_hook( + kwargs={ + "tools": [{"type": "function", "name": "calculator"}], + "custom_llm_provider": "bedrock", + }, + call_type=CallTypes.aresponses, + ) + + assert result is None + + @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_provider_from_top_level_kwargs(): """Test that async_pre_call_deployment_hook finds custom_llm_provider at top-level kwargs.