-
Notifications
You must be signed in to change notification settings - Fork 2.4k
Expand file tree
/
Copy pathtools.py
More file actions
126 lines (95 loc) · 3.58 KB
/
tools.py
File metadata and controls
126 lines (95 loc) · 3.58 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
# 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.
"""Tools for the ADK Sampmles Data Science Agent."""
import logging
from google.adk.tools import ToolContext
from google.adk.tools.agent_tool import AgentTool
from .sub_agents import alloydb_agent, analytics_agent, bigquery_agent
logger = logging.getLogger(__name__)
async def call_bigquery_agent(
question: str,
tool_context: ToolContext,
):
"""Tool to call bigquery database (nl2sql) agent."""
logger.debug("call_bigquery_agent: %s", question)
agent_tool = AgentTool(agent=bigquery_agent)
bigquery_agent_output = await agent_tool.run_async(
args={"request": question}, tool_context=tool_context
)
tool_context.state["bigquery_agent_output"] = bigquery_agent_output
return bigquery_agent_output
async def call_alloydb_agent(
question: str,
tool_context: ToolContext,
):
"""Tool to call alloydb database (nl2sql) agent."""
logger.debug("call_alloydb_agent: %s", question)
agent_tool = AgentTool(agent=alloydb_agent)
alloydb_agent_output = await agent_tool.run_async(
args={"request": question}, tool_context=tool_context
)
tool_context.state["alloydb_agent_output"] = alloydb_agent_output
return alloydb_agent_output
async def call_analytics_agent(
question: str,
tool_context: ToolContext,
):
"""
This tool can generate Python code to process and analyze a dataset.
Some of the tasks it can do in Python include:
* Creating graphics for data visualization;
* Processing or filtering existing datasets;
* Combining datasets to create a joined dataset for further analysis.
The Python modules available to it are:
* io
* math
* re
* matplotlib.pyplot
* numpy
* pandas
The tool DOES NOT have the ability to retrieve additional data from
a database. Only the data already retrieved will be analyzed.
Args:
question (str): Natural language question or analytics request.
tool_context (ToolContext): The tool context to use for generating the
SQL query.
Returns:
Response from the analytics agent.
"""
logger.debug("call_analytics_agent: %s", question)
# if question == "N/A":
# return tool_context.state["db_agent_output"]
bigquery_data = ""
alloydb_data = ""
if "bigquery_query_result" in tool_context.state:
bigquery_data = tool_context.state["bigquery_query_result"]
if "alloydb_query_result" in tool_context.state:
alloydb_data = tool_context.state["alloydb_query_result"]
question_with_data = f"""
Question to answer: {question}
Actual data to analyze this question is available in the following data
tables:
<BIGQUERY>
{bigquery_data}
</BIGQUERY>
<ALLOYDB>
{alloydb_data}
</ALLOYDB>
"""
agent_tool = AgentTool(agent=analytics_agent)
analytics_agent_output = await agent_tool.run_async(
args={"request": question_with_data}, tool_context=tool_context
)
tool_context.state["analytics_agent_output"] = analytics_agent_output
return analytics_agent_output