chore: Adapt oauth calendar agent to use authenticated tool

PiperOrigin-RevId: 775433950

Author: seanzhougoogle

contributing/samples/oauth_calendar_agent/agent.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 from datetime import datetime
-import json
 import os
 
 from dotenv import load_dotenv
@@ -27,8 +26,8 @@
 from google.adk.auth import AuthCredentialTypes
 from google.adk.auth import OAuth2Auth
 from google.adk.tools import ToolContext
+from google.adk.tools.authenticated_function_tool import AuthenticatedFunctionTool
 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
 
@@ -56,6 +55,7 @@ def list_calendar_events(
     end_time: str,
     limit: int,
     tool_context: ToolContext,
+    credential: AuthCredential,
 ) -> list[dict]:
   """Search for calendar events.
 
@@ -80,84 +80,11 @@ def list_calendar_events(
   Returns:
       list[dict]: A list of events that match the search criteria.
   """
-  creds = None
-
-  # Check if the tokes were already in the session state, which means the user
-  # has already gone through the OAuth flow and successfully authenticated and
-  # authorized the tool to access their calendar.
-  if "calendar_tool_tokens" in tool_context.state:
-    creds = Credentials.from_authorized_user_info(
-        tool_context.state["calendar_tool_tokens"], SCOPES
-    )
-  if not creds or not creds.valid:
-    # If the access token is expired, refresh it with the refresh token.
-    if creds and creds.expired and creds.refresh_token:
-      creds.refresh(Request())
-    else:
-      auth_scheme = OAuth2(
-          flows=OAuthFlows(
-              authorizationCode=OAuthFlowAuthorizationCode(
-                  authorizationUrl="https://accounts.google.com/o/oauth2/auth",
-                  tokenUrl="https://oauth2.googleapis.com/token",
-                  scopes={
-                      "https://www.googleapis.com/auth/calendar": (
-                          "See, edit, share, and permanently delete all the"
-                          " calendars you can access using Google Calendar"
-                      )
-                  },
-              )
-          )
-      )
-      auth_credential = AuthCredential(
-          auth_type=AuthCredentialTypes.OAUTH2,
-          oauth2=OAuth2Auth(
-              client_id=oauth_client_id, client_secret=oauth_client_secret
-          ),
-      )
-      # If the user has not gone through the OAuth flow before, or the refresh
-      # token also expired, we need to ask users to go through the OAuth flow.
-      # First we check whether the user has just gone through the OAuth flow and
-      # Oauth response is just passed back.
-      auth_response = tool_context.get_auth_response(
-          AuthConfig(
-              auth_scheme=auth_scheme, raw_auth_credential=auth_credential
-          )
-      )
-      if auth_response:
-        # ADK exchanged the access token already for us
-        access_token = auth_response.oauth2.access_token
-        refresh_token = auth_response.oauth2.refresh_token
-
-        creds = Credentials(
-            token=access_token,
-            refresh_token=refresh_token,
-            token_uri=auth_scheme.flows.authorizationCode.tokenUrl,
-            client_id=oauth_client_id,
-            client_secret=oauth_client_secret,
-            scopes=list(auth_scheme.flows.authorizationCode.scopes.keys()),
-        )
-      else:
-        # If there are no auth response which means the user has not gone
-        # through the OAuth flow yet, we need to ask users to go through the
-        # OAuth flow.
-        tool_context.request_credential(
-            AuthConfig(
-                auth_scheme=auth_scheme,
-                raw_auth_credential=auth_credential,
-            )
-        )
-        # The return value is optional and could be any dict object. It will be
-        # wrapped in a dict with key as 'result' and value as the return value
-        # if the object returned is not a dict. This response will be passed
-        # to LLM to generate a user friendly message. e.g. LLM will tell user:
-        # "I need your authorization to access your calendar. Please authorize
-        # me so I can check your meetings for today."
-        return "Need User Authorization to access their calendar."
-    # We store the access token and refresh token in the session state for the
-    # next runs. This is just an example. On production, a tool should store
-    # those credentials in some secure store or properly encrypt it before store
-    # it in the session state.
-    tool_context.state["calendar_tool_tokens"] = json.loads(creds.to_json())
+
+  creds = Credentials(
+      token=credential.oauth2.access_token,
+      refresh_token=credential.oauth2.refresh_token,
+  )
 
   service = build("calendar", "v3", credentials=creds)
   events_result = (
@@ -208,6 +135,33 @@ def update_time(callback_context: CallbackContext):
 
       Currnet time: {_time}
 """,
-    tools=[list_calendar_events, calendar_toolset],
+    tools=[
+        AuthenticatedFunctionTool(
+            func=list_calendar_events,
+            auth_config=AuthConfig(
+                auth_scheme=OAuth2(
+                    flows=OAuthFlows(
+                        authorizationCode=OAuthFlowAuthorizationCode(
+                            authorizationUrl=(
+                                "https://accounts.google.com/o/oauth2/auth"
+                            ),
+                            tokenUrl="https://oauth2.googleapis.com/token",
+                            scopes={
+                                "https://www.googleapis.com/auth/calendar": "",
+                            },
+                        )
+                    )
+                ),
+                raw_auth_credential=AuthCredential(
+                    auth_type=AuthCredentialTypes.OAUTH2,
+                    oauth2=OAuth2Auth(
+                        client_id=oauth_client_id,
+                        client_secret=oauth_client_secret,
+                    ),
+                ),
+            ),
+        ),
+        calendar_toolset,
+    ],
     before_agent_callback=update_time,
 )