feat: Add bypass_multi_tools_limit option to GoogleSearchTool and VertexAiSearchTool

PiperOrigin-RevId: 817493869

Author: xuanyang15

src/google/adk/agents/llm_agent.py

@@ -118,17 +118,19 @@ async def _convert_tool_union_to_tools(
     model: Union[str, BaseLlm],
     multiple_tools: bool = False,
 ) -> list[BaseTool]:
-  from ..tools.google_search_tool import google_search
+  from ..tools.google_search_tool import GoogleSearchTool
   from ..tools.vertex_ai_search_tool import VertexAiSearchTool
 
   # Wrap google_search tool with AgentTool if there are multiple tools because
   # the built-in tools cannot be used together with other tools.
   # TODO(b/448114567): Remove once the workaround is no longer needed.
-  if multiple_tools and tool_union is google_search:
+  if multiple_tools and isinstance(tool_union, GoogleSearchTool):
     from ..tools.google_search_agent_tool import create_google_search_agent
     from ..tools.google_search_agent_tool import GoogleSearchAgentTool
 
-    return [GoogleSearchAgentTool(create_google_search_agent(model))]
+    search_tool = cast(GoogleSearchTool, tool_union)
+    if search_tool.bypass_multi_tools_limit:
+      return [GoogleSearchAgentTool(create_google_search_agent(model))]
 
   # Replace VertexAiSearchTool with DiscoveryEngineSearchTool if there are
   # multiple tools because the built-in tools cannot be used together with
@@ -138,15 +140,16 @@ async def _convert_tool_union_to_tools(
     from ..tools.discovery_engine_search_tool import DiscoveryEngineSearchTool
 
     vais_tool = cast(VertexAiSearchTool, tool_union)
-    return [
-        DiscoveryEngineSearchTool(
-            data_store_id=vais_tool.data_store_id,
-            data_store_specs=vais_tool.data_store_specs,
-            search_engine_id=vais_tool.search_engine_id,
-            filter=vais_tool.filter,
-            max_results=vais_tool.max_results,
-        )
-    ]
+    if vais_tool.bypass_multi_tools_limit:
+      return [
+          DiscoveryEngineSearchTool(
+              data_store_id=vais_tool.data_store_id,
+              data_store_specs=vais_tool.data_store_specs,
+              search_engine_id=vais_tool.search_engine_id,
+              filter=vais_tool.filter,
+              max_results=vais_tool.max_results,
+          )
+      ]
 
   if isinstance(tool_union, BaseTool):
     return [tool_union]

src/google/adk/tools/google_search_tool.py

@@ -35,9 +35,17 @@ class GoogleSearchTool(BaseTool):
   local code execution.
   """
 
-  def __init__(self):
+  def __init__(self, *, bypass_multi_tools_limit: bool = True):
+    """Initializes the Google search tool.
+
+    Args:
+      bypass_multi_tools_limit: Whether to bypass the multi tools limitation,
+        so that the tool can be used with other tools in the same agent.
+    """
+
     # Name and description are not used because this is a model built-in tool.
     super().__init__(name='google_search', description='google_search')
+    self.bypass_multi_tools_limit = bypass_multi_tools_limit
 
   @override
   async def process_llm_request(

src/google/adk/tools/vertex_ai_search_tool.py

@@ -47,6 +47,7 @@ def __init__(
       search_engine_id: Optional[str] = None,
       filter: Optional[str] = None,
       max_results: Optional[int] = None,
+      bypass_multi_tools_limit: bool = True,
   ):
     """Initializes the Vertex AI Search tool.
 
@@ -58,6 +59,10 @@ def __init__(
         searched. It should only be set if engine is used.
       search_engine_id: The Vertex AI search engine resource ID in the format of
         "projects/{project}/locations/{location}/collections/{collection}/engines/{engine}".
+      filter: The filter to apply to the search results.
+      max_results: The maximum number of results to return.
+      bypass_multi_tools_limit: Whether to bypass the multi tools limitation,
+        so that the tool can be used with other tools in the same agent.
 
     Raises:
       ValueError: If both data_store_id and search_engine_id are not specified
@@ -80,6 +85,7 @@ def __init__(
     self.search_engine_id = search_engine_id
     self.filter = filter
     self.max_results = max_results
+    self.bypass_multi_tools_limit = bypass_multi_tools_limit
 
   @override
   async def process_llm_request(

tests/unittests/agents/test_llm_agent_fields.py

@@ -26,6 +26,7 @@
 from google.adk.models.registry import LLMRegistry
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.tools.google_search_tool import google_search
+from google.adk.tools.google_search_tool import GoogleSearchTool
 from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool
 from google.genai import types
 from pydantic import BaseModel
@@ -310,6 +311,25 @@ async def test_handle_google_search_with_other_tools(self):
     assert tools[1].name == 'google_search_agent'
     assert tools[1].__class__.__name__ == 'GoogleSearchAgentTool'
 
+  async def test_handle_google_search_with_other_tools_no_bypass(self):
+    """Test that google_search is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            GoogleSearchTool(bypass_multi_tools_limit=False),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'google_search'
+    assert tools[1].__class__.__name__ == 'GoogleSearchTool'
+
   async def test_handle_google_search_only(self):
     """Test that google_search is not wrapped into an agent."""
     agent = LlmAgent(
@@ -346,8 +366,8 @@ async def test_function_tool_only(self):
       'google.auth.default',
       mock.MagicMock(return_value=('credentials', 'project')),
   )
-  async def test_handle_google_vais_with_other_tools(self):
-    """Test that VertexAiSearchTool is wrapped into an agent."""
+  async def test_handle_vais_with_other_tools(self):
+    """Test that VertexAiSearchTool is replaced with Discovery Engine Search."""
     agent = LlmAgent(
         name='test_agent',
         model='gemini-pro',
@@ -365,6 +385,28 @@ async def test_handle_google_vais_with_other_tools(self):
     assert tools[1].name == 'discovery_engine_search'
     assert tools[1].__class__.__name__ == 'DiscoveryEngineSearchTool'
 
+  async def test_handle_vais_with_other_tools_no_bypass(self):
+    """Test that VertexAiSearchTool is not replaced."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            VertexAiSearchTool(
+                data_store_id='test_data_store_id',
+                bypass_multi_tools_limit=False,
+            ),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'vertex_ai_search'
+    assert tools[1].__class__.__name__ == 'VertexAiSearchTool'
+
   async def test_handle_vais_only(self):
     """Test that VertexAiSearchTool is not wrapped into an agent."""
     agent = LlmAgent(