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.

litellm/llms/vertex_ai/vector_stores/search_api/transformation.py

@@ -16,6 +16,8 @@ from litellm.types.vector_stores import (
     VectorStoreSearchOptionalRequestParams,
     VectorStoreSearchResponse,
     VectorStoreSearchResult,
+    VertexSearchDataStoreExtraBody,
+    VertexSearchEngineExtraBody,
 )
 
 if TYPE_CHECKING:
@@ -26,49 +28,28 @@ 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(
     {
-        "dataStoreSpecs",
         "branch",
         "servingConfig",
         "entity",
     }
 )
 
-VERTEX_SEARCH_SUPPORTED_EXTRA_BODY_FIELDS = frozenset(
-    {
-        "query",
-        "pageSize",
-        "pageToken",
-        "offset",
-        "oneBoxPageSize",
-        "numResultsPerDataStore",
-        "pageCategories",
-        "imageQuery",
-        "filter",
-        "canonicalFilter",
-        "orderBy",
-        "userInfo",
-        "languageCode",
-        "facetSpecs",
-        "boostSpec",
-        "params",
-        "queryExpansionSpec",
-        "spellCorrectionSpec",
-        "userPseudoId",
-        "contentSearchSpec",
-        "rankingExpression",
-        "rankingExpressionBackend",
-        "safeSearch",
-        "userLabels",
-        "naturalLanguageQueryUnderstandingSpec",
-        "searchAsYouTypeSpec",
-        "displaySpec",
-        "crowdingSpecs",
-        "relevanceThreshold",
-        "relevanceScoreSpec",
-        "customRankingParams",
-    }
+# 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__
 )
 
 
@@ -83,20 +64,34 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
         super().__init__()
 
     @staticmethod
-    def get_supported_extra_body_fields() -> frozenset:
-        """Native SearchRequest fields callers may forward via ``extra_body``."""
-        return VERTEX_SEARCH_SUPPORTED_EXTRA_BODY_FIELDS
+    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]) -> Dict[str, Any]:
+    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.
+        Validate ``extra_body`` against the supported-field allowlist for the
+        active serving config (engine/app vs data store).
 
         Raises ``ValueError`` if the caller includes a target-selecting field
-        (e.g. ``dataStoreSpecs``) or any field not on the allowlist, so the
-        request fails loudly instead of silently searching the wrong store.
+        (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()
+        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
         }
@@ -111,9 +106,10 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
 
         unsupported = set(filtered) - supported
         if unsupported:
+            mode = "engine/app" if is_engine else "data store"
             raise ValueError(
-                f"Unsupported Vertex AI Search extra_body fields {sorted(unsupported)}. "
-                f"Supported fields: {sorted(supported)}."
+                f"Unsupported Vertex AI Search extra_body fields {sorted(unsupported)} "
+                f"for {mode} mode. Supported fields: {sorted(supported)}."
             )
 
         return filtered
@@ -223,21 +219,29 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
         callers can send native Discovery Engine tuning fields such as filter,
         boostSpec, or contentSearchSpec.
 
-        Target-selecting fields (e.g. dataStoreSpecs, branch) are rejected: the
-        data store is scoped by the URL path (vector_store_id / vertex_engine_id)
-        and must not be overridable per request.
+        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.
         """
         if isinstance(query, list):
             query = " ".join(query)
 
         url = f"{api_base}:search"
 
+        is_engine = bool(litellm_params.get("vertex_engine_id"))
+
         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))
+            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

litellm/types/vector_stores.py

@@ -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"""

tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py

@@ -135,14 +135,36 @@ _ENGINE_BASE = (
     "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)
@@ -166,25 +188,48 @@ def test_search_request_maps_max_num_results_to_pagesize():
     assert body["pageSize"] == 25
 
 
-def test_search_request_rejects_datastorespecs_in_extra_body():
+def test_engine_search_request_forwards_datastorespecs():
     specs = [
         {
             "dataStore": "projects/p/locations/global/collections/default_collection/dataStores/ds-beta"
         }
     ]
 
-    with pytest.raises(ValueError, match="target-selecting"):
-        _search_request(extra_body={"dataStoreSpecs": specs})
+    _, body = _search_request(extra_body={"dataStoreSpecs": specs})
+
+    assert body["dataStoreSpecs"] == specs
 
 
-@pytest.mark.parametrize(
-    "field", ["dataStoreSpecs", "branch", "servingConfig", "entity"]
-)
+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(ValueError, 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(ValueError, 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(ValueError, 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(ValueError, match="target-selecting"):
+        _datastore_search_request(extra_body={field: "x"})
+
+
 def test_search_request_rejects_unsupported_extra_body_field():
     with pytest.raises(ValueError, match="Unsupported Vertex AI Search extra_body"):
         _search_request(extra_body={"notARealField": True})
@@ -203,6 +248,14 @@ def test_search_request_forwards_supported_extra_body_fields():
     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})