diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index aefd4fa3b47..81310b9ddc3 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -416,25 +416,15 @@ class WebSearchInterceptionLogger(CustomLogger): if not tools: return None - 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 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) + 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 @@ -1641,7 +1631,9 @@ class WebSearchInterceptionLogger(CustomLogger): return None def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None": - search_tools: Final = list(getattr(llm_router, "search_tools", []) or []) + if llm_router is None or not hasattr(llm_router, "search_tools"): + return None + 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( @@ -1651,26 +1643,20 @@ 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 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" + 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, ) - + return matching_tools[0] verbose_logger.debug( - "WebSearchInterception: Found search tool '%s' from %s with provider '%s'", + "WebSearchInterception: Search tool '%s' not found in %s, falling back to first available or perplexity", 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 ec4bc1f49eb..f39f41a6d12 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 CallTypes, LlmProviders +from litellm.types.utils import LlmProviders def test_initialize_from_proxy_config(): @@ -230,124 +230,6 @@ 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. @@ -515,72 +397,6 @@ 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.