From 393ec0858baa61c5ebb8b733cd29e1ff673ead19 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Thu, 20 Aug 2026 14:48:17 -0700 Subject: [PATCH] feat: add location parameter to list_accessible_data_agents in data_agent toolset 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 --- .../adk/tools/data_agent/data_agent_tool.py | 50 ++++++++++------ .../tools/data_agent/test_data_agent_tool.py | 58 +++++++++++++++++-- .../data_agent/test_data_agent_toolset.py | 21 +++++++ 3 files changed, 108 insertions(+), 21 deletions(-) diff --git a/src/google/adk/tools/data_agent/data_agent_tool.py b/src/google/adk/tools/data_agent/data_agent_tool.py index e60006e3..4bb50c54 100644 --- a/src/google/adk/tools/data_agent/data_agent_tool.py +++ b/src/google/adk/tools/data_agent/data_agent_tool.py @@ -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" diff --git a/tests/unittests/tools/data_agent/test_data_agent_tool.py b/tests/unittests/tools/data_agent/test_data_agent_tool.py index 148051ac..04a1381d 100644 --- a/tests/unittests/tools/data_agent/test_data_agent_tool.py +++ b/tests/unittests/tools/data_agent/test_data_agent_tool.py @@ -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.""" diff --git a/tests/unittests/tools/data_agent/test_data_agent_toolset.py b/tests/unittests/tools/data_agent/test_data_agent_toolset.py index a7306eaa..d9cbb798 100644 --- a/tests/unittests/tools/data_agent/test_data_agent_toolset.py +++ b/tests/unittests/tools/data_agent/test_data_agent_toolset.py @@ -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