diff --git a/pyproject.toml b/pyproject.toml index 14756309..95c69f1f 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 9a964786..21c82774 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/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): diff --git a/src/google/adk/tools/__init__.py b/src/google/adk/tools/__init__.py index c45b72b9..50d01dfe 100644 --- a/src/google/adk/tools/__init__.py +++ b/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', diff --git a/src/google/adk/tools/discovery_engine_search_tool.py b/src/google/adk/tools/discovery_engine_search_tool.py new file mode 100644 index 00000000..3eeb40f3 --- /dev/null +++ b/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} diff --git a/tests/unittests/agents/test_llm_agent_fields.py b/tests/unittests/agents/test_llm_agent_fields.py index bac0aace..42d02b4b 100644 --- a/tests/unittests/agents/test_llm_agent_fields.py +++ b/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 @@ class TestCanonicalTools: 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 @@ class TestCanonicalTools: 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 @@ class TestCanonicalTools: 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' diff --git a/tests/unittests/tools/test_discovery_engine_search_tool.py b/tests/unittests/tools/test_discovery_engine_search_tool.py new file mode 100644 index 00000000..08313789 --- /dev/null +++ b/tests/unittests/tools/test_discovery_engine_search_tool.py @@ -0,0 +1,134 @@ +# 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 unittest import mock + +from google.adk.tools.discovery_engine_search_tool import DiscoveryEngineSearchTool +from google.api_core import exceptions +from google.cloud import discoveryengine_v1beta as discoveryengine +import pytest + + +@mock.patch( + "google.auth.default", + mock.MagicMock(return_value=("credentials", "project")), +) +class TestDiscoveryEngineSearchTool: + """Test the DiscoveryEngineSearchTool class.""" + + def test_init_with_data_store_id(self): + """Test initialization with data_store_id.""" + tool = DiscoveryEngineSearchTool(data_store_id="test_data_store") + assert ( + tool._serving_config == "test_data_store/servingConfigs/default_config" + ) + + def test_init_with_search_engine_id(self): + """Test initialization with search_engine_id.""" + tool = DiscoveryEngineSearchTool(search_engine_id="test_search_engine") + assert ( + tool._serving_config + == "test_search_engine/servingConfigs/default_config" + ) + + def test_init_with_no_ids_raises_error(self): + """Test that initialization with no IDs raises ValueError.""" + with pytest.raises( + ValueError, + match="Either data_store_id or search_engine_id must be specified.", + ): + DiscoveryEngineSearchTool() + + def test_init_with_both_ids_raises_error(self): + """Test that initialization with both IDs raises ValueError.""" + with pytest.raises( + ValueError, + match="Either data_store_id or search_engine_id must be specified.", + ): + DiscoveryEngineSearchTool( + data_store_id="test_data_store", + search_engine_id="test_search_engine", + ) + + def test_init_with_data_store_specs_without_search_engine_id_raises_error( + self, + ): + """Test that data_store_specs without search_engine_id raises ValueError.""" + with pytest.raises( + ValueError, + match=( + "search_engine_id must be specified if data_store_specs is" + " specified." + ), + ): + DiscoveryEngineSearchTool( + data_store_id="test_data_store", data_store_specs=[{"id": "123"}] + ) + + @mock.patch( + "google.cloud.discoveryengine_v1beta.SearchServiceClient", + ) + def test_discovery_engine_search_success(self, mock_search_client): + """Test successful discovery engine search.""" + mock_response = discoveryengine.SearchResponse() + mock_response.results = [ + discoveryengine.SearchResponse.SearchResult( + chunk=discoveryengine.Chunk( + document_metadata={ + "title": "Test Title", + "uri": "http://example.com", + }, + content="Test Content", + ) + ) + ] + mock_search_client.return_value.search.return_value = mock_response + + tool = DiscoveryEngineSearchTool(data_store_id="test_data_store") + result = tool.discovery_engine_search("test query") + + assert result["status"] == "success" + assert len(result["results"]) == 1 + assert result["results"][0]["title"] == "Test Title" + assert result["results"][0]["url"] == "http://example.com" + assert result["results"][0]["content"] == "Test Content" + + @mock.patch( + "google.cloud.discoveryengine_v1beta.SearchServiceClient", + ) + def test_discovery_engine_search_api_error(self, mock_search_client): + """Test discovery engine search with API error.""" + mock_search_client.return_value.search.side_effect = ( + exceptions.GoogleAPICallError("API error") + ) + + tool = DiscoveryEngineSearchTool(data_store_id="test_data_store") + result = tool.discovery_engine_search("test query") + + assert result["status"] == "error" + assert result["error_message"] == "None API error" + + @mock.patch( + "google.cloud.discoveryengine_v1beta.SearchServiceClient", + ) + def test_discovery_engine_search_no_results(self, mock_search_client): + """Test discovery engine search with no results.""" + mock_response = discoveryengine.SearchResponse() + mock_search_client.return_value.search.return_value = mock_response + + tool = DiscoveryEngineSearchTool(data_store_id="test_data_store") + result = tool.discovery_engine_search("test query") + + assert result["status"] == "success" + assert not result["results"]