feat: implement enterprise search agent tool and related functionality

:

Author: m-baden

src/google/adk/agents/llm_agent.py

@@ -138,6 +138,7 @@ async def _convert_tool_union_to_tools(
     model: Union[str, BaseLlm],
     multiple_tools: bool = False,
 ) -> list[BaseTool]:
+  from ..tools.enterprise_search_tool import EnterpriseWebSearchTool
   from ..tools.google_search_tool import GoogleSearchTool
   from ..tools.vertex_ai_search_tool import VertexAiSearchTool
 
@@ -171,6 +172,17 @@ async def _convert_tool_union_to_tools(
           )
       ]
 
+  # Wrap enterprise_web_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 isinstance(tool_union, EnterpriseWebSearchTool):
+    from ..tools.enterprise_search_agent_tool import create_enterprise_search_agent
+    from ..tools.enterprise_search_agent_tool import EnterpriseSearchAgentTool
+
+    enterprise_tool = cast(EnterpriseWebSearchTool, tool_union)
+    if enterprise_tool.bypass_multi_tools_limit:
+      return [EnterpriseSearchAgentTool(create_enterprise_search_agent(model))]
+
   if isinstance(tool_union, BaseTool):
     return [tool_union]
   if callable(tool_union):

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -846,7 +846,11 @@ async def _maybe_add_grounding_metadata(
         tools = await agent.canonical_tools(readonly_context)
         invocation_context.canonical_tools_cache = tools
 
-      if not any(tool.name == 'google_search_agent' for tool in tools):
+      if not any(
+          tool.name == 'google_search_agent'
+          or tool.name == 'enterprise_search_agent'
+          for tool in tools
+      ):
         return response
       ground_metadata = invocation_context.session.state.get(
           'temp:_adk_grounding_metadata', None

src/google/adk/tools/enterprise_search_agent_tool.py

@@ -0,0 +1,141 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any
+from typing import Union
+
+from google.genai import types
+from typing_extensions import override
+
+from ..agents.llm_agent import LlmAgent
+from ..memory.in_memory_memory_service import InMemoryMemoryService
+from ..models.base_llm import BaseLlm
+from ..runners import Runner
+from ..sessions.in_memory_session_service import InMemorySessionService
+from ..utils.context_utils import Aclosing
+from ._forwarding_artifact_service import ForwardingArtifactService
+from .agent_tool import AgentTool
+from .enterprise_search_tool import enterprise_web_search_tool
+from .tool_context import ToolContext
+
+
+def create_enterprise_search_agent(model: Union[str, BaseLlm]) -> LlmAgent:
+  """Create a sub-agent that only uses enterprise_web_search tool."""
+  return LlmAgent(
+      name='enterprise_search_agent',
+      model=model,
+      description=(
+          'An agent for performing Enterprise search using the'
+          ' `enterprise_web_search` tool'
+      ),
+      instruction="""
+        You are a specialized Enterprise search agent.
+
+        When given a search query, use the `enterprise_web_search` tool to find the related information.
+      """,
+      tools=[enterprise_web_search_tool],
+  )
+
+
+class EnterpriseSearchAgentTool(AgentTool):
+  """A tool that wraps a sub-agent that only uses enterprise_web_search tool.
+
+  This is a workaround to support using enterprise_web_search tool with other tools.
+  TODO(b/448114567): Remove once the workaround is no longer needed.
+
+  Attributes:
+    model: The model to use for the sub-agent.
+  """
+
+  def __init__(self, agent: LlmAgent):
+    self.agent = agent
+    super().__init__(agent=self.agent)
+
+  @override
+  async def run_async(
+      self,
+      *,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+  ) -> Any:
+    from ..agents.llm_agent import LlmAgent
+
+    if isinstance(self.agent, LlmAgent) and self.agent.input_schema:
+      input_value = self.agent.input_schema.model_validate(args)
+      content = types.Content(
+          role='user',
+          parts=[
+              types.Part.from_text(
+                  text=input_value.model_dump_json(exclude_none=True)
+              )
+          ],
+      )
+    else:
+      content = types.Content(
+          role='user',
+          parts=[types.Part.from_text(text=args['request'])],
+      )
+    runner = Runner(
+        app_name=self.agent.name,
+        agent=self.agent,
+        artifact_service=ForwardingArtifactService(tool_context),
+        session_service=InMemorySessionService(),
+        memory_service=InMemoryMemoryService(),
+        credential_service=tool_context._invocation_context.credential_service,
+        plugins=list(tool_context._invocation_context.plugin_manager.plugins),
+    )
+
+    state_dict = {
+        k: v
+        for k, v in tool_context.state.to_dict().items()
+        if not k.startswith('_adk')  # Filter out adk internal states
+    }
+    session = await runner.session_service.create_session(
+        app_name=self.agent.name,
+        user_id=tool_context._invocation_context.user_id,
+        state=state_dict,
+    )
+
+    last_content = None
+    last_grounding_metadata = None
+    async with Aclosing(
+        runner.run_async(
+            user_id=session.user_id, session_id=session.id, new_message=content
+        )
+    ) as agen:
+      async for event in agen:
+        # Forward state delta to parent session.
+        if event.actions.state_delta:
+          tool_context.state.update(event.actions.state_delta)
+        if event.content:
+          last_content = event.content
+          last_grounding_metadata = event.grounding_metadata
+
+    if not last_content:
+      return ''
+    merged_text = '\n'.join(p.text for p in last_content.parts if p.text)
+    if isinstance(self.agent, LlmAgent) and self.agent.output_schema:
+      tool_result = self.agent.output_schema.model_validate_json(
+          merged_text
+      ).model_dump(exclude_none=True)
+    else:
+      tool_result = merged_text
+
+    if last_grounding_metadata:
+      tool_context.state['temp:_adk_grounding_metadata'] = (
+          last_grounding_metadata
+      )
+    return tool_result

src/google/adk/tools/enterprise_search_tool.py

@@ -35,12 +35,18 @@ class EnterpriseWebSearchTool(BaseTool):
   https://cloud.google.com/vertex-ai/generative-ai/docs/grounding/web-grounding-enterprise.
   """
 
-  def __init__(self):
-    """Initializes the Vertex AI Search tool."""
+  def __init__(self, *, bypass_multi_tools_limit: bool = False):
+    """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='enterprise_web_search', description='enterprise_web_search'
     )
+    self.bypass_multi_tools_limit = bypass_multi_tools_limit
 
   @override
   async def process_llm_request(

tests/unittests/tools/test_enterprise_search_agent_tool.py

@@ -0,0 +1,110 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from unittest import mock
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.run_config import RunConfig
+from google.adk.agents.sequential_agent import SequentialAgent
+from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
+from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
+from google.adk.plugins.plugin_manager import PluginManager
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.enterprise_search_agent_tool import create_enterprise_search_agent
+from google.adk.tools.enterprise_search_agent_tool import EnterpriseSearchAgentTool
+from google.adk.tools.tool_context import ToolContext
+from pytest import mark
+
+
+async def _create_tool_context() -> ToolContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+  agent = SequentialAgent(name='test_agent')
+  invocation_context = InvocationContext(
+      invocation_id='invocation_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+      artifact_service=InMemoryArtifactService(),
+      memory_service=InMemoryMemoryService(),
+      plugin_manager=PluginManager(),
+      run_config=RunConfig(),
+  )
+  return ToolContext(invocation_context=invocation_context)
+
+
+class TestEnterpriseSearchAgentTool:
+  """Test the EnterpriseSearchAgentTool class."""
+
+  def test_create_enterprise_search_agent(self):
+    """Test that create_enterprise_search_agent creates a valid agent."""
+    agent = create_enterprise_search_agent('gemini-pro')
+    assert isinstance(agent, LlmAgent)
+    assert agent.name == 'enterprise_search_agent'
+    assert 'enterprise_web_search' in [t.name for t in agent.tools]
+
+  def test_enterprise_search_agent_tool_init(self):
+    """Test initialization of EnterpriseSearchAgentTool."""
+    mock_agent = mock.MagicMock(spec=LlmAgent)
+    mock_agent.name = 'test_agent'
+    mock_agent.description = 'test_description'
+    tool = EnterpriseSearchAgentTool(mock_agent)
+    assert tool.agent == mock_agent
+
+  @mark.asyncio
+  @mock.patch('google.adk.tools.enterprise_search_agent_tool.Runner')
+  async def test_run_async_succeeds(self, mock_runner_class):
+    """Test that run_async executes the sub-agent and returns the result."""
+    # Arrange
+    mock_agent = mock.MagicMock(spec=LlmAgent)
+    mock_agent.name = 'enterprise_search_agent'
+    mock_agent.description = 'test_description'
+    mock_agent.input_schema = None
+    mock_agent.output_schema = None
+
+    tool = EnterpriseSearchAgentTool(mock_agent)
+    tool_context = await _create_tool_context()
+
+    async def mock_run_async_gen():
+      yield mock.MagicMock(
+          actions=mock.MagicMock(state_delta={'key': 'value'}), content=None
+      )
+      yield mock.MagicMock(
+          actions=mock.MagicMock(state_delta=None),
+          content=mock.MagicMock(parts=[mock.MagicMock(text='test response')]),
+      )
+
+    mock_runner_instance = mock.MagicMock()
+    mock_runner_instance.run_async.return_value = mock_run_async_gen()
+    mock_runner_instance.session_service = mock.AsyncMock()
+    mock_runner_instance.session_service.create_session.return_value = (
+        tool_context._invocation_context.session
+    )
+    mock_runner_class.return_value = mock_runner_instance
+
+    # Act
+    result = await tool.run_async(
+        args={'request': 'test query'}, tool_context=tool_context
+    )
+
+    # Assert
+    mock_runner_class.assert_called_once()
+    mock_runner_instance.run_async.assert_called_once()
+    assert tool_context.state['key'] == 'value'
+    assert result == 'test response'