diff --git a/src/google/adk/tools/vertex_ai_search_tool.py b/src/google/adk/tools/vertex_ai_search_tool.py index 7344dc39..91fe60e5 100644 --- a/src/google/adk/tools/vertex_ai_search_tool.py +++ b/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 ) ) ) diff --git a/tests/unittests/tools/test_vertex_ai_search_tool.py b/tests/unittests/tools/test_vertex_ai_search_tool.py index 6a743152..3ade634d 100644 --- a/tests/unittests/tools/test_vertex_ai_search_tool.py +++ b/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