Merge pull request #40399 from adssoccer1/feat/websearch-multi-query-schema

feat(websearch): let the model emit objective + multi-query search shapes

litellm/integrations/websearch_interception/handler.py

@@ -44,6 +44,7 @@ from litellm.types.integrations.custom_logger import (
 from litellm.types.integrations.websearch_interception import (
     AnthropicSearchQuery,
     AnthropicServerToolUseBlock,
+    RichWebSearchInput,
     SearchFailed,
     SearchOutcome,
     WebSearchInterceptionConfig,
@@ -1144,7 +1145,9 @@ class WebSearchInterceptionLogger(CustomLogger):
         """Execute litellm.asearch() and build a Responses API rerun patch."""
         search_tasks: Final = [
             (
-                self._execute_search(tool_call["input"]["query"], kwargs=kwargs)
+                self._execute_search(
+                    tool_call["input"]["query"], kwargs=kwargs, rich=self._rich_search_input(tool_call["input"])
+                )
                 if isinstance(tool_call.get("input"), dict) and tool_call["input"].get("query")
                 else self._create_empty_search_result()
             )
@@ -1362,7 +1365,9 @@ class WebSearchInterceptionLogger(CustomLogger):
             query = tool_call["input"].get("query")
             if query:
                 verbose_logger.debug("WebSearchInterception: Queuing search for query='%s'", query)
-                search_tasks.append(self._execute_search(query, kwargs=kwargs))
+                search_tasks.append(
+                    self._execute_search(query, kwargs=kwargs, rich=self._rich_search_input(tool_call["input"]))
+                )
             else:
                 verbose_logger.debug("WebSearchInterception: Tool call %s has no query", tool_call["id"])
                 # Add empty result for tools without query
@@ -1431,8 +1436,53 @@ class WebSearchInterceptionLogger(CustomLogger):
             return WebSearchTransformation.search_outcome(e)
         return WebSearchTransformation.search_outcome(result)
 
+    @staticmethod
+    def _rich_search_input(tool_input: object) -> RichWebSearchInput | None:
+        """
+        Extract the optional objective/search_queries pair from a tool input.
+
+        Returns None when the input carries neither, so callers can pass the
+        result straight through as ``_execute_search``'s ``rich`` argument.
+        """
+        if not isinstance(tool_input, Mapping):
+            return None
+        objective = tool_input.get("objective")
+        valid_objective = objective if isinstance(objective, str) and objective.strip() else None
+        raw_queries = tool_input.get("search_queries")
+        valid_queries: list[str] | None = None  # mutable-ok: matches litellm.asearch's list[str] query parameter
+        if isinstance(raw_queries, Sequence) and not isinstance(raw_queries, str):
+            queries = [q for q in raw_queries if isinstance(q, str) and q.strip()]
+            if queries:
+                # Providers cap multi-query requests (Parallel drops queries
+                # past the fifth); trim here so nothing is silently ignored.
+                valid_queries = queries[:5]
+        if valid_objective is not None and valid_queries is not None:
+            return {"objective": valid_objective, "search_queries": valid_queries}
+        if valid_objective is not None:
+            return {"objective": valid_objective}
+        if valid_queries is not None:
+            return {"search_queries": valid_queries}
+        return None
+
+    @staticmethod
+    def _provider_supports_rich_search(search_provider: str | None) -> bool:
+        """Whether the provider's search config accepts objective + multi-query input."""
+        if not search_provider:
+            return False
+        try:
+            from litellm.utils import ProviderConfigManager
+        except ImportError:
+            return False
+        # SearchProviders is a str enum, so an unknown provider string simply
+        # misses the config map and returns None rather than raising.
+        config = ProviderConfigManager.get_provider_search_config(search_provider)  # pyright: ignore[reportArgumentType] -- SearchProviders is a str enum, so the router's provider string hashes to the matching member; unknown strings miss the map and yield None
+        return config is not None and config.supports_rich_search_input()
+
     async def _execute_search(
-        self, query: str, kwargs: Mapping[str, object] | None = None
+        self,
+        query: str,
+        kwargs: Mapping[str, object] | None = None,
+        rich: RichWebSearchInput | None = None,
     ) -> tuple[str, SearchResponse | None]:
         """
         Execute a single web search using router's search tools.
@@ -1490,13 +1540,24 @@ class WebSearchInterceptionLogger(CustomLogger):
                 for key, value in search_litellm_params.items()
                 if key != "search_provider" and value is not None
             }
+            # Forward the model's richer shape (objective + keyword queries)
+            # only to providers whose search API takes it natively; everyone
+            # else keeps the single query string the model also provided.
+            query_arg: str | list[str] = query  # mutable-ok: litellm.asearch declares query as str | list[str]
+            if rich and self._provider_supports_rich_search(search_provider):
+                rich_queries = rich.get("search_queries")
+                if rich_queries:
+                    query_arg = rich_queries
+                rich_objective = rich.get("objective")
+                if rich_objective and "objective" not in search_kwargs:
+                    search_kwargs["objective"] = rich_objective
             result: Final = (
                 await litellm.asearch(
-                    query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
+                    query=query_arg, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
                 )
                 if search_metadata is None
                 else await litellm.asearch(
-                    query=query,
+                    query=query_arg,
                     search_provider=search_provider,
                     litellm_metadata=search_metadata,
                     **_NO_ASEARCH_NAMED,
@@ -1701,18 +1762,21 @@ class WebSearchInterceptionLogger(CustomLogger):
         for tool_call in tool_calls:
             # Handle both Anthropic-style input and OpenAI-style function.arguments
             query = None
+            tool_args: dict | None = None  # mutable-ok: the tool call's own arguments dict
             if "input" in tool_call and isinstance(tool_call["input"], dict):
-                query = tool_call["input"].get("query")
+                tool_args = tool_call["input"]
+                query = tool_args.get("query")
             elif "function" in tool_call:
                 func = tool_call["function"]
                 if isinstance(func, dict):
                     args = func.get("arguments", {})
                     if isinstance(args, dict):
+                        tool_args = args
                         query = args.get("query")
 
             if query:
                 verbose_logger.debug("WebSearchInterception: Queuing search for query='%s'", query)
-                search_tasks.append(self._execute_search(query, kwargs=kwargs))
+                search_tasks.append(self._execute_search(query, kwargs=kwargs, rich=self._rich_search_input(tool_args)))
             else:
                 verbose_logger.debug("WebSearchInterception: Tool call %s has no query", tool_call.get("id"))
                 # Add empty result for tools without query

litellm/integrations/websearch_interception/tools.py

@@ -11,6 +11,50 @@ from typing import Any, Final
 
 from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME
 
+_WEB_SEARCH_TOOL_DESCRIPTION: Final = (
+    "Search the web for information. Use this when you need current "
+    "information or answers to questions that require up-to-date data."
+)
+
+
+def _web_search_input_schema() -> dict[str, object]:  # mutable-ok: plain-dict tool shape, as the get_* builders
+    """
+    JSON schema for the web search tool's input, shared by every tool format.
+
+    ``query`` stays required so providers and callers that only understand a
+    single query string keep working unchanged. ``objective`` and
+    ``search_queries`` are optional richer inputs; they are forwarded only to
+    search providers that support them (see
+    ``BaseSearchConfig.supports_rich_search_input``).
+    """
+    return {
+        "type": "object",
+        "properties": {
+            "query": {
+                "type": "string",
+                "description": "The search query to execute",
+            },
+            "objective": {
+                "type": "string",
+                "description": (
+                    "Natural-language description of the goal behind the "
+                    "search, including any source or freshness requirements."
+                ),
+            },
+            "search_queries": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": (
+                    "Two to five short keyword queries (3-6 words each) "
+                    "covering different angles of the objective, e.g. varying "
+                    "names, synonyms, or phrasings. Provide together with "
+                    "objective for the best results."
+                ),
+            },
+        },
+        "required": ["query"],
+    }
+
 
 def get_litellm_web_search_tool() -> dict[str, object]:
     """
@@ -33,20 +77,8 @@ def get_litellm_web_search_tool() -> dict[str, object]:
     """
     return {
         "name": LITELLM_WEB_SEARCH_TOOL_NAME,
-        "description": (
-            "Search the web for information. Use this when you need current "
-            "information or answers to questions that require up-to-date data."
-        ),
-        "input_schema": {
-            "type": "object",
-            "properties": {
-                "query": {
-                    "type": "string",
-                    "description": "The search query to execute",
-                }
-            },
-            "required": ["query"],
-        },
+        "description": _WEB_SEARCH_TOOL_DESCRIPTION,
+        "input_schema": _web_search_input_schema(),
     }
 
 
@@ -65,20 +97,8 @@ def get_litellm_web_search_tool_openai() -> dict[str, object]:
         "type": "function",
         "function": {
             "name": LITELLM_WEB_SEARCH_TOOL_NAME,
-            "description": (
-                "Search the web for information. Use this when you need current "
-                "information or answers to questions that require up-to-date data."
-            ),
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "query": {
-                        "type": "string",
-                        "description": "The search query to execute",
-                    }
-                },
-                "required": ["query"],
-            },
+            "description": _WEB_SEARCH_TOOL_DESCRIPTION,
+            "parameters": _web_search_input_schema(),
         },
     }
 
@@ -98,20 +118,8 @@ def get_litellm_web_search_tool_responses() -> dict[str, object]:
     return {
         "type": "function",
         "name": LITELLM_WEB_SEARCH_TOOL_NAME,
-        "description": (
-            "Search the web for information. Use this when you need current "
-            "information or answers to questions that require up-to-date data."
-        ),
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "query": {
-                    "type": "string",
-                    "description": "The search query to execute",
-                }
-            },
-            "required": ["query"],
-        },
+        "description": _WEB_SEARCH_TOOL_DESCRIPTION,
+        "parameters": _web_search_input_schema(),
     }
 
 

litellm/llms/base_llm/search/transformation.py

@@ -95,6 +95,18 @@ class BaseSearchConfig:
         """
         return "Unknown Search Provider"
 
+    def supports_rich_search_input(self) -> bool:
+        """
+        Whether this provider's search API accepts a natural-language
+        objective plus multiple keyword queries in one request.
+
+        Integrations that collect the richer shape (e.g. websearch
+        interception) forward ``query`` as a list plus an ``objective``
+        optional param to providers that return True; every other provider
+        keeps receiving the single query string.
+        """
+        return False
+
     def get_http_method(self) -> Literal["GET", "POST"]:
         """
         Get HTTP method for search requests.

litellm/llms/parallel_ai/search/transformation.py

@@ -90,6 +90,11 @@ class ParallelAISearchConfig(BaseSearchConfig):
     def ui_friendly_name() -> str:
         return "Parallel AI"
 
+    def supports_rich_search_input(self) -> bool:
+        # The v1 search API takes `objective` + multiple `search_queries`
+        # natively; sending both is the documented best practice.
+        return True
+
     def validate_environment(
         self,
         headers: dict,

litellm/types/integrations/websearch_interception.py

@@ -31,6 +31,22 @@ class AnthropicServerToolUseBlock(BaseModel):
     input: AnthropicSearchQuery
 
 
+class RichWebSearchInput(TypedDict, total=False):
+    """
+    Optional richer search shape a model may emit alongside ``query``.
+
+    Collected from the intercepted tool call and forwarded only to search
+    providers whose config reports ``supports_rich_search_input()``; every
+    other provider keeps receiving the single ``query`` string.
+    """
+
+    objective: ReadOnly[str]
+    """Natural-language description of the goal behind the search."""
+
+    search_queries: ReadOnly[list[str]]  # mutable-ok: forwarded verbatim as litellm.asearch's list[str] query argument
+    """Two to five short keyword queries covering different angles."""
+
+
 WebSearchToolResultErrorCode: TypeAlias = Literal[
     "invalid_tool_input",
     "unavailable",

tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py

@@ -595,7 +595,7 @@ class TestFailedSearchEndsTheTurn:
     async def test_mixed_iteration_keeps_the_follow_up_call(self, monkeypatch):
         monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate)
 
-        async def search(query, kwargs=None):
+        async def search(query, kwargs=None, rich=None):
             if query == "fails":
                 raise RateLimitError("slow down", llm_provider="tavily", model="tavily")
             found = SearchResult(title="Result", url="https://example.com", snippet="A result.", date=None)

tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py

@@ -419,7 +419,7 @@ class TestFailedSearchOutcome:
             {"id": "toolu_two", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "works"}},
         ]
 
-        async def search(query, kwargs=None):
+        async def search(query, kwargs=None, rich=None):
             if query == "fails":
                 raise RateLimitError("slow down", llm_provider="tavily", model="tavily")
             return ("Title: x", _make_search_response())

tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py

@@ -0,0 +1,246 @@
+"""
+Unit tests for the rich web-search input shape (objective + search_queries).
+
+The intercepted web search tool exposes optional `objective` and
+`search_queries` fields alongside the required single `query` string. The
+handler forwards the richer shape only to search providers whose config
+reports supports_rich_search_input(); every other provider keeps receiving
+the single query string the model also provided.
+"""
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.integrations.websearch_interception.handler import (
+    WebSearchInterceptionLogger,
+)
+from litellm.integrations.websearch_interception.tools import (
+    get_litellm_web_search_tool,
+    get_litellm_web_search_tool_openai,
+    get_litellm_web_search_tool_responses,
+)
+from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse
+from litellm.llms.parallel_ai.search.transformation import ParallelAISearchConfig
+
+RICH_INPUT = {
+    "query": "stripe node sdk v14 authentication",
+    "objective": "Find the current authentication flow for the Stripe Node SDK v14",
+    "search_queries": ["stripe node sdk v14 auth", "stripe api key rotation node"],
+}
+
+
+def _search_response() -> SearchResponse:
+    return SearchResponse(object="search", results=[])
+
+
+def _mock_router(search_provider: str) -> MagicMock:
+    """Router stub exposing one configured search tool."""
+    router = MagicMock()
+    router.search_tools = [
+        {
+            "search_tool_name": "test-search",
+            "litellm_params": {
+                "search_provider": search_provider,
+                "api_key": "sk-test",
+            },
+        }
+    ]
+    return router
+
+
+class TestToolSchema:
+    def test_all_formats_expose_rich_fields_and_keep_query_required(self):
+        anthropic_schema = get_litellm_web_search_tool()["input_schema"]
+        openai_schema = get_litellm_web_search_tool_openai()["function"]["parameters"]
+        responses_schema = get_litellm_web_search_tool_responses()["parameters"]
+
+        for schema in (anthropic_schema, openai_schema, responses_schema):
+            assert schema["required"] == ["query"]
+            assert "objective" in schema["properties"]
+            assert "search_queries" in schema["properties"]
+            assert schema["properties"]["search_queries"]["type"] == "array"
+
+
+class TestRichInputExtraction:
+    def test_extracts_objective_and_queries(self):
+        rich = WebSearchInterceptionLogger._rich_search_input(RICH_INPUT)
+        assert rich == {
+            "objective": RICH_INPUT["objective"],
+            "search_queries": RICH_INPUT["search_queries"],
+        }
+
+    def test_returns_none_when_only_query_present(self):
+        assert WebSearchInterceptionLogger._rich_search_input({"query": "plain"}) is None
+
+    def test_returns_none_for_non_mapping_input(self):
+        assert WebSearchInterceptionLogger._rich_search_input(None) is None
+        assert WebSearchInterceptionLogger._rich_search_input("query") is None
+
+    def test_drops_invalid_queries_and_caps_at_five(self):
+        rich = WebSearchInterceptionLogger._rich_search_input(
+            {
+                "query": "q",
+                "search_queries": ["a", "", 3, "b", "c", "d", "e", "f"],
+            }
+        )
+        assert rich == {"search_queries": ["a", "b", "c", "d", "e"]}
+
+    def test_ignores_string_valued_search_queries(self):
+        # A string is a Sequence; it must not be treated as a list of queries.
+        assert WebSearchInterceptionLogger._rich_search_input({"query": "q", "search_queries": "not a list"}) is None
+
+
+class TestProviderSupport:
+    def test_parallel_ai_supports_rich_input(self):
+        assert ParallelAISearchConfig().supports_rich_search_input() is True
+
+    def test_base_config_defaults_to_unsupported(self):
+        assert BaseSearchConfig().supports_rich_search_input() is False
+
+    def test_unknown_provider_is_unsupported(self):
+        assert WebSearchInterceptionLogger._provider_supports_rich_search(None) is False
+        assert WebSearchInterceptionLogger._provider_supports_rich_search("not_a_provider") is False
+
+
+class TestExecuteSearchShape:
+    @pytest.mark.asyncio
+    async def test_rich_shape_reaches_supporting_provider(self, monkeypatch):
+        """Parallel AI receives the query list plus objective."""
+        import litellm
+        from litellm.proxy import proxy_server
+
+        logger = WebSearchInterceptionLogger()
+        mock_asearch = AsyncMock(return_value=_search_response())
+        monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai"))
+        monkeypatch.setattr(litellm, "asearch", mock_asearch)
+
+        rich = WebSearchInterceptionLogger._rich_search_input(RICH_INPUT)
+        await logger._execute_search(RICH_INPUT["query"], rich=rich)
+
+        call_kwargs = mock_asearch.await_args.kwargs
+        assert call_kwargs["query"] == RICH_INPUT["search_queries"]
+        assert call_kwargs["objective"] == RICH_INPUT["objective"]
+        assert call_kwargs["search_provider"] == "parallel_ai"
+
+    @pytest.mark.asyncio
+    async def test_string_only_provider_keeps_single_query(self, monkeypatch):
+        """A provider without rich support receives the plain query string."""
+        import litellm
+        from litellm.proxy import proxy_server
+
+        logger = WebSearchInterceptionLogger()
+        mock_asearch = AsyncMock(return_value=_search_response())
+        monkeypatch.setattr(proxy_server, "llm_router", _mock_router("perplexity"))
+        monkeypatch.setattr(litellm, "asearch", mock_asearch)
+
+        rich = WebSearchInterceptionLogger._rich_search_input(RICH_INPUT)
+        await logger._execute_search(RICH_INPUT["query"], rich=rich)
+
+        call_kwargs = mock_asearch.await_args.kwargs
+        assert call_kwargs["query"] == RICH_INPUT["query"]
+        assert "objective" not in call_kwargs
+
+    @pytest.mark.asyncio
+    async def test_single_string_callers_unchanged(self, monkeypatch):
+        """No rich input: behavior is identical to before for any provider."""
+        import litellm
+        from litellm.proxy import proxy_server
+
+        logger = WebSearchInterceptionLogger()
+        mock_asearch = AsyncMock(return_value=_search_response())
+        monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai"))
+        monkeypatch.setattr(litellm, "asearch", mock_asearch)
+
+        await logger._execute_search("plain query")
+
+        call_kwargs = mock_asearch.await_args.kwargs
+        assert call_kwargs["query"] == "plain query"
+        assert "objective" not in call_kwargs
+
+    @pytest.mark.asyncio
+    async def test_configured_objective_not_overwritten(self, monkeypatch):
+        """An objective set on the search tool's litellm_params wins over the model's."""
+        import litellm
+        from litellm.proxy import proxy_server
+
+        logger = WebSearchInterceptionLogger()
+        router = _mock_router("parallel_ai")
+        router.search_tools[0]["litellm_params"]["objective"] = "configured objective"
+        mock_asearch = AsyncMock(return_value=_search_response())
+        monkeypatch.setattr(proxy_server, "llm_router", router)
+        monkeypatch.setattr(litellm, "asearch", mock_asearch)
+
+        rich = WebSearchInterceptionLogger._rich_search_input(RICH_INPUT)
+        await logger._execute_search(RICH_INPUT["query"], rich=rich)
+
+        call_kwargs = mock_asearch.await_args.kwargs
+        assert call_kwargs["objective"] == "configured objective"
+
+
+class TestCallSiteWiring:
+    """Drive the patch builders end to end so regressions in the tool-call ->
+    _rich_search_input wiring are caught, not just _execute_search itself."""
+
+    @pytest.mark.asyncio
+    async def test_anthropic_tool_call_forwards_rich_shape(self, monkeypatch):
+        import litellm
+        from litellm.proxy import proxy_server
+
+        logger = WebSearchInterceptionLogger()
+        mock_asearch = AsyncMock(return_value=_search_response())
+        monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai"))
+        monkeypatch.setattr(litellm, "asearch", mock_asearch)
+
+        tool_calls = [{"id": "toolu_1", "name": "litellm_web_search", "input": dict(RICH_INPUT)}]
+        await logger._build_anthropic_request_patch(
+            model="claude",
+            messages=[{"role": "user", "content": "hi"}],
+            tool_calls=tool_calls,
+            thinking_blocks=[],
+            anthropic_messages_optional_request_params={},
+            logging_obj=None,
+            kwargs={},
+        )
+
+        call_kwargs = mock_asearch.await_args.kwargs
+        assert call_kwargs["query"] == RICH_INPUT["search_queries"]
+        assert call_kwargs["objective"] == RICH_INPUT["objective"]
+
+    @pytest.mark.asyncio
+    async def test_chat_completion_tool_call_forwards_rich_shape(self, monkeypatch):
+        import json
+
+        import litellm
+        from litellm.proxy import proxy_server
+
+        logger = WebSearchInterceptionLogger()
+        mock_asearch = AsyncMock(return_value=_search_response())
+        monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai"))
+        monkeypatch.setattr(litellm, "asearch", mock_asearch)
+
+        # The normalized shape transform_request produces for OpenAI responses:
+        # function.arguments (raw) plus top-level name/input (parsed).
+        tool_calls = [
+            {
+                "id": "call_1",
+                "type": "function",
+                "name": "litellm_web_search",
+                "function": {
+                    "name": "litellm_web_search",
+                    "arguments": json.dumps(RICH_INPUT),
+                },
+                "input": dict(RICH_INPUT),
+            }
+        ]
+        await logger._build_chat_completion_request_patch(
+            model="claude",
+            messages=[{"role": "user", "content": "hi"}],
+            tool_calls=tool_calls,
+            optional_params={},
+            kwargs={},
+        )
+
+        call_kwargs = mock_asearch.await_args.kwargs
+        assert call_kwargs["query"] == RICH_INPUT["search_queries"]
+        assert call_kwargs["objective"] == RICH_INPUT["objective"]