From 6037ac76688561b56d459596f16d490e1971c4ab Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 2 Jun 2026 10:52:33 -0700 Subject: [PATCH] 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. --- .../search_api/transformation.py | 102 +++++++++--------- litellm/types/vector_stores.py | 60 +++++++++++ ...x_ai_search_vector_store_transformation.py | 65 +++++++++-- 3 files changed, 172 insertions(+), 55 deletions(-) diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 8f459d05d03..2e55aed3c0b 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/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 diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index ce247fc900f..6adfbf4fd35 100644 --- a/litellm/types/vector_stores.py +++ b/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""" diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py index 1a22e49778a..8a83ae3a894 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py +++ b/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})