fix(tools): add optional location override validation

Author: Gautier Masse

src/google/adk/tools/discovery_engine_search_tool.py

@@ -31,16 +31,42 @@
 _LOCATION_PATTERN = re.compile(r"/locations/([^/]+)(?:/|$)")
 
 
+def _resolve_location(resource_id: str, location: Optional[str]) -> str:
+  """Resolves the Discovery Engine location to use for the endpoint."""
+  location_match = _LOCATION_PATTERN.search(resource_id)
+  inferred_location = (
+      location_match.group(1).lower() if location_match else None
+  )
+
+  if location is not None:
+    normalized_location = location.strip().lower()
+    if not normalized_location:
+      raise ValueError("location must not be empty if specified.")
+    if inferred_location and normalized_location != inferred_location:
+      raise ValueError(
+          "location must match the location in data_store_id or "
+          "search_engine_id."
+      )
+    return normalized_location
+
+  if inferred_location:
+    return inferred_location
+  return _GLOBAL_LOCATION
+
+
 def _build_client_options(
-    resource_id: str, quota_project_id: Optional[str]
+    resource_id: str,
+    quota_project_id: Optional[str],
+    location: Optional[str],
 ) -> Optional[client_options.ClientOptions]:
   """Builds client options for Discovery Engine requests."""
   client_options_kwargs = {}
-  location_match = _LOCATION_PATTERN.search(resource_id)
-  location = location_match.group(1) if location_match else _GLOBAL_LOCATION
+  resolved_location = _resolve_location(resource_id, location)
 
-  if location != _GLOBAL_LOCATION:
-    client_options_kwargs["api_endpoint"] = f"{location}-{_DEFAULT_ENDPOINT}"
+  if resolved_location != _GLOBAL_LOCATION:
+    client_options_kwargs["api_endpoint"] = (
+        f"{resolved_location}-{_DEFAULT_ENDPOINT}"
+    )
   if quota_project_id:
     client_options_kwargs["quota_project_id"] = quota_project_id
 
@@ -61,6 +87,7 @@ def __init__(
       search_engine_id: Optional[str] = None,
       filter: Optional[str] = None,
       max_results: Optional[int] = None,
+      location: Optional[str] = None,
   ):
     """Initializes the DiscoveryEngineSearchTool.
 
@@ -74,6 +101,9 @@ def __init__(
         "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.
+      location: Optional endpoint location override.
+        Examples: "global", "us", "eu". If not specified, location is inferred
+        from `data_store_id` or `search_engine_id` and defaults to "global".
     """
     super().__init__(self.discovery_engine_search)
     if (data_store_id is None and search_engine_id is None) or (
@@ -94,11 +124,16 @@ def __init__(
     self._search_engine_id = search_engine_id
     self._filter = filter
     self._max_results = max_results
+    self._location = location
 
     credentials, _ = google.auth.default()
     quota_project_id = getattr(credentials, "quota_project_id", None)
     resource_id = data_store_id or search_engine_id or ""
-    options = _build_client_options(resource_id, quota_project_id)
+    options = _build_client_options(
+        resource_id=resource_id,
+        quota_project_id=quota_project_id,
+        location=location,
+    )
     self._discovery_engine_client = discoveryengine.SearchServiceClient(
         credentials=credentials, client_options=options
     )

tests/unittests/tools/test_discovery_engine_search_tool.py

@@ -131,6 +131,61 @@ def test_init_with_regional_location_uses_regional_endpoint(
         client_options=mock_client_options.ClientOptions.return_value,
     )
 
+  @mock.patch.object(discovery_engine_search_tool, "client_options")
+  @mock.patch.object(discoveryengine, "SearchServiceClient")
+  def test_init_with_explicit_location_override_uses_input_location(
+      self, mock_search_client, mock_client_options
+  ):
+    """Test initialization uses explicit location when resource has none."""
+    DiscoveryEngineSearchTool(
+        data_store_id="test_data_store",
+        location="eu",
+    )
+
+    mock_client_options.ClientOptions.assert_called_once_with(
+        api_endpoint="eu-discoveryengine.googleapis.com"
+    )
+    mock_search_client.assert_called_once_with(
+        credentials="credentials",
+        client_options=mock_client_options.ClientOptions.return_value,
+    )
+
+  @mock.patch.object(discoveryengine, "SearchServiceClient")
+  def test_init_with_mismatched_location_raises_error(self, mock_search_client):
+    """Test initialization rejects mismatched location overrides."""
+    with pytest.raises(
+        ValueError,
+        match=(
+            "location must match the location in data_store_id or "
+            "search_engine_id."
+        ),
+    ):
+      DiscoveryEngineSearchTool(
+          data_store_id=(
+              "projects/test/locations/us/collections/default_collection/"
+              "dataStores/test_data_store"
+          ),
+          location="eu",
+      )
+
+    mock_search_client.assert_not_called()
+
+  @mock.patch.object(discoveryengine, "SearchServiceClient")
+  def test_init_with_empty_location_raises_error(self, mock_search_client):
+    """Test initialization rejects an empty location override."""
+    with pytest.raises(
+        ValueError, match="location must not be empty if specified."
+    ):
+      DiscoveryEngineSearchTool(
+          data_store_id=(
+              "projects/test/locations/us/collections/default_collection/"
+              "dataStores/test_data_store"
+          ),
+          location=" ",
+      )
+
+    mock_search_client.assert_not_called()
+
   @mock.patch.object(discovery_engine_search_tool, "client_options")
   @mock.patch.object(discoveryengine, "SearchServiceClient")
   def test_init_with_global_location_keeps_default_endpoint(