Skip to content

Commit 47b88d2

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Add the ask_data_insights tool for natural language queries on BigQuery data
PiperOrigin-RevId: 799267061
1 parent 6806dea commit 47b88d2

File tree

5 files changed

+34
-7
lines changed

5 files changed

+34
-7
lines changed

contributing/samples/bigquery/README.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,16 @@ distributed via the `google.adk.tools.bigquery` module. These tools include:
2525

2626
Runs a SQL query in BigQuery.
2727

28+
1. `ask_data_insights`
29+
30+
Natural language-in, natural language-out tool that answers questions
31+
about structured data in BigQuery. Provides a one-stop solution for generating
32+
insights from data.
33+
34+
**Note**: This tool requires additional setup in your project. Please refer to
35+
the official [Conversational Analytics API documentation](https://cloud.google.com/gemini/docs/conversational-analytics-api/overview)
36+
for instructions.
37+
2838
## How to use
2939

3040
Set up environment variables in your `.env` file for using

src/google/adk/tools/bigquery/bigquery_toolset.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from google.adk.agents.readonly_context import ReadonlyContext
2222
from typing_extensions import override
2323

24+
from . import data_insights_tool
2425
from . import metadata_tool
2526
from . import query_tool
2627
from ...tools.base_tool import BaseTool
@@ -80,6 +81,7 @@ async def get_tools(
8081
metadata_tool.list_dataset_ids,
8182
metadata_tool.list_table_ids,
8283
query_tool.get_execute_sql(self._tool_settings),
84+
data_insights_tool.ask_data_insights,
8385
]
8486
]
8587

src/google/adk/tools/bigquery/data_insights_tool.py

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,14 +34,16 @@ def ask_data_insights(
3434
) -> Dict[str, Any]:
3535
"""Answers questions about structured data in BigQuery tables using natural language.
3636
37-
This function takes auser's question (which can include conversational
38-
history for context) andreferences to specific BigQuery tables, and sends
37+
This function takes a user's question (which can include conversational
38+
history for context) and references to specific BigQuery tables, and sends
3939
them to a stateless conversational API.
4040
4141
The API uses a GenAI agent to understand the question, generate and execute
4242
SQL queries and Python code, and formulate an answer. This function returns a
4343
detailed, sequential log of this entire process, which includes any generated
44-
SQL or Python code, the data retrieved, and the final text answer.
44+
SQL or Python code, the data retrieved, and the final text answer. The final
45+
answer is always in plain text, as the underlying API is instructed not to
46+
generate any charts, graphs, images, or other visualizations.
4547
4648
Use this tool to perform data analysis, get insights, or answer complex
4749
questions about the contents of specific BigQuery tables.
@@ -123,9 +125,22 @@ def ask_data_insights(
123125
}
124126
ca_url = f"https://geminidataanalytics.googleapis.com/v1alpha/projects/{project_id}/locations/{location}:chat"
125127

128+
instructions = """**INSTRUCTIONS - FOLLOW THESE RULES:**
129+
1. **CONTENT:** Your answer should present the supporting data and then provide a conclusion based on that data.
130+
2. **OUTPUT FORMAT:** Your entire response MUST be in plain text format ONLY.
131+
3. **NO CHARTS:** You are STRICTLY FORBIDDEN from generating any charts, graphs, images, or any other form of visualization.
132+
"""
133+
134+
final_query_text = f"""
135+
{instructions}
136+
137+
**User Query and Context:**
138+
{user_query_with_context}
139+
"""
140+
126141
ca_payload = {
127142
"project": f"projects/{project_id}",
128-
"messages": [{"userMessage": {"text": user_query_with_context}}],
143+
"messages": [{"userMessage": {"text": final_query_text}}],
129144
"inlineContext": {
130145
"datasourceReferences": {
131146
"bq": {"tableReferences": table_references}
@@ -289,7 +304,7 @@ def _handle_data_response(
289304
schema = resp["result"]["schema"]
290305
headers = [field.get("name") for field in schema.get("fields", [])]
291306

292-
all_rows = resp["result"]["data"]
307+
all_rows = resp["result"].get("data", [])
293308
total_rows = len(all_rows)
294309

295310
compact_rows = []

tests/unittests/tools/bigquery/test_bigquery_client.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
from google.adk.tools.bigquery.client import get_bigquery_client
2222
from google.auth.exceptions import DefaultCredentialsError
2323
from google.oauth2.credentials import Credentials
24-
import pytest
2524

2625

2726
def test_bigquery_client_project():

tests/unittests/tools/bigquery/test_bigquery_toolset.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ async def test_bigquery_toolset_tools_default():
4141
tools = await toolset.get_tools()
4242
assert tools is not None
4343

44-
assert len(tools) == 5
44+
assert len(tools) == 6
4545
assert all([isinstance(tool, GoogleTool) for tool in tools])
4646

4747
expected_tool_names = set([
@@ -50,6 +50,7 @@ async def test_bigquery_toolset_tools_default():
5050
"list_table_ids",
5151
"get_table_info",
5252
"execute_sql",
53+
"ask_data_insights",
5354
])
5455
actual_tool_names = set([tool.name for tool in tools])
5556
assert actual_tool_names == expected_tool_names

0 commit comments

Comments
 (0)