fix: adapt oauth calendar agent and bigquery agent to new Google API toolset interface

PiperOrigin-RevId: 759298612
This commit is contained in:
Xiang (Sean) Zhou 2025-05-15 14:20:59 -07:00 committed by Copybara-Service
parent 4c6820e78c
commit fccd17df6f
2 changed files with 13 additions and 15 deletions

View File

@ -16,7 +16,7 @@ import os
from dotenv import load_dotenv
from google.adk import Agent
from google.adk.tools.google_api_tool import bigquery_toolset
from google.adk.tools.google_api_tool import BigQueryToolset
# Load environment variables from .env file
load_dotenv()
@ -24,8 +24,6 @@ load_dotenv()
# Access the variable
oauth_client_id = os.getenv("OAUTH_CLIENT_ID")
oauth_client_secret = os.getenv("OAUTH_CLIENT_SECRET")
bigquery_toolset.configure_auth(oauth_client_id, oauth_client_secret)
tools_to_expose = [
"bigquery_datasets_list",
"bigquery_datasets_get",
@ -34,8 +32,10 @@ tools_to_expose = [
"bigquery_tables_get",
"bigquery_tables_insert",
]
bigquery_toolset.set_tool_filter(
lambda tool, ctx=None: tool.name in tools_to_expose
bigquery_toolset = BigQueryToolset(
client_id=oauth_client_id,
client_secret=oauth_client_secret,
tool_filter=tools_to_expose,
)
root_agent = Agent(

View File

@ -27,7 +27,7 @@ from google.adk.auth import AuthCredential
from google.adk.auth import AuthCredentialTypes
from google.adk.auth import OAuth2Auth
from google.adk.tools import ToolContext
from google.adk.tools.google_api_tool import calendar_toolset
from google.adk.tools.google_api_tool import CalendarToolset
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
@ -42,16 +42,14 @@ oauth_client_secret = os.getenv("OAUTH_CLIENT_SECRET")
SCOPES = ["https://www.googleapis.com/auth/calendar"]
calendar_toolset.configure_auth(
client_id=oauth_client_id, client_secret=oauth_client_secret
calendar_toolset = CalendarToolset(
# you can also replace below customized `list_calendar_events` with build-in
# google calendar tool by adding `calendar_events_list` in the filter list
client_id=oauth_client_id,
client_secret=oauth_client_secret,
tool_filter=["calendar_events_get"],
)
get_calendar_events = calendar_toolset.get_tool("calendar_events_get")
# list_calendar_events = calendar_toolset.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.
def list_calendar_events(
start_time: str,
@ -210,6 +208,6 @@ root_agent = Agent(
Currnet time: {_time}
""",
tools=[list_calendar_events, get_calendar_events],
tools=[list_calendar_events, calendar_toolset],
before_agent_callback=update_time,
)