mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: add new conversational analytics api tool set
PiperOrigin-RevId: 853489874
This commit is contained in:
Google Team Member
committed by
Copybara-Service
Copybara-Service
parent
aaf76a6a51
commit
c34feb4c0e
@@ -1,29 +0,0 @@
|
||||
# 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.
|
||||
|
||||
import sys
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
MOCK_GEMINI_DATA_ANALYTICS = mock.MagicMock()
|
||||
sys.modules["google.cloud.geminidataanalytics"] = MOCK_GEMINI_DATA_ANALYTICS
|
||||
|
||||
# The mock.patch calls require 'geminidataanalytics' to be an attribute of
|
||||
# 'google.cloud' module for patching by string to work.
|
||||
try:
|
||||
import google.cloud
|
||||
except ImportError:
|
||||
sys.modules["google.cloud"] = types.ModuleType("google.cloud")
|
||||
finally:
|
||||
sys.modules["google.cloud"].geminidataanalytics = MOCK_GEMINI_DATA_ANALYTICS
|
||||
@@ -1,130 +0,0 @@
|
||||
# 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 unittest import mock
|
||||
|
||||
from google.adk.tools.bigquery import BigQueryCredentialsConfig
|
||||
from google.adk.tools.bigquery import BigQueryDataAgentToolset
|
||||
from google.adk.tools.bigquery.config import BigQueryToolConfig
|
||||
from google.adk.tools.google_tool import GoogleTool
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bigquery_data_agent_toolset_tools_default():
|
||||
"""Test default BigQueryDataAgentToolset.
|
||||
|
||||
This test verifies the behavior of the BigQueryDataAgentToolset when no filter is
|
||||
specified.
|
||||
"""
|
||||
credentials_config = BigQueryCredentialsConfig(
|
||||
client_id="abc", client_secret="def"
|
||||
)
|
||||
toolset = BigQueryDataAgentToolset(
|
||||
credentials_config=credentials_config, bigquery_tool_config=None
|
||||
)
|
||||
# Verify that the tool config is initialized to default values.
|
||||
assert isinstance(toolset._tool_settings, BigQueryToolConfig) # pylint: disable=protected-access
|
||||
assert toolset._tool_settings.__dict__ == BigQueryToolConfig().__dict__ # pylint: disable=protected-access
|
||||
|
||||
tools = await toolset.get_tools()
|
||||
assert tools is not None
|
||||
|
||||
assert len(tools) == 3
|
||||
assert all(isinstance(tool, GoogleTool) for tool in tools)
|
||||
|
||||
expected_tool_names = set([
|
||||
"list_accessible_data_agents",
|
||||
"get_data_agent_info",
|
||||
"ask_data_agent",
|
||||
])
|
||||
actual_tool_names = {tool.name for tool in tools}
|
||||
assert actual_tool_names == expected_tool_names
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"selected_tools",
|
||||
[
|
||||
pytest.param([], id="None"),
|
||||
pytest.param(
|
||||
["list_accessible_data_agents", "get_data_agent_info"],
|
||||
id="list_and_get",
|
||||
),
|
||||
pytest.param(["ask_data_agent"], id="ask"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_bigquery_data_agent_toolset_tools_selective(selected_tools):
|
||||
"""Test BigQueryDataAgentToolset with filter.
|
||||
|
||||
This test verifies the behavior of the BigQueryDataAgentToolset when filter is
|
||||
specified. A use case for this would be when the agent builder wants to
|
||||
use only a subset of the tools provided by the toolset.
|
||||
"""
|
||||
credentials_config = BigQueryCredentialsConfig(
|
||||
client_id="abc", client_secret="def"
|
||||
)
|
||||
toolset = BigQueryDataAgentToolset(
|
||||
credentials_config=credentials_config, tool_filter=selected_tools
|
||||
)
|
||||
tools = await toolset.get_tools()
|
||||
assert tools is not None
|
||||
|
||||
assert len(tools) == len(selected_tools)
|
||||
assert all(isinstance(tool, GoogleTool) for tool in tools)
|
||||
|
||||
expected_tool_names = set(selected_tools)
|
||||
actual_tool_names = {tool.name for tool in tools}
|
||||
assert actual_tool_names == expected_tool_names
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("selected_tools", "returned_tools"),
|
||||
[
|
||||
pytest.param(["unknown"], [], id="all-unknown"),
|
||||
pytest.param(
|
||||
["unknown", "ask_data_agent"],
|
||||
["ask_data_agent"],
|
||||
id="mixed-known-unknown",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_bigquery_data_agent_toolset_unknown_tool(
|
||||
selected_tools, returned_tools
|
||||
):
|
||||
"""Test BigQueryDataAgentToolset with filter.
|
||||
|
||||
This test verifies the behavior of the BigQueryDataAgentToolset when filter is
|
||||
specified with an unknown tool.
|
||||
"""
|
||||
credentials_config = BigQueryCredentialsConfig(
|
||||
client_id="abc", client_secret="def"
|
||||
)
|
||||
|
||||
toolset = BigQueryDataAgentToolset(
|
||||
credentials_config=credentials_config, tool_filter=selected_tools
|
||||
)
|
||||
|
||||
tools = await toolset.get_tools()
|
||||
assert tools is not None
|
||||
|
||||
assert len(tools) == len(returned_tools)
|
||||
assert all(isinstance(tool, GoogleTool) for tool in tools)
|
||||
|
||||
expected_tool_names = set(returned_tools)
|
||||
actual_tool_names = {tool.name for tool in tools}
|
||||
assert actual_tool_names == expected_tool_names
|
||||
@@ -16,7 +16,6 @@ import pathlib
|
||||
from unittest import mock
|
||||
|
||||
from google.adk.tools.bigquery import data_insights_tool
|
||||
from google.adk.tools.tool_context import ToolContext
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
@@ -270,134 +269,3 @@ def test_handle_error(response_dict, expected_output):
|
||||
"""Tests the error response handler."""
|
||||
result = data_insights_tool._handle_error(response_dict) # pylint: disable=protected-access
|
||||
assert result == expected_output
|
||||
|
||||
|
||||
@mock.patch.object(
|
||||
data_insights_tool.geminidataanalytics, "DataAgentServiceClient"
|
||||
)
|
||||
def test_list_accessible_data_agents_success(mock_data_agent_client):
|
||||
"""Tests list_accessible_data_agents success path."""
|
||||
mock_creds = mock.Mock()
|
||||
mock_agent1 = mock.MagicMock()
|
||||
mock_agent1.__str__.return_value = "agent1"
|
||||
mock_agent2 = mock.MagicMock()
|
||||
mock_agent2.__str__.return_value = "agent2"
|
||||
mock_data_agent_client.return_value.list_accessible_data_agents.return_value = [
|
||||
mock_agent1,
|
||||
mock_agent2,
|
||||
]
|
||||
result = data_insights_tool.list_accessible_data_agents(
|
||||
"test-project", mock_creds
|
||||
)
|
||||
assert result["status"] == "SUCCESS"
|
||||
assert result["response"] == ["agent1", "agent2"]
|
||||
mock_data_agent_client.assert_called_once_with(credentials=mock_creds)
|
||||
|
||||
|
||||
@mock.patch.object(
|
||||
data_insights_tool.geminidataanalytics, "DataAgentServiceClient"
|
||||
)
|
||||
def test_list_accessible_data_agents_exception(mock_data_agent_client):
|
||||
"""Tests list_accessible_data_agents exception path."""
|
||||
mock_creds = mock.Mock()
|
||||
mock_data_agent_client.return_value.list_accessible_data_agents.side_effect = Exception(
|
||||
"List failed!"
|
||||
)
|
||||
result = data_insights_tool.list_accessible_data_agents(
|
||||
"test-project", mock_creds
|
||||
)
|
||||
assert result["status"] == "ERROR"
|
||||
assert "List failed!" in result["error_details"]
|
||||
mock_data_agent_client.assert_called_once_with(credentials=mock_creds)
|
||||
|
||||
|
||||
@mock.patch.object(
|
||||
data_insights_tool.geminidataanalytics, "DataAgentServiceClient"
|
||||
)
|
||||
def test_get_data_agent_info_success(mock_data_agent_client):
|
||||
"""Tests get_data_agent_info success path."""
|
||||
mock_creds = mock.Mock()
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.__str__.return_value = "agent_info"
|
||||
mock_data_agent_client.return_value.get_data_agent.return_value = (
|
||||
mock_response
|
||||
)
|
||||
result = data_insights_tool.get_data_agent_info("agent_name", mock_creds)
|
||||
assert result["status"] == "SUCCESS"
|
||||
assert result["response"] == "agent_info"
|
||||
mock_data_agent_client.assert_called_once_with(credentials=mock_creds)
|
||||
|
||||
|
||||
@mock.patch.object(
|
||||
data_insights_tool.geminidataanalytics, "DataAgentServiceClient"
|
||||
)
|
||||
def test_get_data_agent_info_exception(mock_data_agent_client):
|
||||
"""Tests get_data_agent_info exception path."""
|
||||
mock_creds = mock.Mock()
|
||||
mock_data_agent_client.return_value.get_data_agent.side_effect = Exception(
|
||||
"Get failed!"
|
||||
)
|
||||
result = data_insights_tool.get_data_agent_info("agent_name", mock_creds)
|
||||
assert result["status"] == "ERROR"
|
||||
assert "Get failed!" in result["error_details"]
|
||||
mock_data_agent_client.assert_called_once_with(credentials=mock_creds)
|
||||
|
||||
|
||||
@mock.patch.object(
|
||||
data_insights_tool.geminidataanalytics, "DataChatServiceClient"
|
||||
)
|
||||
def test_ask_data_agent_success(mock_data_chat_client):
|
||||
"""Tests ask_data_agent success path."""
|
||||
mock_creds = mock.Mock()
|
||||
mock_invocation_context = mock.Mock()
|
||||
mock_invocation_context.session.state = {}
|
||||
mock_context = ToolContext(mock_invocation_context)
|
||||
mock_response1 = mock.MagicMock()
|
||||
mock_response1.system_message.text.parts = ["response1"]
|
||||
mock_response1.system_message.data.generated_sql = None
|
||||
mock_response1.system_message.data.result = None
|
||||
mock_response1.system_message.error = None
|
||||
mock_response2 = mock.MagicMock()
|
||||
mock_response2.system_message.text.parts = ["response2"]
|
||||
mock_response2.system_message.data.generated_sql = None
|
||||
mock_response2.system_message.data.result = None
|
||||
mock_response2.system_message.error = None
|
||||
mock_data_chat_client.return_value.chat.return_value = [
|
||||
mock_response1,
|
||||
mock_response2,
|
||||
]
|
||||
result = data_insights_tool.ask_data_agent(
|
||||
"projects/p/locations/l/dataAgents/a",
|
||||
"query",
|
||||
credentials=mock_creds,
|
||||
tool_context=mock_context,
|
||||
)
|
||||
assert result["status"] == "SUCCESS"
|
||||
assert result["response"] == [
|
||||
{"Answer": "response1"},
|
||||
{"Answer": "response2"},
|
||||
]
|
||||
mock_data_chat_client.assert_called_once_with(credentials=mock_creds)
|
||||
|
||||
|
||||
@mock.patch.object(
|
||||
data_insights_tool.geminidataanalytics, "DataChatServiceClient"
|
||||
)
|
||||
def test_ask_data_agent_exception(mock_data_chat_client):
|
||||
"""Tests ask_data_agent exception path."""
|
||||
mock_creds = mock.Mock()
|
||||
mock_invocation_context = mock.Mock()
|
||||
mock_invocation_context.session.state = {}
|
||||
mock_context = ToolContext(mock_invocation_context)
|
||||
mock_data_chat_client.return_value.chat.side_effect = Exception(
|
||||
"Chat failed!"
|
||||
)
|
||||
result = data_insights_tool.ask_data_agent(
|
||||
"projects/p/locations/l/dataAgents/a",
|
||||
"query",
|
||||
credentials=mock_creds,
|
||||
tool_context=mock_context,
|
||||
)
|
||||
assert result["status"] == "ERROR"
|
||||
assert "Chat failed!" in result["error_details"]
|
||||
mock_data_chat_client.assert_called_once_with(credentials=mock_creds)
|
||||
|
||||
Reference in New Issue
Block a user