feat: Support dynamic config for VertexAiSearchTool

Close: https://github.com/google/adk-python/issues/4067

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 862965769

src/google/adk/tools/vertex_ai_search_tool.py

@@ -21,6 +21,7 @@ from typing import TYPE_CHECKING
 from google.genai import types
 from typing_extensions import override
 
+from ..agents.readonly_context import ReadonlyContext
 from ..utils.model_name_utils import is_gemini_1_model
 from ..utils.model_name_utils import is_gemini_model
 from .base_tool import BaseTool
@@ -38,6 +39,25 @@ class VertexAiSearchTool(BaseTool):
   Attributes:
     data_store_id: The Vertex AI search data store resource ID.
     search_engine_id: The Vertex AI search engine resource ID.
+
+  To dynamically customize the search configuration at runtime (e.g., set
+  filter based on user context), subclass this tool and override the
+  `_build_vertex_ai_search_config` method.
+
+  Example:
+    ```python
+    class DynamicFilterSearchTool(VertexAiSearchTool):
+      def _build_vertex_ai_search_config(
+          self, ctx: ReadonlyContext
+      ) -> types.VertexAISearch:
+        user_id = ctx.state.get('user_id')
+        return types.VertexAISearch(
+            datastore=self.data_store_id,
+            engine=self.search_engine_id,
+            filter=f"user_id = '{user_id}'",
+            max_results=self.max_results,
+        )
+    ```
   """
 
   def __init__(
@@ -90,6 +110,30 @@ class VertexAiSearchTool(BaseTool):
     self.max_results = max_results
     self.bypass_multi_tools_limit = bypass_multi_tools_limit
 
+  def _build_vertex_ai_search_config(
+      self, readonly_context: ReadonlyContext
+  ) -> types.VertexAISearch:
+    """Builds the VertexAISearch configuration.
+
+    Override this method in a subclass to dynamically customize the search
+    configuration based on the context (e.g., set filter based on session
+    state).
+
+    Args:
+      readonly_context: The readonly context with access to state and session
+        info.
+
+    Returns:
+      The VertexAISearch configuration to use for this request.
+    """
+    return types.VertexAISearch(
+        datastore=self.data_store_id,
+        data_store_specs=self.data_store_specs,
+        engine=self.search_engine_id,
+        filter=self.filter,
+        max_results=self.max_results,
+    )
+
   @override
   async def process_llm_request(
       self,
@@ -106,14 +150,20 @@ class VertexAiSearchTool(BaseTool):
       llm_request.config = llm_request.config or types.GenerateContentConfig()
       llm_request.config.tools = llm_request.config.tools or []
 
+      # Build the search config (can be overridden by subclasses)
+      vertex_ai_search_config = self._build_vertex_ai_search_config(
+          tool_context
+      )
+
       # Format data_store_specs concisely for logging
-      if self.data_store_specs:
+      if vertex_ai_search_config.data_store_specs:
         spec_ids = [
             spec.data_store.split('/')[-1] if spec.data_store else 'unnamed'
-            for spec in self.data_store_specs
+            for spec in vertex_ai_search_config.data_store_specs
         ]
         specs_info = (
-            f'{len(self.data_store_specs)} spec(s): [{", ".join(spec_ids)}]'
+            f'{len(vertex_ai_search_config.data_store_specs)} spec(s):'
+            f' [{", ".join(spec_ids)}]'
         )
       else:
         specs_info = None
@@ -122,23 +172,17 @@ class VertexAiSearchTool(BaseTool):
           'Adding Vertex AI Search tool config to LLM request: '
           'datastore=%s, engine=%s, filter=%s, max_results=%s, '
           'data_store_specs=%s',
-          self.data_store_id,
-          self.search_engine_id,
-          self.filter,
-          self.max_results,
+          vertex_ai_search_config.datastore,
+          vertex_ai_search_config.engine,
+          vertex_ai_search_config.filter,
+          vertex_ai_search_config.max_results,
           specs_info,
       )
 
       llm_request.config.tools.append(
           types.Tool(
               retrieval=types.Retrieval(
-                  vertex_ai_search=types.VertexAISearch(
-                      datastore=self.data_store_id,
-                      data_store_specs=self.data_store_specs,
-                      engine=self.search_engine_id,
-                      filter=self.filter,
-                      max_results=self.max_results,
-                  )
+                  vertex_ai_search=vertex_ai_search_config
               )
           )
       )

tests/unittests/tools/test_vertex_ai_search_tool.py

@@ -449,3 +449,110 @@ class TestVertexAiSearchTool:
     assert 'filter=None' in log_message
     assert 'max_results=None' in log_message
     assert 'data_store_specs=None' in log_message
+
+  @pytest.mark.asyncio
+  async def test_subclass_with_dynamic_filter(self):
+    """Test subclassing to provide dynamic filter based on context."""
+
+    class DynamicFilterSearchTool(VertexAiSearchTool):
+      """Custom search tool with dynamic filter."""
+
+      def _build_vertex_ai_search_config(self, ctx):
+        user_id = ctx.state.get('user_id', 'default_user')
+        return types.VertexAISearch(
+            datastore=self.data_store_id,
+            engine=self.search_engine_id,
+            filter=f"user_id = '{user_id}'",
+            max_results=self.max_results,
+        )
+
+    tool = DynamicFilterSearchTool(data_store_id='test_data_store')
+    tool_context = await _create_tool_context()
+    tool_context.state['user_id'] = 'test_user_123'
+
+    llm_request = LlmRequest(
+        model='gemini-2.5-pro', config=types.GenerateContentConfig()
+    )
+
+    await tool.process_llm_request(
+        tool_context=tool_context, llm_request=llm_request
+    )
+
+    assert llm_request.config.tools is not None
+    assert len(llm_request.config.tools) == 1
+    retrieval_tool = llm_request.config.tools[0]
+    assert retrieval_tool.retrieval is not None
+    assert retrieval_tool.retrieval.vertex_ai_search is not None
+    # Verify the filter was dynamically set
+    assert (
+        retrieval_tool.retrieval.vertex_ai_search.filter
+        == "user_id = 'test_user_123'"
+    )
+
+  @pytest.mark.asyncio
+  async def test_subclass_with_dynamic_max_results(self):
+    """Test subclassing to provide dynamic max_results based on context."""
+
+    class DynamicMaxResultsSearchTool(VertexAiSearchTool):
+      """Custom search tool with dynamic max_results."""
+
+      def _build_vertex_ai_search_config(self, ctx):
+        # Use a larger max_results for premium users
+        is_premium = ctx.state.get('is_premium', False)
+        dynamic_max_results = 20 if is_premium else 5
+        return types.VertexAISearch(
+            datastore=self.data_store_id,
+            engine=self.search_engine_id,
+            filter=self.filter,
+            max_results=dynamic_max_results,
+        )
+
+    tool = DynamicMaxResultsSearchTool(
+        data_store_id='test_data_store', max_results=10
+    )
+    tool_context = await _create_tool_context()
+    tool_context.state['is_premium'] = True
+
+    llm_request = LlmRequest(
+        model='gemini-2.5-pro', config=types.GenerateContentConfig()
+    )
+
+    await tool.process_llm_request(
+        tool_context=tool_context, llm_request=llm_request
+    )
+
+    retrieval_tool = llm_request.config.tools[0]
+    # Verify max_results was dynamically set to premium value
+    assert retrieval_tool.retrieval.vertex_ai_search.max_results == 20
+
+  @pytest.mark.asyncio
+  async def test_subclass_receives_readonly_context(self):
+    """Test that subclass receives the context correctly."""
+    received_contexts = []
+
+    class ContextCapturingSearchTool(VertexAiSearchTool):
+      """Custom search tool that captures the context."""
+
+      def _build_vertex_ai_search_config(self, ctx):
+        received_contexts.append(ctx)
+        return types.VertexAISearch(
+            datastore=self.data_store_id,
+            engine=self.search_engine_id,
+            filter=self.filter,
+            max_results=self.max_results,
+        )
+
+    tool = ContextCapturingSearchTool(data_store_id='test_data_store')
+    tool_context = await _create_tool_context()
+
+    llm_request = LlmRequest(
+        model='gemini-2.5-pro', config=types.GenerateContentConfig()
+    )
+
+    await tool.process_llm_request(
+        tool_context=tool_context, llm_request=llm_request
+    )
+
+    # Verify the context was passed to _build_vertex_ai_search_config
+    assert len(received_contexts) == 1
+    assert received_contexts[0] is tool_context