update bigquery agent and oauth calendar agent according to new toolset interface

PiperOrigin-RevId: 757949732
This commit is contained in:
Xiang (Sean) Zhou 2025-05-12 16:05:45 -07:00 committed by Copybara-Service
parent 6a04ff84ba
commit d19927b1bc
2 changed files with 18 additions and 24 deletions

View File

@ -26,17 +26,18 @@ oauth_client_id = os.getenv("OAUTH_CLIENT_ID")
oauth_client_secret = os.getenv("OAUTH_CLIENT_SECRET")
bigquery_tool_set.configure_auth(oauth_client_id, oauth_client_secret)
bigquery_datasets_list = bigquery_tool_set.get_tool("bigquery_datasets_list")
bigquery_datasets_get = bigquery_tool_set.get_tool("bigquery_datasets_get")
bigquery_datasets_insert = bigquery_tool_set.get_tool(
"bigquery_datasets_insert"
tools_to_expose = [
"bigquery_datasets_list",
"bigquery_datasets_get",
"bigquery_datasets_insert",
"bigquery_tables_list",
"bigquery_tables_get",
"bigquery_tables_insert",
]
bigquery_tool_set.set_tool_filter(
lambda tool, ctx=None: tool.name in tools_to_expose
)
bigquery_tables_list = bigquery_tool_set.get_tool("bigquery_tables_list")
bigquery_tables_get = bigquery_tool_set.get_tool("bigquery_tables_get")
bigquery_tables_insert = bigquery_tool_set.get_tool("bigquery_tables_insert")
root_agent = Agent(
model="gemini-2.0-flash",
name="bigquery_agent",
@ -73,12 +74,5 @@ root_agent = Agent(
{userInfo?}
</User>
""",
tools=[
bigquery_datasets_list,
bigquery_datasets_get,
bigquery_datasets_insert,
bigquery_tables_list,
bigquery_tables_get,
bigquery_tables_insert,
],
tools=[bigquery_tool_set],
)

View File

@ -45,12 +45,12 @@ SCOPES = ["https://www.googleapis.com/auth/calendar"]
calendar_tool_set.configure_auth(
client_id=oauth_client_id, client_secret=oauth_client_secret
)
get_calendar_events = calendar_tool_set.get_tool("calendar_events_get")
# list_calendar_events = calendar_tool_set.get_tool("calendar_events_list")
# you can replace below customized list_calendar_events tool with above ADK
# build-in google calendar tool which is commented for now to acheive same
# effect.
calendar_tool_set.set_tool_filter(
# you can also replace below customized `list_calendar_events` with build-in
# google calendar tool by adding `calendar_events_list` in the filter list
lambda tool, ctx=None: tool.name
in ["calendar_events_get"]
)
def list_calendar_events(
@ -210,6 +210,6 @@ root_agent = Agent(
Currnet time: {_time}
""",
tools=[list_calendar_events, get_calendar_events],
tools=[list_calendar_events, calendar_tool_set],
before_agent_callback=update_time,
)