ADK changes

PiperOrigin-RevId: 816288113

src/google/adk/agents/llm_agent.py

@@ -21,6 +21,7 @@ from typing import Any
 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 ..auth.auth_tool import AuthToolArguments
 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 @@ __all__ = [
     '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}