feat: expose more config of VertexAiSearchTool from latest Google GenAI SDK

seanzhougoogle · copybara-github · commit 2b5c89b3a94e · 2025-05-28T22:32:44.000-07:00
for fixing: #955 PiperOrigin-RevId: 764568728
diff --git a/src/google/adk/tools/vertex_ai_search_tool.py b/src/google/adk/tools/vertex_ai_search_tool.py
@@ -40,6 +40,8 @@ def __init__(
       *,
       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 @@ def __init__(
       )
     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 @@ async def process_llm_request(
           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,
                   )
               )
           )
