ADK changes

PiperOrigin-RevId: 816288113

Author: xuanyang15

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
   "google-api-python-client>=2.157.0, <3.0.0",              # Google API client discovery
   "google-cloud-aiplatform[agent_engines]>=1.112.0, <2.0.0",# For VertexAI integrations, e.g. example store.
   "google-cloud-bigtable>=2.32.0",                          # For Bigtable database
+  "google-cloud-discoveryengine>=0.13.12, <0.14.0",         # For Discovery Engine Search Tool
   "google-cloud-secret-manager>=2.22.0, <3.0.0",            # Fetching secrets in RestAPI Tool
   "google-cloud-spanner>=3.56.0, <4.0.0",                   # For Spanner database
   "google-cloud-speech>=2.30.0, <3.0.0",                    # For Audio Transcription

src/google/adk/agents/llm_agent.py

@@ -21,6 +21,7 @@
 from typing import AsyncGenerator
 from typing import Awaitable
 from typing import Callable
+from typing import cast
 from typing import ClassVar
 from typing import Dict
 from typing import Literal
@@ -118,6 +119,7 @@ async def _convert_tool_union_to_tools(
     multiple_tools: bool = False,
 ) -> list[BaseTool]:
   from ..tools.google_search_tool import google_search
+  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.
@@ -128,6 +130,24 @@ async def _convert_tool_union_to_tools(
 
     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
+  # other tools.
+  # TODO(b/448114567): Remove once the workaround is no longer needed.
+  if multiple_tools and isinstance(tool_union, VertexAiSearchTool):
+    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 isinstance(tool_union, BaseTool):
     return [tool_union]
   if callable(tool_union):

src/google/adk/tools/__init__.py

@@ -18,6 +18,7 @@
 from .agent_tool import AgentTool
 from .apihub_tool.apihub_toolset import APIHubToolset
 from .base_tool import BaseTool
+from .discovery_engine_search_tool import DiscoveryEngineSearchTool
 from .enterprise_search_tool import enterprise_web_search_tool as enterprise_web_search
 from .example_tool import ExampleTool
 from .exit_loop_tool import exit_loop
@@ -39,6 +40,7 @@
     'APIHubToolset',
     'AuthToolArguments',
     'BaseTool',
+    'DiscoveryEngineSearchTool',
     'enterprise_web_search',
     'google_maps_grounding',
     'google_search',

src/google/adk/tools/discovery_engine_search_tool.py

@@ -0,0 +1,126 @@
+# 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 Optional
+
+from google.api_core.exceptions import GoogleAPICallError
+import google.auth
+from google.cloud import discoveryengine_v1beta as discoveryengine
+from google.genai import types
+
+from .function_tool import FunctionTool
+
+
+class DiscoveryEngineSearchTool(FunctionTool):
+  """Tool for searching the discovery engine."""
+
+  def __init__(
+      self,
+      data_store_id: Optional[str] = None,
+      data_store_specs: Optional[
+          list[types.VertexAISearchDataStoreSpec]
+      ] = None,
+      search_engine_id: Optional[str] = None,
+      filter: Optional[str] = None,
+      max_results: Optional[int] = None,
+  ):
+    """Initializes the DiscoveryEngineSearchTool.
+
+    Args:
+      data_store_id: The Vertex AI search data store resource ID in the format
+        of
+        "projects/{project}/locations/{location}/collections/{collection}/dataStores/{dataStore}".
+      data_store_specs: Specifications that define the specific DataStores to be
+        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 be applied to the search request. Default is None.
+      max_results: The maximum number of results to return. Default is None.
+    """
+    super().__init__(self.discovery_engine_search)
+    if (data_store_id is None and search_engine_id is None) or (
+        data_store_id is not None and search_engine_id is not None
+    ):
+      raise ValueError(
+          "Either data_store_id or search_engine_id must be specified."
+      )
+    if data_store_specs is not None and search_engine_id is None:
+      raise ValueError(
+          "search_engine_id must be specified if data_store_specs is specified."
+      )
+
+    self._serving_config = (
+        f"{data_store_id or search_engine_id}/servingConfigs/default_config"
+    )
+    self._data_store_specs = data_store_specs
+    self._search_engine_id = search_engine_id
+    self._filter = filter
+    self._max_results = max_results
+
+    credentials, _ = google.auth.default()
+    self._discovery_engine_client = discoveryengine.SearchServiceClient(
+        credentials=credentials
+    )
+
+  def discovery_engine_search(
+      self,
+      query: str,
+  ) -> dict[str, Any]:
+    """Search the discovery engine.
+
+    Args:
+      query: The search query.
+
+    Returns:
+      A dictionary containing the status of the request and the list of search
+      results, which contains the title, url and content.
+    """
+    request = discoveryengine.SearchRequest(
+        serving_config=self._serving_config,
+        query=query,
+        content_search_spec=discoveryengine.SearchRequest.ContentSearchSpec(
+            search_result_mode=discoveryengine.SearchRequest.ContentSearchSpec.SearchResultMode.CHUNKS,
+            chunk_spec=discoveryengine.SearchRequest.ContentSearchSpec.ChunkSpec(
+                num_previous_chunks=0,
+                num_next_chunks=0,
+            ),
+        ),
+    )
+
+    if self._data_store_specs:
+      request.data_store_specs = self._data_store_specs
+    if self._filter:
+      request.filter = self._filter
+    if self._max_results:
+      request.page_size = self._max_results
+
+    results = []
+    try:
+      response = self._discovery_engine_client.search(request)
+      for item in response.results:
+        chunk = item.chunk
+        if not chunk or not chunk.document_metadata:
+          continue
+
+        results.append({
+            "title": chunk.document_metadata.title,
+            "url": chunk.document_metadata.uri,
+            "content": chunk.content,
+        })
+    except GoogleAPICallError as e:
+      return {"status": "error", "error_message": str(e)}
+    return {"status": "success", "results": results}

tests/unittests/agents/test_llm_agent_fields.py

@@ -15,18 +15,17 @@
 """Unit tests for canonical_xxx fields in LlmAgent."""
 
 from typing import Any
-from typing import cast
 from typing import Optional
 
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.llm_agent import LlmAgent
-from google.adk.agents.loop_agent import LoopAgent
 from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.models.llm_request import LlmRequest
 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.vertex_ai_search_tool import VertexAiSearchTool
 from google.genai import types
 from pydantic import BaseModel
 import pytest
@@ -306,6 +305,7 @@ async def test_handle_google_search_with_other_tools(self):
 
     assert len(tools) == 2
     assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
     assert tools[1].name == 'google_search_agent'
     assert tools[1].__class__.__name__ == 'GoogleSearchAgentTool'
 
@@ -325,8 +325,8 @@ async def test_handle_google_search_only(self):
     assert tools[0].name == 'google_search'
     assert tools[0].__class__.__name__ == 'GoogleSearchTool'
 
-  async def test_no_google_search(self):
-    """Test other tools are not affected."""
+  async def test_function_tool_only(self):
+    """Test that function tool is not affected."""
     agent = LlmAgent(
         name='test_agent',
         model='gemini-pro',
@@ -340,3 +340,38 @@ async def test_no_google_search(self):
     assert len(tools) == 1
     assert tools[0].name == '_my_tool'
     assert tools[0].__class__.__name__ == 'FunctionTool'
+
+  async def test_handle_google_vais_with_other_tools(self):
+    """Test that VertexAiSearchTool is wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            VertexAiSearchTool(data_store_id='test_data_store_id'),
+        ],
+    )
+    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 == 'discovery_engine_search'
+    assert tools[1].__class__.__name__ == 'DiscoveryEngineSearchTool'
+
+  async def test_handle_vais_only(self):
+    """Test that VertexAiSearchTool is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            VertexAiSearchTool(data_store_id='test_data_store_id'),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == 'vertex_ai_search'
+    assert tools[0].__class__.__name__ == 'VertexAiSearchTool'