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") oauth_client_secret = os.getenv("OAUTH_CLIENT_SECRET")
bigquery_tool_set.configure_auth(oauth_client_id, 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") tools_to_expose = [
bigquery_datasets_get = bigquery_tool_set.get_tool("bigquery_datasets_get") "bigquery_datasets_list",
bigquery_datasets_insert = bigquery_tool_set.get_tool( "bigquery_datasets_get",
"bigquery_datasets_insert" "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( root_agent = Agent(
model="gemini-2.0-flash", model="gemini-2.0-flash",
name="bigquery_agent", name="bigquery_agent",
@ -73,12 +74,5 @@ root_agent = Agent(
{userInfo?} {userInfo?}
</User> </User>
""", """,
tools=[ tools=[bigquery_tool_set],
bigquery_datasets_list,
bigquery_datasets_get,
bigquery_datasets_insert,
bigquery_tables_list,
bigquery_tables_get,
bigquery_tables_insert,
],
) )

View File

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