From 675faefc670b5cd41991939fe0fc604df331111a Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Mon, 16 Jun 2025 08:34:53 -0700 Subject: [PATCH] feat: Allow data_store_specs pass into ADK VAIS built-in tool PiperOrigin-RevId: 772039465 --- src/google/adk/tools/vertex_ai_search_tool.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/google/adk/tools/vertex_ai_search_tool.py b/src/google/adk/tools/vertex_ai_search_tool.py index 5449f5090..b00cd0329 100644 --- a/src/google/adk/tools/vertex_ai_search_tool.py +++ b/src/google/adk/tools/vertex_ai_search_tool.py @@ -39,6 +39,9 @@ def __init__( self, *, data_store_id: Optional[str] = None, + data_store_specs: Optional[ + list[types.VertexAISearchDataStoreSpec] + ] = None, search_engine_id: Optional[str] = None, filter: Optional[str] = None, max_results: Optional[int] = None, @@ -49,6 +52,8 @@ def __init__( data_store_id: The Vertex AI search data store resource ID in the format of "projects/{project}/locations/{location}/collections/{collection}/dataStores/{dataStore}". + data_store_specs: Specifications that define the specific DataStores to be + searched. It should only be set if engine is used. search_engine_id: The Vertex AI search engine resource ID in the format of "projects/{project}/locations/{location}/collections/{collection}/engines/{engine}". @@ -64,7 +69,12 @@ def __init__( raise ValueError( 'Either data_store_id or search_engine_id must be specified.' ) + if data_store_specs is not None and search_engine_id is None: + raise ValueError( + 'search_engine_id must be specified if data_store_specs is specified.' + ) self.data_store_id = data_store_id + self.data_store_specs = data_store_specs self.search_engine_id = search_engine_id self.filter = filter self.max_results = max_results @@ -89,6 +99,7 @@ async def process_llm_request( retrieval=types.Retrieval( vertex_ai_search=types.VertexAISearch( datastore=self.data_store_id, + data_store_specs=self.data_store_specs, engine=self.search_engine_id, filter=self.filter, max_results=self.max_results,