@@ -39,6 +39,9 @@ def __init__(
39
39
self ,
40
40
* ,
41
41
data_store_id : Optional [str ] = None ,
42
+ data_store_specs : Optional [
43
+ list [types .VertexAISearchDataStoreSpec ]
44
+ ] = None ,
42
45
search_engine_id : Optional [str ] = None ,
43
46
filter : Optional [str ] = None ,
44
47
max_results : Optional [int ] = None ,
@@ -49,6 +52,8 @@ def __init__(
49
52
data_store_id: The Vertex AI search data store resource ID in the format
50
53
of
51
54
"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.
52
57
search_engine_id: The Vertex AI search engine resource ID in the format of
53
58
"projects/{project}/locations/{location}/collections/{collection}/engines/{engine}".
54
59
@@ -64,7 +69,12 @@ def __init__(
64
69
raise ValueError (
65
70
'Either data_store_id or search_engine_id must be specified.'
66
71
)
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
+ )
67
76
self .data_store_id = data_store_id
77
+ self .data_store_specs = data_store_specs
68
78
self .search_engine_id = search_engine_id
69
79
self .filter = filter
70
80
self .max_results = max_results
@@ -89,6 +99,7 @@ async def process_llm_request(
89
99
retrieval = types .Retrieval (
90
100
vertex_ai_search = types .VertexAISearch (
91
101
datastore = self .data_store_id ,
102
+ data_store_specs = self .data_store_specs ,
92
103
engine = self .search_engine_id ,
93
104
filter = self .filter ,
94
105
max_results = self .max_results ,
0 commit comments