@@ -39,6 +39,9 @@ def __init__(
3939 self ,
4040 * ,
4141 data_store_id : Optional [str ] = None ,
42+ data_store_specs : Optional [
43+ list [types .VertexAISearchDataStoreSpec ]
44+ ] = None ,
4245 search_engine_id : Optional [str ] = None ,
4346 filter : Optional [str ] = None ,
4447 max_results : Optional [int ] = None ,
@@ -49,6 +52,8 @@ def __init__(
4952 data_store_id: The Vertex AI search data store resource ID in the format
5053 of
5154 "projects/{project}/locations/{location}/collections/{collection}/dataStores/{dataStore}".
55+ data_store_specs: Specifications that define the specific DataStores to be
56+ searched. It should only be set if engine is used.
5257 search_engine_id: The Vertex AI search engine resource ID in the format of
5358 "projects/{project}/locations/{location}/collections/{collection}/engines/{engine}".
5459
@@ -64,7 +69,12 @@ def __init__(
6469 raise ValueError (
6570 'Either data_store_id or search_engine_id must be specified.'
6671 )
72+ if data_store_specs is not None and search_engine_id is None :
73+ raise ValueError (
74+ 'search_engine_id must be specified if data_store_specs is specified.'
75+ )
6776 self .data_store_id = data_store_id
77+ self .data_store_specs = data_store_specs
6878 self .search_engine_id = search_engine_id
6979 self .filter = filter
7080 self .max_results = max_results
@@ -89,6 +99,7 @@ async def process_llm_request(
8999 retrieval = types .Retrieval (
90100 vertex_ai_search = types .VertexAISearch (
91101 datastore = self .data_store_id ,
102+ data_store_specs = self .data_store_specs ,
92103 engine = self .search_engine_id ,
93104 filter = self .filter ,
94105 max_results = self .max_results ,
0 commit comments