feat(vector-stores): forward per-request params to Vertex AI Search (#29459)

* feat(vector-stores): forward per-request params to Vertex AI Search The vertex_ai/search_api search transform hardcoded the request body to query plus pageSize 10, dropping max_num_results and extra_body. Map max_num_results to pageSize and merge extra_body through with precedence, so callers can send native Discovery Engine fields such as dataStoreSpecs. Resolves LIT-3506 * fix(vector-stores): log effective query when extra_body overrides it When a caller passes a query inside extra_body, the outbound Vertex Search request used that value but model_call_details recorded the original, so the echoed search_query was stale. Log the effective query from the request body. * fix(vector-stores): allowlist Vertex AI Search extra_body fields Raw-merging extra_body let callers set dataStoreSpecs/branch to search a different Discovery Engine data store with the proxy's Vertex credentials, bypassing the vector_store_id path authorization. Reject target-selecting fields and forward only allowlisted per-request tuning fields. Resolves LIT-3506 * refactor(vector-stores): split Vertex AI Search extra_body allowlists by mode Data-store and engine/app serving configs accept different SearchRequest fields, so derive two TypedDicts (VertexSearchDataStoreExtraBody and VertexSearchEngineExtraBody) in types/vector_stores.py and make _filter_extra_body mode-aware via vertex_engine_id. dataStoreSpecs and numResultsPerDataStore now pass through in engine/app mode (where an app fans out across stores) and are rejected in data-store mode. branch/servingConfig/entity remain rejected in both modes. * fix(vector-stores): raise BadRequestError (400) for invalid Vertex Search extra_body Rejecting unsupported or target-selecting extra_body fields previously raised a bare ValueError, which the vector store error path mapped to a generic APIConnectionError (HTTP 500). Raise litellm.BadRequestError so invalid per-request input surfaces as HTTP 400 with a clear message.
2026-08-09 06:23:14 +00:00 · 2026-06-02 12:51:20 -07:00
parent 6d6eda8101
commit ebbc5cc787
3 changed files with 347 additions and 10 deletions
@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 import httpx

 from litellm import get_model_info
+from litellm.exceptions import BadRequestError
 from litellm.litellm_core_utils.url_utils import encode_url_path_segment
 from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
 from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
@@ -16,6 +17,8 @@ from litellm.types.vector_stores import (
    VectorStoreSearchOptionalRequestParams,
    VectorStoreSearchResponse,
    VectorStoreSearchResult,
+    VertexSearchDataStoreExtraBody,
+    VertexSearchEngineExtraBody,
 )

 if TYPE_CHECKING:
@@ -26,6 +29,31 @@ else:
    LiteLLMLoggingObj = Any


+# Fields that select which data store / serving config to search. These are
+# always determined by the request URL path (vector_store_id / vertex_engine_id),
+# so allowing them per request could silently redirect the search to a different
+# target. Rejected in both data-store and engine/app modes.
+VERTEX_SEARCH_TARGET_SELECTING_FIELDS = frozenset(
+    {
+        "branch",
+        "servingConfig",
+        "entity",
+    }
+)
+
+# Allowlists of native Discovery Engine SearchRequest fields callers may forward
+# via extra_body, derived from the TypedDicts so the type is the source of truth.
+# Engine/app mode is a superset (adds dataStoreSpecs, numResultsPerDataStore),
+# since an app fans out across multiple member data stores.
+VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS = frozenset(
+    VertexSearchDataStoreExtraBody.__annotations__
+)
+
+VERTEX_SEARCH_ENGINE_EXTRA_BODY_FIELDS = frozenset(
+    VertexSearchEngineExtraBody.__annotations__
+)
+
+
 class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
    """
    Configuration for Vertex AI Search API Vector Store
@@ -36,6 +64,66 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
    def __init__(self):
        super().__init__()

+    @staticmethod
+    def get_supported_extra_body_fields(is_engine: bool = False) -> frozenset:
+        """
+        Native SearchRequest fields callers may forward via ``extra_body``.
+
+        The set depends on which serving config the request targets:
+        - engine/app mode (``is_engine=True``): includes multi-store fields such
+          as ``dataStoreSpecs`` and ``numResultsPerDataStore``.
+        - data-store mode: the engine-only fields are excluded.
+        """
+        if is_engine:
+            return VERTEX_SEARCH_ENGINE_EXTRA_BODY_FIELDS
+        return VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS
+
+    @classmethod
+    def _filter_extra_body(
+        cls, extra_body: Dict[str, Any], is_engine: bool = False
+    ) -> Dict[str, Any]:
+        """
+        Validate ``extra_body`` against the supported-field allowlist for the
+        active serving config (engine/app vs data store).
+
+        Raises ``BadRequestError`` (HTTP 400) if the caller includes a
+        target-selecting field (e.g. ``servingConfig``) or any field not
+        supported for the active mode, so the request fails loudly instead of
+        silently searching the wrong target. Engine-only fields
+        (``dataStoreSpecs``, ``numResultsPerDataStore``) are rejected in
+        data-store mode where they are meaningless.
+        """
+        supported = cls.get_supported_extra_body_fields(is_engine=is_engine)
+        filtered = {
+            key: value for key, value in extra_body.items() if value is not None
+        }
+
+        target_selecting = set(filtered) & VERTEX_SEARCH_TARGET_SELECTING_FIELDS
+        if target_selecting:
+            raise BadRequestError(
+                message=(
+                    "Vertex AI Search extra_body may not set target-selecting fields "
+                    f"{sorted(target_selecting)}: the data store is scoped by "
+                    "vector_store_id / vertex_engine_id and cannot be overridden per request."
+                ),
+                model="vertex_ai/search_api",
+                llm_provider="vertex_ai",
+            )
+
+        unsupported = set(filtered) - supported
+        if unsupported:
+            mode = "engine/app" if is_engine else "data store"
+            raise BadRequestError(
+                message=(
+                    f"Unsupported Vertex AI Search extra_body fields {sorted(unsupported)} "
+                    f"for {mode} mode. Supported fields: {sorted(supported)}."
+                ),
+                model="vertex_ai/search_api",
+                llm_provider="vertex_ai",
+            )
+
+        return filtered
+
    def get_auth_credentials(
        self, litellm_params: dict
    ) -> BaseVectorStoreAuthCredentials:
@@ -133,23 +221,41 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
        extra_body: Optional[Dict[str, Any]] = None,
    ) -> Tuple[str, Dict[str, Any]]:
        """
-        Transform search request for Vertex AI RAG API
+        Transform a search request for the Vertex AI Search (Discovery Engine) API.
+
+        Per-request params pass through to the engine: max_num_results maps to
+        pageSize, and extra_body fields on the supported allowlist
+        (`get_supported_extra_body_fields`) are merged in with precedence, so
+        callers can send native Discovery Engine tuning fields such as filter,
+        boostSpec, or contentSearchSpec.
+
+        The allowlist depends on the serving config: engine/app mode (when
+        `vertex_engine_id` is set) additionally accepts multi-store fields like
+        `dataStoreSpecs` and `numResultsPerDataStore`, while data-store mode
+        rejects them. Target-selecting fields (e.g. servingConfig, branch) are
+        rejected in both modes: the target is scoped by the URL path
+        (vector_store_id / vertex_engine_id) and must not be overridable per
+        request.
        """
-        # Convert query to string if it's a list
        if isinstance(query, list):
            query = " ".join(query)

-        # Vertex AI RAG API endpoint for retrieving contexts
        url = f"{api_base}:search"

-        # Construct full rag corpus path
-        # Build the request body for Vertex AI Search API
-        request_body = {"query": query, "pageSize": 10}
+        is_engine = bool(litellm_params.get("vertex_engine_id"))

-        #########################################################
-        # Update logging object with details of the request
-        #########################################################
-        litellm_logging_obj.model_call_details["query"] = query
+        request_body: Dict[str, Any] = {"query": query, "pageSize": 10}
+        max_num_results = vector_store_search_optional_params.get("max_num_results")
+        if max_num_results is not None:
+            request_body["pageSize"] = max_num_results
+        if isinstance(extra_body, dict):
+            request_body.update(
+                self._filter_extra_body(extra_body, is_engine=is_engine)
+            )
+
+        litellm_logging_obj.model_call_details["query"] = request_body.get(
+            "query", query
+        )

        return url, request_body

@@ -112,6 +112,66 @@ class VectorStoreSearchRequest(VectorStoreSearchOptionalRequestParams, total=Fal
    query: Union[str, List[str]]


+class VertexSearchDataStoreExtraBody(TypedDict, total=False):
+    """
+    Native Discovery Engine ``SearchRequest`` fields callers may forward via
+    ``extra_body`` when searching a Vertex AI Search **data store** serving
+    config (``.../dataStores/{id}/servingConfigs/default_config``).
+
+    The data store is scoped by the request URL path, so target-selecting
+    fields (``servingConfig``, ``branch``, ``entity``) are intentionally
+    omitted and rejected by the transformation layer. Engine/app-only fields
+    such as ``dataStoreSpecs`` and ``numResultsPerDataStore`` live on
+    ``VertexSearchEngineExtraBody`` instead.
+    """
+
+    query: str
+    pageSize: int
+    pageToken: str
+    offset: int
+    oneBoxPageSize: int
+    pageCategories: List[str]
+    imageQuery: Dict[str, Any]
+    filter: str
+    canonicalFilter: str
+    orderBy: str
+    userInfo: Dict[str, Any]
+    languageCode: str
+    facetSpecs: List[Dict[str, Any]]
+    boostSpec: Dict[str, Any]
+    params: Dict[str, Any]
+    queryExpansionSpec: Dict[str, Any]
+    spellCorrectionSpec: Dict[str, Any]
+    userPseudoId: str
+    contentSearchSpec: Dict[str, Any]
+    rankingExpression: str
+    rankingExpressionBackend: str
+    safeSearch: bool
+    userLabels: Dict[str, str]
+    naturalLanguageQueryUnderstandingSpec: Dict[str, Any]
+    searchAsYouTypeSpec: Dict[str, Any]
+    displaySpec: Dict[str, Any]
+    crowdingSpecs: List[Dict[str, Any]]
+    relevanceThreshold: str
+    relevanceScoreSpec: Dict[str, Any]
+    customRankingParams: Dict[str, Any]
+
+
+class VertexSearchEngineExtraBody(VertexSearchDataStoreExtraBody, total=False):
+    """
+    Native Discovery Engine ``SearchRequest`` fields callers may forward via
+    ``extra_body`` when searching a Vertex AI Search **engine/app** serving
+    config (``.../engines/{id}/servingConfigs/default_serving_config``).
+
+    Inherits every data-store field and adds fields that only make sense when
+    an app fans out across multiple member data stores, e.g. ``dataStoreSpecs``
+    (per-store scoping/filtering) and ``numResultsPerDataStore``.
+    """
+
+    dataStoreSpecs: List[Dict[str, Any]]
+    numResultsPerDataStore: int
+
+
 # Vector Store Creation Types
 class VectorStoreExpirationPolicy(TypedDict, total=False):
    """The expiration policy for a vector store"""
@@ -1,5 +1,8 @@
+from types import SimpleNamespace
+
 import pytest

+from litellm.exceptions import BadRequestError
 from litellm.llms.vertex_ai.vector_stores.search_api.transformation import (
    VertexSearchAPIVectorStoreConfig,
 )
@@ -126,3 +129,171 @@ def test_should_raise_when_neither_engine_id_nor_vector_store_id_provided():
                "vertex_location": "global",
            },
        )
+
+
+_ENGINE_BASE = (
+    "https://discoveryengine.googleapis.com/v1/projects/p/locations/global/"
+    "collections/default_collection/engines/app-2/servingConfigs/default_serving_config"
+)
+
+_DATASTORE_BASE = (
+    "https://discoveryengine.googleapis.com/v1/projects/p/locations/global/"
+    "collections/default_collection/dataStores/ds-1/servingConfigs/default_config"
+)
+
+
+def _search_request(**overrides):
+    """Engine/app-mode search request (vertex_engine_id set)."""
+    kwargs = dict(
+        vector_store_id="vs",
+        query="hello",
+        vector_store_search_optional_params={},
+        api_base=_ENGINE_BASE,
+        litellm_logging_obj=SimpleNamespace(model_call_details={}),
+        litellm_params={"vertex_engine_id": "app-2"},
+    )
+    kwargs.update(overrides)
+    return VertexSearchAPIVectorStoreConfig().transform_search_vector_store_request(
+        **kwargs
+    )
+
+
+def _datastore_search_request(**overrides):
+    """Data-store-mode search request (no vertex_engine_id)."""
+    kwargs = dict(
+        vector_store_id="ds-1",
+        query="hello",
+        vector_store_search_optional_params={},
+        api_base=_DATASTORE_BASE,
+        litellm_logging_obj=SimpleNamespace(model_call_details={}),
+        litellm_params={},
+    )
+    kwargs.update(overrides)
+    return VertexSearchAPIVectorStoreConfig().transform_search_vector_store_request(
+        **kwargs
+    )
+
+
+def test_search_request_defaults_to_query_and_pagesize_10():
+    url, body = _search_request()
+
+    assert url == _ENGINE_BASE + ":search"
+    assert body == {"query": "hello", "pageSize": 10}
+
+
+def test_search_request_maps_max_num_results_to_pagesize():
+    _, body = _search_request(
+        vector_store_search_optional_params={"max_num_results": 25}
+    )
+
+    assert body["pageSize"] == 25
+
+
+def test_engine_search_request_forwards_datastorespecs():
+    specs = [
+        {
+            "dataStore": "projects/p/locations/global/collections/default_collection/dataStores/ds-beta"
+        }
+    ]
+
+    _, body = _search_request(extra_body={"dataStoreSpecs": specs})
+
+    assert body["dataStoreSpecs"] == specs
+
+
+def test_engine_search_request_forwards_num_results_per_data_store():
+    _, body = _search_request(extra_body={"numResultsPerDataStore": 3})
+
+    assert body["numResultsPerDataStore"] == 3
+
+
+def test_datastore_search_request_rejects_datastorespecs():
+    specs = [{"dataStore": "projects/p/.../dataStores/ds-beta"}]
+
+    with pytest.raises(BadRequestError, match="data store mode"):
+        _datastore_search_request(extra_body={"dataStoreSpecs": specs})
+
+
+def test_datastore_search_request_rejects_num_results_per_data_store():
+    with pytest.raises(BadRequestError, match="data store mode"):
+        _datastore_search_request(extra_body={"numResultsPerDataStore": 3})
+
+
+@pytest.mark.parametrize("field", ["branch", "servingConfig", "entity"])
+def test_search_request_rejects_target_selecting_fields(field):
+    with pytest.raises(BadRequestError, match="target-selecting"):
+        _search_request(extra_body={field: "x"})
+
+
+@pytest.mark.parametrize("field", ["branch", "servingConfig", "entity"])
+def test_datastore_search_request_rejects_target_selecting_fields(field):
+    with pytest.raises(BadRequestError, match="target-selecting"):
+        _datastore_search_request(extra_body={field: "x"})
+
+
+def test_search_request_rejects_unsupported_extra_body_field():
+    with pytest.raises(BadRequestError, match="Unsupported Vertex AI Search extra_body"):
+        _search_request(extra_body={"notARealField": True})
+
+
+def test_rejected_extra_body_raises_http_400():
+    with pytest.raises(BadRequestError) as exc_info:
+        _search_request(extra_body={"notARealField": True})
+
+    assert exc_info.value.status_code == 400
+
+
+def test_search_request_forwards_supported_extra_body_fields():
+    _, body = _search_request(
+        extra_body={
+            "filter": 'category: ANY("docs")',
+            "boostSpec": {"conditionBoostSpecs": []},
+        }
+    )
+
+    assert body["filter"] == 'category: ANY("docs")'
+    assert body["boostSpec"] == {"conditionBoostSpecs": []}
+    assert body["query"] == "hello"
+
+
+def test_datastore_search_request_forwards_supported_extra_body_fields():
+    _, body = _datastore_search_request(
+        extra_body={"filter": 'category: ANY("docs")'}
+    )
+
+    assert body["filter"] == 'category: ANY("docs")'
+
+
+def test_search_request_ignores_none_valued_extra_body_fields():
+    _, body = _search_request(extra_body={"filter": None})
+
+    assert "filter" not in body
+
+
+def test_search_request_extra_body_takes_precedence_over_defaults():
+    _, body = _search_request(
+        vector_store_search_optional_params={"max_num_results": 5},
+        extra_body={"pageSize": 50, "filter": 'category: ANY("docs")'},
+    )
+
+    assert body["pageSize"] == 50
+    assert body["filter"] == 'category: ANY("docs")'
+
+
+def test_search_request_joins_list_query():
+    _, body = _search_request(query=["foo", "bar"])
+
+    assert body["query"] == "foo bar"
+
+
+def test_search_request_logs_effective_query_when_extra_body_overrides_query():
+    log = SimpleNamespace(model_call_details={})
+
+    _, body = _search_request(
+        query="original",
+        extra_body={"query": "from-extra-body"},
+        litellm_logging_obj=log,
+    )
+
+    assert body["query"] == "from-extra-body"
+    assert log.model_call_details["query"] == "from-extra-body"