encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: add chat first-party tool

Browse Source 
This tool answers questions about structured data in BigQuery using natural language.

PiperOrigin-RevId: 789000987
This commit is contained in:
Google Team Member
2025-07-30 12:35:20 -07:00
committed by Copybara-Service
parent 0c6086cb15
commit 7c9b0a2567
12 changed files with 1008 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unittests/tools/bigquery/test_bigquery_client.py
+1
View File
@@ -1,3 +1,4 @@
# 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.
tests/unittests/tools/bigquery/test_bigquery_data_insights_tool.py
+273
View File
@@ -0,0 +1,273 @@
# 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 pathlib
from unittest import mock
from google.adk.tools.bigquery import data_insights_tool
import pytest
import yaml
@pytest.mark.parametrize(
"case_file_path",
[
pytest.param("test_data/ask_data_insights_penguins_highest_mass.yaml"),
],
)
@mock.patch(
"google.adk.tools.bigquery.data_insights_tool.requests.Session.post"
)
def test_ask_data_insights_pipeline_from_file(mock_post, case_file_path):
"""Runs a full integration test for the ask_data_insights pipeline using data from a specific file."""
# 1. Construct the full, absolute path to the data file
full_path = pathlib.Path(__file__).parent / case_file_path
# 2. Load the test case data from the specified YAML file
with open(full_path, "r", encoding="utf-8") as f:
case_data = yaml.safe_load(f)
# 3. Prepare the mock stream and expected output from the loaded data
mock_stream_str = case_data["mock_api_stream"]
fake_stream_lines = [
line.encode("utf-8") for line in mock_stream_str.splitlines()
]
# Load the expected output as a list of dictionaries, not a single string
expected_final_list = case_data["expected_output"]
# 4. Configure the mock for requests.post
mock_response = mock.Mock()
mock_response.iter_lines.return_value = fake_stream_lines
# Add raise_for_status mock which is called in the updated code
mock_response.raise_for_status.return_value = None
mock_post.return_value.__enter__.return_value = mock_response
# 5. Call the function under test
result = data_insights_tool._get_stream( # pylint: disable=protected-access
url="fake_url",
ca_payload={},
headers={},
max_query_result_rows=50,
)
# 6. Assert that the final list of dicts matches the expected output
assert result == expected_final_list
@mock.patch("google.adk.tools.bigquery.data_insights_tool._get_stream")
def test_ask_data_insights_success(mock_get_stream):
"""Tests the success path of ask_data_insights using decorators."""
# 1. Configure the behavior of the mocked functions
mock_get_stream.return_value = "Final formatted string from stream"
# 2. Create mock inputs for the function call
mock_creds = mock.Mock()
mock_creds.token = "fake-token"
mock_config = mock.Mock()
mock_config.max_query_result_rows = 100
# 3. Call the function under test
result = data_insights_tool.ask_data_insights(
project_id="test-project",
user_query_with_context="test query",
table_references=[],
credentials=mock_creds,
config=mock_config,
)
# 4. Assert the results are as expected
assert result["status"] == "SUCCESS"
assert result["response"] == "Final formatted string from stream"
mock_get_stream.assert_called_once()
@mock.patch("google.adk.tools.bigquery.data_insights_tool._get_stream")
def test_ask_data_insights_handles_exception(mock_get_stream):
"""Tests the exception path of ask_data_insights using decorators."""
# 1. Configure one of the mocks to raise an error
mock_get_stream.side_effect = Exception("API call failed!")
# 2. Create mock inputs
mock_creds = mock.Mock()
mock_creds.token = "fake-token"
mock_config = mock.Mock()
# 3. Call the function
result = data_insights_tool.ask_data_insights(
project_id="test-project",
user_query_with_context="test query",
table_references=[],
credentials=mock_creds,
config=mock_config,
)
# 4. Assert that the error was caught and formatted correctly
assert result["status"] == "ERROR"
assert "API call failed!" in result["error_details"]
mock_get_stream.assert_called_once()
@pytest.mark.parametrize(
"initial_messages, new_message, expected_list",
[
pytest.param(
[{"Thinking": None}, {"Schema Resolved": {}}],
{"SQL Generated": "SELECT 1"},
[
{"Thinking": None},
{"Schema Resolved": {}},
{"SQL Generated": "SELECT 1"},
],
id="append_when_last_message_is_not_data",
),
pytest.param(
[{"Thinking": None}, {"Data Retrieved": {"rows": [1]}}],
{"Data Retrieved": {"rows": [1, 2]}},
[{"Thinking": None}, {"Data Retrieved": {"rows": [1, 2]}}],
id="replace_when_last_message_is_data",
),
pytest.param(
[],
{"Answer": "First Message"},
[{"Answer": "First Message"}],
id="append_to_an_empty_list",
),
pytest.param(
[{"Data Retrieved": {}}],
{},
[{"Data Retrieved": {}}],
id="should_not_append_an_empty_new_message",
),
],
)
def test_append_message(initial_messages, new_message, expected_list):
"""Tests the logic of replacing the last message if it's a data message."""
messages_copy = initial_messages.copy()
data_insights_tool._append_message(messages_copy, new_message) # pylint: disable=protected-access
assert messages_copy == expected_list
@pytest.mark.parametrize(
"response_dict, expected_output",
[
pytest.param(
{"parts": ["The answer", " is 42."]},
{"Answer": "The answer is 42."},
id="multiple_parts",
),
pytest.param(
{"parts": ["Hello"]}, {"Answer": "Hello"}, id="single_part"
),
pytest.param({}, {"Answer": ""}, id="empty_response"),
],
)
def test_handle_text_response(response_dict, expected_output):
"""Tests the text response handler."""
result = data_insights_tool._handle_text_response(response_dict) # pylint: disable=protected-access
assert result == expected_output
@pytest.mark.parametrize(
"response_dict, expected_output",
[
pytest.param(
{"query": {"question": "What is the schema?"}},
{"Question": "What is the schema?"},
id="schema_query_path",
),
pytest.param(
{
"result": {
"datasources": [{
"bigqueryTableReference": {
"projectId": "p",
"datasetId": "d",
"tableId": "t",
},
"schema": {
"fields": [{"name": "col1", "type": "STRING"}]
},
}]
}
},
{
"Schema Resolved": [{
"source_name": "p.d.t",
"schema": {
"headers": ["Column", "Type", "Description", "Mode"],
"rows": [["col1", "STRING", "", ""]],
},
}]
},
id="schema_result_path",
),
],
)
def test_handle_schema_response(response_dict, expected_output):
"""Tests different paths of the schema response handler."""
result = data_insights_tool._handle_schema_response(response_dict) # pylint: disable=protected-access
assert result == expected_output
@pytest.mark.parametrize(
"response_dict, expected_output",
[
pytest.param(
{"generatedSql": "SELECT 1;"},
{"SQL Generated": "SELECT 1;"},
id="format_generated_sql",
),
pytest.param(
{
"result": {
"schema": {"fields": [{"name": "id"}, {"name": "name"}]},
"data": [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}],
}
},
{
"Data Retrieved": {
"headers": ["id", "name"],
"rows": [[1, "A"], [2, "B"]],
"summary": "Showing all 2 rows.",
}
},
id="format_data_result_table",
),
],
)
def test_handle_data_response(response_dict, expected_output):
"""Tests different paths of the data response handler, including truncation."""
result = data_insights_tool._handle_data_response(response_dict, 100) # pylint: disable=protected-access
assert result == expected_output
@pytest.mark.parametrize(
"response_dict, expected_output",
[
pytest.param(
{"code": 404, "message": "Not Found"},
{"Error": {"Code": 404, "Message": "Not Found"}},
id="full_error_message",
),
pytest.param(
{"code": 500},
{"Error": {"Code": 500, "Message": "No message provided."}},
id="error_with_missing_message",
),
],
)
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
tests/unittests/tools/bigquery/test_bigquery_tool.py
+33
View File
@@ -19,6 +19,7 @@ from unittest.mock import patch
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsManager
from google.adk.tools.bigquery.bigquery_tool import BigQueryTool
from google.adk.tools.bigquery.config import BigQueryToolConfig
from google.adk.tools.tool_context import ToolContext
# Mock the Google OAuth and API dependencies
from google.oauth2.credentials import Credentials
@@ -267,3 +268,35 @@ class TestBigQueryTool:
assert "required_param" in mandatory_args
assert "credentials" not in mandatory_args
assert "optional_param" not in mandatory_args
@pytest.mark.parametrize(
"input_config, expected_config",
[
pytest.param(
BigQueryToolConfig(
write_mode="blocked", max_query_result_rows=50
),
BigQueryToolConfig(
write_mode="blocked", max_query_result_rows=50
),
id="with_provided_config",
),
pytest.param(
None,
BigQueryToolConfig(),
id="with_none_config_creates_default",
),
],
)
def test_tool_config_initialization(self, input_config, expected_config):
"""Tests that self._tool_config is correctly initialized by comparing its
final state to an expected configuration object.
"""
# 1. Initialize the tool with the parameterized config
tool = BigQueryTool(func=None, bigquery_tool_config=input_config)
# 2. Assert that the tool's config has the same attribute values
# as the expected config. Comparing the __dict__ is a robust
# way to check for value equality.
assert tool._tool_config.__dict__ == expected_config.__dict__ # pylint: disable=protected-access
tests/unittests/tools/bigquery/test_bigquery_toolset.py
+2 -1
View File
@@ -34,7 +34,7 @@ async def test_bigquery_toolset_tools_default():
tools = await toolset.get_tools()
assert tools is not None
assert len(tools) == 5
assert len(tools) == 6
assert all([isinstance(tool, BigQueryTool) for tool in tools])
expected_tool_names = set([
@@ -43,6 +43,7 @@ async def test_bigquery_toolset_tools_default():
"list_table_ids",
"get_table_info",
"execute_sql",
"ask_data_insights",
])
actual_tool_names = set([tool.name for tool in tools])
assert actual_tool_names == expected_tool_names
tests/unittests/tools/bigquery/test_data/ask_data_insights_penguins_highest_mass.yaml
+336
View File
@@ -0,0 +1,336 @@
description: "Tests a full, realistic stream about finding the penguin island with the highest body mass."
user_question: "Penguins on which island have the highest average body mass?"
mock_api_stream: |
[{
"timestamp": "2025-07-17T17:25:28.231Z",
"systemMessage": {
"schema": {
"query": {
"question": "Penguins on which island have the highest average body mass?"
}
}
}
}
,
{
"timestamp": "2025-07-17T17:25:29.406Z",
"systemMessage": {
"schema": {
"result": {
"datasources": [
{
"bigqueryTableReference": {
"projectId": "bigframes-dev-perf",
"datasetId": "bigframes_testing_eu",
"tableId": "penguins"
},
"schema": {
"fields": [
{
"name": "species",
"type": "STRING",
"mode": "NULLABLE"
},
{
"name": "island",
"type": "STRING",
"mode": "NULLABLE"
},
{
"name": "culmen_length_mm",
"type": "FLOAT64",
"mode": "NULLABLE"
},
{
"name": "culmen_depth_mm",
"type": "FLOAT64",
"mode": "NULLABLE"
},
{
"name": "flipper_length_mm",
"type": "FLOAT64",
"mode": "NULLABLE"
},
{
"name": "body_mass_g",
"type": "FLOAT64",
"mode": "NULLABLE"
},
{
"name": "sex",
"type": "STRING",
"mode": "NULLABLE"
}
]
}
}
]
}
}
}
}
,
{
"timestamp": "2025-07-17T17:25:30.431Z",
"systemMessage": {
"data": {
"query": {
"question": "What is the average body mass for each island?",
"datasources": [
{
"bigqueryTableReference": {
"projectId": "bigframes-dev-perf",
"datasetId": "bigframes_testing_eu",
"tableId": "penguins"
},
"schema": {
"fields": [
{
"name": "species",
"type": "STRING",
"mode": "NULLABLE"
},
{
"name": "island",
"type": "STRING",
"mode": "NULLABLE"
},
{
"name": "culmen_length_mm",
"type": "FLOAT64",
"mode": "NULLABLE"
},
{
"name": "culmen_depth_mm",
"type": "FLOAT64",
"mode": "NULLABLE"
},
{
"name": "flipper_length_mm",
"type": "FLOAT64",
"mode": "NULLABLE"
},
{
"name": "body_mass_g",
"type": "FLOAT64",
"mode": "NULLABLE"
},
{
"name": "sex",
"type": "STRING",
"mode": "NULLABLE"
}
]
}
}
],
"name": "average_body_mass_by_island"
}
}
}
}
,
{
"timestamp": "2025-07-17T17:25:31.171Z",
"systemMessage": {
"data": {
"generatedSql": "SELECT island, AVG(body_mass_g) AS average_body_mass\nFROM `bigframes-dev-perf`.`bigframes_testing_eu`.`penguins`\nGROUP BY island;"
}
}
}
,
{
"timestamp": "2025-07-17T17:25:32.378Z",
"systemMessage": {
"data": {
"bigQueryJob": {
"projectId": "bigframes-dev-perf",
"jobId": "job_S4PGRwxO78_FrVmCHW_sklpeZFps",
"destinationTable": {
"projectId": "bigframes-dev-perf",
"datasetId": "_376b2bd1b83171a540d39ff3d58f39752e2724c9",
"tableId": "anonev_4a9PK1uHzAHwAOpSNOxMVhpUppM2sllR68riN6t41kM"
},
"location": "EU",
"schema": {
"fields": [
{
"name": "island",
"type": "STRING",
"mode": "NULLABLE"
},
{
"name": "average_body_mass",
"type": "FLOAT",
"mode": "NULLABLE"
}
]
}
}
}
}
}
,
{
"timestamp": "2025-07-17T17:25:32.664Z",
"systemMessage": {
"data": {
"result": {
"data": [
{
"island": "Biscoe",
"average_body_mass": "4716.017964071853"
},
{
"island": "Dream",
"average_body_mass": "3712.9032258064512"
},
{
"island": "Torgersen",
"average_body_mass": "3706.3725490196075"
}
],
"name": "average_body_mass_by_island",
"schema": {
"fields": [
{
"name": "island",
"type": "STRING",
"mode": "NULLABLE"
},
{
"name": "average_body_mass",
"type": "FLOAT",
"mode": "NULLABLE"
}
]
}
}
}
}
}
,
{
"timestamp": "2025-07-17T17:25:33.808Z",
"systemMessage": {
"chart": {
"query": {
"instructions": "Create a bar chart showing the average body mass for each island. The island should be on the x axis and the average body mass should be on the y axis.",
"dataResultName": "average_body_mass_by_island"
}
}
}
}
,
{
"timestamp": "2025-07-17T17:25:38.999Z",
"systemMessage": {
"chart": {
"result": {
"vegaConfig": {
"mark": {
"type": "bar",
"tooltip": true
},
"encoding": {
"x": {
"field": "island",
"type": "nominal",
"title": "Island",
"axis": {
"labelOverlap": true
},
"sort": {}
},
"y": {
"field": "average_body_mass",
"type": "quantitative",
"title": "Average Body Mass",
"axis": {
"labelOverlap": true
},
"sort": {}
}
},
"title": "Average Body Mass for Each Island",
"data": {
"values": [
{
"island": "Biscoe",
"average_body_mass": 4716.0179640718534
},
{
"island": "Dream",
"average_body_mass": 3712.9032258064512
},
{
"island": "Torgersen",
"average_body_mass": 3706.3725490196075
}
]
}
},
"image": {}
}
}
}
}
,
{
"timestamp": "2025-07-17T17:25:40.018Z",
"systemMessage": {
"text": {
"parts": [
"Penguins on Biscoe island have the highest average body mass, with an average of 4716.02g."
]
}
}
}
]
expected_output:
- Question: Penguins on which island have the highest average body mass?
- Schema Resolved:
- source_name: bigframes-dev-perf.bigframes_testing_eu.penguins
schema:
headers:
- Column
- Type
- Description
- Mode
rows:
- - species
- STRING
- ''
- NULLABLE
- - island
- STRING
- ''
- NULLABLE
- - culmen_length_mm
- FLOAT64
- ''
- NULLABLE
- - culmen_depth_mm
- FLOAT64
- ''
- NULLABLE
- - flipper_length_mm
- FLOAT64
- ''
- NULLABLE
- - body_mass_g
- FLOAT64
- ''
- NULLABLE
- - sex
- STRING
- ''
- NULLABLE
- Retrieval Query:
Query Name: average_body_mass_by_island
Question: What is the average body mass for each island?
- SQL Generated: "SELECT island, AVG(body_mass_g) AS average_body_mass\nFROM `bigframes-dev-perf`.`bigframes_testing_eu`.`penguins`\nGROUP BY island;"
- Answer: Penguins on Biscoe island have the highest average body mass, with an average of 4716.02g.