10000 chore: fix oauth_calendar_agent example · google/adk-python@2329cde · GitHub
[go: up one dir, main page]

Skip to content

Commit 2329cde

Browse files
seanzhougooglecopybara-github
authored andcommitted
chore: fix oauth_calendar_agent example
PiperOrigin-RevId: 771756083
1 parent a4d432a commit 2329cde

File tree

1 file changed

+79
-41
lines changed
  • contributing/samples/oauth_calendar_agent

1 file changed

+79
-41
lines changed

contributing/samples/oauth_calendar_agent/agent.py

Lines changed: 79 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,6 @@
2727
from google.adk.auth import AuthCredentialTypes
2828
from google.adk.auth import OAuth2Auth
2929
from google.adk.tools import ToolContext
30-
from google.adk.tools.authenticated_tool.base_authenticated_tool import AuthenticatedFunctionTool
31-
from google.adk.tools.authenticated_tool.credentials_store import ToolContextCredentialsStore
3230
from google.adk.tools.google_api_tool import CalendarToolset
3331
from google.auth.transport.requests import Request
3432
from google.oauth2.credentials import Credentials
@@ -58,7 +56,6 @@ def list_calendar_events(
5856
end_time: str,
5957
limit: int,
6058
tool_context: ToolContext,
61-
credential: AuthCredential,
6259
) -> list[dict]:
6360
"""Search for calendar events.
6461
@@ -83,11 +80,84 @@ def list_calendar_events(
8380
Returns:
8481
list[dict]: A list of events that match the search criteria.
8582
"""
86-
87-
creds = Credentials(
88-
token=credential.oauth2.access_token,
89-
refresh_token=credential.oauth2.refresh_token,
90-
)
83+
creds = None
84+
85+
# Check if the tokes were already in the session state, which means the user
86+
# has already gone through the OAuth flow and successfully authenticated and
87+
# authorized the tool to access their calendar.
88+
if "calendar_tool_tokens" in tool_context.state:
89+
creds = Credentials.from_authorized_user_info(
90+
tool_context.state["calendar_tool_tokens"], SCOPES
91+
)
92+
if not creds or not creds.valid:
93+
# If the access token is expired, refresh it with the refresh token.
94+
if creds and creds.expired and creds.refresh_token:
95+
creds.refresh(Request())
96+
else:
97+
auth_scheme = OAuth2(
98+
flows=OAuthFlows(
99+
authorizationCode=OAuthFlowAuthorizationCode(
100+
authorizationUrl="https://accounts.google.com/o/oauth2/auth",
101+
tokenUrl="https://oauth2.googleapis.com/token",
102+
scopes={
103+
"https://www.googleapis.com/auth/calendar": (
104+
"See, edit, share, and permanently delete all the"
105+
" calendars you can access using Google Calendar"
106+
)
107+
},
108+
)
109+
)
110+
)
111+
auth_credential = AuthCredential(
112+
auth_type=AuthCredentialTypes.OAUTH2,
113+
oauth2=OAuth2Auth(
114+
client_id=oauth_client_id, client_secret=oauth_client_secret
115+
),
116+
)
117+
# If the user has not gone through the OAuth flow before, or the refresh
118+
# token also expired, we need to ask users to go through the OAuth flow.
119+
# First we check whether the user has just gone through the OAuth flow and
120+
# Oauth response is just passed back.
121+
auth_response = tool_context.get_auth_response(
122+
AuthConfig(
123+
auth_scheme=auth_scheme, raw_auth_credential=auth_credential
124+
)
125+
)
126+
if auth_response:
127+
# ADK exchanged the access token already for us
128+
access_token = auth_response.oauth2.access_token
129+
refresh_token = auth_response.oauth2.refresh_token
130+
131+
creds = Credentials(
132+
token=access_token,
133+
refresh_token=refresh_token,
134+
token_uri=auth_scheme.flows.authorizationCode.tokenUrl,
135+
client_id=oauth_client_id,
136+
client_secret=oauth_client_secret,
137+
scopes=list(auth_scheme.flows.authorizationCode.scopes.keys()),
138+
)
139+
else:
140+
# If there are no auth response which means the user has not gone
141+
# through the OAuth flow yet, we need to ask users to go through the
142+
# OAuth flow.
143+
tool_context.request_credential(
144+
AuthConfig(
145+
auth_scheme=auth_scheme,
146+
raw_auth_credential=auth_credential,
147+
)
148+
)
149+
# The return value is optional and could be any dict object. It will be
150+
# wrapped in a dict with key as 'result' and value as the return value
151+
# if the object returned is not a dict. This response will be passed
152+
# to LLM to generate a user friendly message. e.g. LLM will tell user:
153+
# "I need your authorization to access your calendar. Please authorize
154+
# me so I can check your meetings for today."
155+
return "Need User Authorization to access their calendar."
156+
# We store the access token and refresh token in the session state for the
157+
# next runs. This is just an example. On production, a tool should store
158+
# those credentials in some secure store or properly encrypt it before store
159+
# it in the session state.
160+
tool_context.state["calendar_tool_tokens"] = json.loads(creds.to_json())
91161

92162
service = build("calendar", "v3", credentials=creds)
93163
events_result = (
@@ -138,38 +208,6 @@ def update_time(callback_context: CallbackContext):
138208
139209
Currnet time: {_time}
140210
""",
141-
tools=[
142-
AuthenticatedFunctionTool(
143-
func=list_calendar_events,
144-
auth_config=AuthConfig(
145-
auth_scheme=OAuth2(
146-
flows=OAuthFlows(
147-
authorizationCode=OAuthFlowAuthorizationCode(
148-
authorizationUrl=(
149-
"https://accounts.google.com/o/oauth2/auth"
150-
),
151-
tokenUrl="https://oauth2.googleapis.com/token",
152-
scopes={
153-
"https://www.googleapis.com/auth/calendar": (
154-
"See, edit, share, and permanently delete"
155-
" all the calendars you can access using"
156-
" Google Calendar"
157-
)
158-
},
159-
)
160-
)
161-
),
162-
raw_auth_credential=AuthCredential(
163-
auth_type=AuthCredentialTypes.OAUTH2,
164-
oauth2=OAuth2Auth(
165-
client_id=oauth_client_id,
166-
client_secret=oauth_client_secret,
167-
),
168-
),
169-
),
170-
credential_store=ToolContextCredentialsStore(),
171-
),
172-
calendar_toolset,
173-
],
211+
tools=[list_calendar_events, calendar_toolset],
174212
before_agent_callback=update_time,
175213
)

0 commit comments

Comments
 (0)
0