diff --git a/src/google/adk/tools/discovery_engine_search_tool.py b/src/google/adk/tools/discovery_engine_search_tool.py index 0e771ece4f..1ac77db23d 100644 --- a/src/google/adk/tools/discovery_engine_search_tool.py +++ b/src/google/adk/tools/discovery_engine_search_tool.py @@ -17,6 +17,7 @@ from typing import Any from typing import Optional +from google.api_core import client_options as client_options_lib from google.api_core.exceptions import GoogleAPICallError import google.auth from google.cloud import discoveryengine_v1beta as discoveryengine @@ -37,6 +38,9 @@ def __init__( search_engine_id: Optional[str] = None, filter: Optional[str] = None, max_results: Optional[int] = None, + client_options: Optional[ + Union[client_options_lib.ClientOptions, dict] + ] = None, ): """Initializes the DiscoveryEngineSearchTool. @@ -73,7 +77,8 @@ def __init__( credentials, _ = google.auth.default() self._discovery_engine_client = discoveryengine.SearchServiceClient( - credentials=credentials + client_options=client_options, + credentials=credentials, ) def discovery_engine_search( diff --git a/tests/unittests/tools/test_discovery_engine_search_tool.py b/tests/unittests/tools/test_discovery_engine_search_tool.py index d10da252c7..983d3295c5 100644 --- a/tests/unittests/tools/test_discovery_engine_search_tool.py +++ b/tests/unittests/tools/test_discovery_engine_search_tool.py @@ -76,6 +76,14 @@ def test_init_with_data_store_specs_without_search_engine_id_raises_error( data_store_id="test_data_store", data_store_specs=[{"id": "123"}] ) + def test_init_with_client_options(self): + """Test initialization with client_options""" + tool = DiscoveryEngineSearchTool( + data_store_id="test_data_store", + client_options={"api_endpoint": "test-discoveryengine"}, + ) + assert tool._discovery_engine_client._api_endpoint == "test-discoveryengine" + @mock.patch( "google.cloud.discoveryengine_v1beta.SearchServiceClient", )