github-featured/google--adk-python
2
0
Fork 0
Code Issues Pull Requests Actions 7 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat: add location parameter to list_accessible_data_agents in data_agent toolset

Browse Source 
Allow callers to explicitly specify the Google Cloud location when listing data agents, following standard three-level precedence (explicit argument, toolset config location, falling back to global).

PiperOrigin-RevId: 968067492
This commit is contained in:
Google Team Member
2026-08-20 14:48:17 -07:00
committed by Copybara-Service
parent d42c634bd6
commit 393ec0858b
3 changed files with 108 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/tools/data_agent/data_agent_tool.py
+33 -17
View File
@@ -269,13 +269,17 @@ async def _mutate_data_agent(
kwargs = {}
loc = location or (
settings.location
if settings and isinstance(getattr(settings, "location", None), str)
if settings and isinstance(settings.location, str)
else None
)
if loc:
kwargs["location"] = loc
api_endpoint = getattr(settings, "api_endpoint", None)
if isinstance(api_endpoint, str):
api_endpoint = (
settings.api_endpoint
if settings and isinstance(settings.api_endpoint, str)
else None
)
if api_endpoint:
kwargs["api_endpoint"] = api_endpoint
session, endpoint = _gda_stream_util.get_gda_session(credentials, **kwargs)
base_url = f"{endpoint}/v1"
@@ -326,12 +330,17 @@ def list_accessible_data_agents(
project_id: str,
credentials: Credentials,
settings: DataAgentToolConfig | None = None,
*,
location: str | None = None,
) -> dict[str, Any]:
"""Lists accessible data agents in a project.
Args:
project_id: The project to list agents in.
credentials: The credentials to use for the request.
location: Optional Google Cloud location to list agents from (e.g. "eu" or
"us"). If omitted, uses the toolset's configured location, falling back
to "global".
settings: Optional tool settings containing location or custom endpoint.
Returns:
@@ -355,7 +364,7 @@ def list_accessible_data_agents(
"updateTime": "2025-10-01T22:44:23.094541325Z",
"dataAnalyticsAgent": {
"publishedContext": {
"datasourceReferences": [{
"datasourceReferences": {
"bq": {
"tableReferences": [{
"projectId": "my-project",
@@ -363,7 +372,7 @@ def list_accessible_data_agents(
"tableId": "table1"
}]
}
}]
}
}
}
},
@@ -375,7 +384,7 @@ def list_accessible_data_agents(
"updateTime": "2025-06-23T20:23:49.437095391Z",
"dataAnalyticsAgent": {
"publishedContext": {
"datasourceReferences": [{
"datasourceReferences": {
"bq": {
"tableReferences": [{
"projectId": "another-project",
@@ -383,7 +392,7 @@ def list_accessible_data_agents(
"tableId": "table2"
}]
}
}],
},
"systemInstruction": "You are a helpful assistant.",
"options": {"analysis": {"python": {"enabled": True}}}
}
@@ -393,11 +402,20 @@ def list_accessible_data_agents(
}
"""
try:
location = (
config_location = (
settings.location
if settings and isinstance(settings.location, str)
else None
)
effective_location = location or config_location or "global"
for val, name in (
(project_id, "project_id"),
(effective_location, "location"),
):
invalid_segment_error = _validate_path_segment(val, name)
if invalid_segment_error:
return invalid_segment_error
api_endpoint = (
settings.api_endpoint
if settings and isinstance(settings.api_endpoint, str)
@@ -405,15 +423,15 @@ def list_accessible_data_agents(
)
kwargs: dict[str, str] = {}
if location:
kwargs["location"] = location
if effective_location:
kwargs["location"] = effective_location
if api_endpoint:
kwargs["api_endpoint"] = api_endpoint
session, endpoint = _gda_stream_util.get_gda_session(credentials, **kwargs)
base_url = f"{endpoint}/v1"
target_location = location or "global"
list_url = f"{base_url}/projects/{project_id}/locations/{target_location}/dataAgents:listAccessible"
list_url = f"{base_url}/projects/{project_id}/locations/{effective_location}/dataAgents:listAccessible"
with session:
resp = session.get(
list_url,
@@ -445,14 +463,12 @@ def _get_data_agent_info(
extracted_location = _extract_location_from_resource_name(data_agent_name)
location = extracted_location or (
real_settings.location
if real_settings
and isinstance(getattr(real_settings, "location", None), str)
if real_settings and isinstance(real_settings.location, str)
else None
)
api_endpoint = (
real_settings.api_endpoint
if real_settings
and isinstance(getattr(real_settings, "api_endpoint", None), str)
if real_settings and isinstance(real_settings.api_endpoint, str)
else None
)
@@ -785,7 +801,7 @@ async def create_data_agent(
config_location = (
settings.location
if settings and isinstance(getattr(settings, "location", None), str)
if settings and isinstance(settings.location, str)
else None
)
effective_location = location or config_location or "global"
tests/unittests/tools/data_agent/test_data_agent_tool.py
+54 -4
View File
@@ -42,7 +42,7 @@ def test_list_accessible_data_agents_success(mock_get_session):
)
assert result["status"] == "SUCCESS"
assert result["response"] == ["agent1", "agent2"]
mock_get_session.assert_called_once_with(mock_creds)
mock_get_session.assert_called_once_with(mock_creds, location="global")
mock_session.get.assert_called_once_with(
"https://geminidataanalytics.googleapis.com/v1/projects/test-project/locations/global/dataAgents:listAccessible",
headers={
@@ -70,7 +70,7 @@ def test_list_accessible_data_agents_exception(mock_get_session):
)
assert result["status"] == "ERROR"
assert "List failed!" in result["error_details"]
mock_get_session.assert_called_once_with(mock_creds)
mock_get_session.assert_called_once_with(mock_creds, location="global")
mock_session.get.assert_called_once()
@@ -283,8 +283,6 @@ def test_get_data_agent_info_auto_extract_location(
)
def test_list_accessible_data_agents_regional(mock_get_session):
"""Tests list_accessible_data_agents with regional settings."""
from google.adk.tools.data_agent.config import DataAgentToolConfig
mock_creds = mock.Mock()
mock_session = mock.MagicMock()
mock_response = mock.Mock()
@@ -312,6 +310,58 @@ def test_list_accessible_data_agents_regional(mock_get_session):
)
@mock.patch.object(
data_agent_tool._gda_stream_util, "get_gda_session", autospec=True
)
def test_list_accessible_data_agents_explicit_location(mock_get_session):
"""Tests list_accessible_data_agents with explicit location parameter overriding settings."""
mock_creds = mock.Mock()
mock_session = mock.MagicMock()
mock_response = mock.Mock()
mock_response.json.return_value = {"dataAgents": ["agent_us"]}
mock_response.raise_for_status.return_value = None
mock_session.get.return_value = mock_response
mock_get_session.return_value = (
mock_session,
"https://geminidataanalytics.us.rep.googleapis.com",
)
settings = DataAgentToolConfig(location="eu")
result = data_agent_tool.list_accessible_data_agents(
"test-project", mock_creds, location="us", settings=settings
)
assert result["status"] == "SUCCESS"
assert result["response"] == ["agent_us"]
mock_get_session.assert_called_once_with(mock_creds, location="us")
mock_session.get.assert_called_once_with(
"https://geminidataanalytics.us.rep.googleapis.com/v1/projects/test-project/locations/us/dataAgents:listAccessible",
headers={
"Content-Type": "application/json",
"X-Goog-API-Client": "GOOGLE_ADK",
},
timeout=mock.ANY,
)
def test_list_accessible_data_agents_invalid_location():
"""Tests list_accessible_data_agents with invalid location segment."""
mock_creds = mock.Mock()
result = data_agent_tool.list_accessible_data_agents(
"test-project", mock_creds, location="invalid/segment"
)
assert result["status"] == "ERROR"
assert "Invalid location format" in result["error_details"]
def test_list_accessible_data_agents_invalid_project_id():
"""Tests list_accessible_data_agents with invalid project_id segment."""
mock_creds = mock.Mock()
result = data_agent_tool.list_accessible_data_agents(
"invalid/project", mock_creds
)
assert result["status"] == "ERROR"
assert "Invalid project_id format" in result["error_details"]
class _FakeClock:
"""Virtual clock: only asyncio.sleep advances time, so tests run instantly."""
tests/unittests/tools/data_agent/test_data_agent_toolset.py
+21
View File
@@ -153,3 +153,24 @@ async def test_data_agent_toolset_unknown_tool(selected_tools, returned_tools):
expected_tool_names = set(returned_tools)
actual_tool_names = {tool.name for tool in tools}
assert actual_tool_names == expected_tool_names
@pytest.mark.asyncio
async def test_data_agent_toolset_tools_selective_modification_disabled():
"""Tests that modification tools are excluded when modification is disabled even if in tool_filter."""
credentials_config = DataAgentCredentialsConfig(
client_id="abc", client_secret="def"
)
tool_config = DataAgentToolConfig(enable_data_agent_modification=False)
toolset = DataAgentToolset(
credentials_config=credentials_config,
data_agent_tool_config=tool_config,
tool_filter=[
"create_data_agent",
"update_data_agent",
"delete_data_agent",
],
)
tools = await toolset.get_tools()
assert tools is not None
assert len(tools) == 0