27
27
from google .adk .auth import AuthCredentialTypes
28
28
from google .adk .auth import OAuth2Auth
29
29
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
32
30
from google .adk .tools .google_api_tool import CalendarToolset
33
31
from google .auth .transport .requests import Request
34
32
from google .oauth2 .credentials import Credentials
@@ -58,7 +56,6 @@ def list_calendar_events(
58
56
end_time : str ,
59
57
limit : int ,
60
58
tool_context : ToolContext ,
61
- credential : AuthCredential ,
62
59
) -> list [dict ]:
63
60
"""Search for calendar events.
64
61
@@ -83,11 +80,84 @@ def list_calendar_events(
83
80
Returns:
84
81
list[dict]: A list of events that match the search criteria.
85
82
"""
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 ())
91
161
92
162
service = build ("calendar" , "v3" , credentials = creds )
93
163
events_result = (
@@ -138,38 +208,6 @@ def update_time(callback_context: CallbackContext):
138
208
139
209
Currnet time: {_time}
140
210
""" ,
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 ],
174
212
before_agent_callback = update_time ,
175
213
)
0 commit comments