From 2b5c89b3a94e82ea4a40363ea8de33d9473d7cf0 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Wed, 28 May 2025 22:32:13 -0700 Subject: [PATCH] feat: expose more config of VertexAiSearchTool from latest Google GenAI SDK for fixing: https://github.com/google/adk-python/issues/955 PiperOrigin-RevId: 764568728 --- src/google/adk/tools/vertex_ai_search_tool.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/google/adk/tools/vertex_ai_search_tool.py b/src/google/adk/tools/vertex_ai_search_tool.py index ebe236e..c370e2a 100644 --- a/src/google/adk/tools/vertex_ai_search_tool.py +++ b/src/google/adk/tools/vertex_ai_search_tool.py @@ -40,6 +40,8 @@ class VertexAiSearchTool(BaseTool): *, data_store_id: Optional[str] = None, search_engine_id: Optional[str] = None, + filter: Optional[str] = None, + max_results: Optional[int] = None, ): """Initializes the Vertex AI Search tool. @@ -64,6 +66,8 @@ class VertexAiSearchTool(BaseTool): ) self.data_store_id = data_store_id self.search_engine_id = search_engine_id + self.filter = filter + self.max_results = max_results @override async def process_llm_request( @@ -84,7 +88,10 @@ class VertexAiSearchTool(BaseTool): types.Tool( retrieval=types.Retrieval( vertex_ai_search=types.VertexAISearch( - datastore=self.data_store_id, engine=self.search_engine_id + datastore=self.data_store_id, + engine=self.search_engine_id, + filter=self.filter, + max_results=self.max_results, ) ) )