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
2026-07-09 18:19:28 -07:00 · 2026-01-29 17:09:49 -08:00
parent c0c98d94b3
commit 585ebfdac7
2 changed files with 165 additions and 14 deletions
@@ -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