From 9e57dd40fa8771917e1bddd1f15211e0f3109e36 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Wed, 13 May 2026 17:14:34 -0700 Subject: [PATCH] fix(vector-stores): support engines URL for Vertex AI Search Adds optional vertex_engine_id field to vertex_ai/search_api so users can route through a Discovery Engine search app instead of the data store directly. Required for website, healthcare, and connector-based data stores that return FAILED_PRECONDITION on the existing dataStores URL. Existing data-store-direct callers are unaffected. Resolves LIT-3036 --- .../search_api/transformation.py | 42 ++++++--- ...x_ai_search_vector_store_transformation.py | 88 +++++++++++++++++++ .../VectorStoreForm.tsx | 17 +++- .../src/components/vector_store_providers.tsx | 9 ++ 4 files changed, 140 insertions(+), 16 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 61fb848b40a..14a0a406dff 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -80,31 +80,47 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): litellm_params: dict, ) -> str: """ - Get the Base endpoint for Vertex AI Search API + Get the Base endpoint for Vertex AI Search API. + + Branches on whether a `vertex_engine_id` is configured: + - Engine ID present: route through the search app (engine) — required for website, + healthcare, and connector-based data stores. Note the serving config name differs + (`default_serving_config` vs `default_config` for direct data store search). + - Engine ID absent: query the data store directly via `vector_store_id`. """ + if api_base: + return api_base.rstrip("/") + vertex_location = self.get_vertex_ai_location(litellm_params) vertex_project = self.get_vertex_ai_project(litellm_params) collection_id = ( litellm_params.get("vertex_collection_id") or "default_collection" ) - datastore_id = litellm_params.get("vector_store_id") - if not datastore_id: - raise ValueError("vector_store_id is required") - if api_base: - return api_base.rstrip("/") encoded_collection_id = encode_url_path_segment( collection_id, field_name="vertex_collection_id" ) + base = ( + f"https://discoveryengine.googleapis.com/v1/" + f"projects/{vertex_project}/locations/{vertex_location}/" + f"collections/{encoded_collection_id}" + ) + + engine_id = litellm_params.get("vertex_engine_id") + if engine_id: + encoded_engine_id = encode_url_path_segment( + engine_id, field_name="vertex_engine_id" + ) + return f"{base}/engines/{encoded_engine_id}/servingConfigs/default_serving_config" + + datastore_id = litellm_params.get("vector_store_id") + if not datastore_id: + raise ValueError( + "vector_store_id is required when vertex_engine_id is not set" + ) encoded_datastore_id = encode_url_path_segment( datastore_id, field_name="vector_store_id" ) - - # Vertex AI Search API endpoint for search - return ( - f"https://discoveryengine.googleapis.com/v1/" - f"projects/{vertex_project}/locations/{vertex_location}/" - f"collections/{encoded_collection_id}/dataStores/{encoded_datastore_id}/servingConfigs/default_config" - ) + return f"{base}/dataStores/{encoded_datastore_id}/servingConfigs/default_config" def transform_search_vector_store_request( self, 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 b6329f33ae4..5ca71dc08c3 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 @@ -38,3 +38,91 @@ def test_should_reject_dot_segment_vertex_search_vector_store_id(): "vector_store_id": "..", }, ) + + +def test_should_use_engines_url_when_engine_id_provided(): + config = VertexSearchAPIVectorStoreConfig() + + url = config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + "vertex_engine_id": "test-engine_1234", + }, + ) + + assert url == ( + "https://discoveryengine.googleapis.com/v1/" + "projects/test-project/locations/global/" + "collections/default_collection/engines/test-engine_1234/servingConfigs/default_serving_config" + ) + + +def test_engine_id_takes_precedence_over_vector_store_id(): + config = VertexSearchAPIVectorStoreConfig() + + url = config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + "vertex_engine_id": "test-engine_1234", + "vector_store_id": "ignored-when-engine-set", + }, + ) + + assert "/engines/test-engine_1234/" in url + assert "/dataStores/" not in url + assert url.endswith("/servingConfigs/default_serving_config") + + +def test_should_encode_vertex_engine_id_in_complete_url(): + config = VertexSearchAPIVectorStoreConfig() + + url = config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + "vertex_engine_id": "../../engines/other?x=1#frag", + }, + ) + + assert url == ( + "https://discoveryengine.googleapis.com/v1/" + "projects/test-project/locations/global/" + "collections/default_collection/engines/..%2F..%2Fengines%2Fother%3Fx%3D1%23frag/servingConfigs/default_serving_config" + ) + + +def test_should_reject_dot_segment_vertex_engine_id(): + config = VertexSearchAPIVectorStoreConfig() + + with pytest.raises( + ValueError, match="vertex_engine_id cannot be a dot path segment" + ): + config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + "vertex_engine_id": "..", + }, + ) + + +def test_should_raise_when_neither_engine_id_nor_vector_store_id_provided(): + config = VertexSearchAPIVectorStoreConfig() + + with pytest.raises( + ValueError, + match="vector_store_id is required when vertex_engine_id is not set", + ): + config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + }, + ) diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx index 673061b275d..afd5ac421f3 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx @@ -32,6 +32,7 @@ const VectorStoreForm: React.FC = ({ const [metadataJson, setMetadataJson] = useState("{}"); const [selectedProvider, setSelectedProvider] = useState("bedrock"); const [modelInfo, setModelInfo] = useState([]); + const vertexEngineId = Form.useWatch("vertex_engine_id", form); useEffect(() => { if (!accessToken) return; @@ -230,8 +231,16 @@ const VectorStoreForm: React.FC = ({
  • Pick a supported location: global, us, or eu
    • -
  • Copy the data store ID from the Vertex AI Search console
    • -
  • Enter the data store ID in the Vector Store ID field below
    • +
  • + For most data store types (Cloud Storage, BigQuery, Media): copy the data store ID and enter it in + the Vector Store ID field below. +
    • +
  • + For website, healthcare, and connector-based sources (Drive, Gmail, Slack, Jira, etc.): create a + search app on top of the data store, then copy the Engine ID and enter it in the + Engine ID field. The Vector Store ID is still required as the LiteLLM-side name for this record, + but it isn't used in the GCP URL when Engine ID is set. +
    • } @@ -258,7 +267,9 @@ const VectorStoreForm: React.FC = ({ selectedProvider === "vertex_rag_engine" ? "6917529027641081856 (Get corpus ID from Vertex AI console)" : selectedProvider === "vertex_ai/search_api" - ? "my-datastore_1234567890 (Get data store ID from Vertex AI Search console)" + ? vertexEngineId + ? "Any identifier you'll use to reference this in LiteLLM" + : "my-datastore_1234567890 (Get data store ID from Vertex AI Search console)" : "Enter vector store ID from your provider" } /> diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index c5bdf1d5c2a..fb1f32a00e0 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -97,6 +97,15 @@ export const vectorStoreProviderFields: Record required: false, type: "text", }, + { + name: "vertex_engine_id", + label: "Engine ID (optional)", + tooltip: + "Search app (engine) ID. Required for website, healthcare, and connector-based data stores (Workspace, Slack, Jira, etc.) because these sources route search through an engine. Leave blank to query the data store directly.", + placeholder: "e.g. my-search-app_1234567890", + required: false, + type: "text", + }, ], openai: [ {