Skip to content

Commit fcb1893

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Allow data_store_specs pass into ADK VAIS built-in tool
PiperOrigin-RevId: 770733583
1 parent a4d432a commit fcb1893

File tree

1 file changed

+11
-0
lines changed

1 file changed

+11
-0
lines changed

src/google/adk/tools/vertex_ai_search_tool.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)