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 dotenv import load_dotenv
from google.adk import Agent 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 environment variables from .env file
load_dotenv() load_dotenv()
@ -24,8 +24,6 @@ load_dotenv()
# Access the variable # Access the variable
oauth_client_id = os.getenv("OAUTH_CLIENT_ID") 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_toolset.configure_auth(oauth_client_id, oauth_client_secret)
tools_to_expose = [ tools_to_expose = [
"bigquery_datasets_list", "bigquery_datasets_list",
"bigquery_datasets_get", "bigquery_datasets_get",
@ -34,8 +32,10 @@ tools_to_expose = [
"bigquery_tables_get", "bigquery_tables_get",
"bigquery_tables_insert", "bigquery_tables_insert",
] ]
bigquery_toolset.set_tool_filter( bigquery_toolset = BigQueryToolset(
lambda tool, ctx=None: tool.name in tools_to_expose client_id=oauth_client_id,
client_secret=oauth_client_secret,
tool_filter=tools_to_expose,
) )
root_agent = Agent( 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 AuthCredentialTypes
from google.adk.auth import OAuth2Auth from google.adk.auth import OAuth2Auth
from google.adk.tools import ToolContext 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.auth.transport.requests import Request
from google.oauth2.credentials import Credentials from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build from googleapiclient.discovery import build
@ -42,16 +42,14 @@ oauth_client_secret = os.getenv("OAUTH_CLIENT_SECRET")
SCOPES = ["https://www.googleapis.com/auth/calendar"] SCOPES = ["https://www.googleapis.com/auth/calendar"]
calendar_toolset.configure_auth( calendar_toolset = CalendarToolset(
client_id=oauth_client_id, client_secret=oauth_client_secret # 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( def list_calendar_events(
start_time: str, start_time: str,
@ -210,6 +208,6 @@ root_agent = Agent(
Currnet time: {_time} Currnet time: {_time}
""", """,
tools=[list_calendar_events, get_calendar_events], tools=[list_calendar_events, calendar_toolset],
before_agent_callback=update_time, before_agent_callback=update_time,
) )