From 8ebf229c4787a5e3f559ea418a11f2ed407637e0 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Fri, 13 Jun 2025 23:23:45 -0700 Subject: [PATCH] chore: Add base credential service interface (WIP) PiperOrigin-RevId: 771358480 --- .../samples/oauth_calendar_agent/agent.py | 118 ++++++------------ src/google/adk/auth/auth_handler.py | 4 +- src/google/adk/auth/auth_tool.py | 24 +++- .../adk/auth/credential_service/__init__.py | 13 ++ .../base_credential_service.py | 75 +++++++++++ tests/unittests/auth/test_auth_config.py | 26 +++- tests/unittests/auth/test_auth_handler.py | 6 +- 7 files changed, 175 insertions(+), 91 deletions(-) create mode 100644 src/google/adk/auth/credential_service/__init__.py create mode 100644 src/google/adk/auth/credential_service/base_credential_service.py diff --git a/contributing/samples/oauth_calendar_agent/agent.py b/contributing/samples/oauth_calendar_agent/agent.py index a1b1dea8..9d56d3ff 100644 --- a/contributing/samples/oauth_calendar_agent/agent.py +++ b/contributing/samples/oauth_calendar_agent/agent.py @@ -27,6 +27,8 @@ 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.authenticated_tool.base_authenticated_tool import AuthenticatedFunctionTool +from google.adk.tools.authenticated_tool.credentials_store import ToolContextCredentialsStore from google.adk.tools.google_api_tool import CalendarToolset from google.auth.transport.requests import Request from google.oauth2.credentials import Credentials @@ -56,6 +58,7 @@ def list_calendar_events( end_time: str, limit: int, tool_context: ToolContext, + credential: AuthCredential, ) -> list[dict]: """Search for calendar events. @@ -80,84 +83,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 +138,38 @@ root_agent = Agent( 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": ( + "See, edit, share, and permanently delete" + " all the calendars you can access using" + " Google Calendar" + ) + }, + ) + ) + ), + raw_auth_credential=AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth( + client_id=oauth_client_id, + client_secret=oauth_client_secret, + ), + ), + ), + credential_store=ToolContextCredentialsStore(), + ), + calendar_toolset, + ], before_agent_callback=update_time, ) diff --git a/src/google/adk/auth/auth_handler.py b/src/google/adk/auth/auth_handler.py index 5f80ee3f..3e13cbac 100644 --- a/src/google/adk/auth/auth_handler.py +++ b/src/google/adk/auth/auth_handler.py @@ -49,7 +49,7 @@ class AuthHandler: def parse_and_store_auth_response(self, state: State) -> None: - credential_key = "temp:" + self.auth_config.get_credential_key() + credential_key = "temp:" + self.auth_config.credential_key state[credential_key] = self.auth_config.exchanged_auth_credential if not isinstance( @@ -67,7 +67,7 @@ class AuthHandler: raise ValueError("auth_scheme is empty.") def get_auth_response(self, state: State) -> AuthCredential: - credential_key = "temp:" + self.auth_config.get_credential_key() + credential_key = "temp:" + self.auth_config.credential_key return state.get(credential_key, None) def generate_auth_request(self) -> AuthConfig: diff --git a/src/google/adk/auth/auth_tool.py b/src/google/adk/auth/auth_tool.py index a1a19ab8..53c571d4 100644 --- a/src/google/adk/auth/auth_tool.py +++ b/src/google/adk/auth/auth_tool.py @@ -14,6 +14,10 @@ from __future__ import annotations +from typing import Optional + +from typing_extensions import deprecated + from .auth_credential import AuthCredential from .auth_credential import BaseModelWithConfig from .auth_schemes import AuthScheme @@ -45,11 +49,23 @@ class AuthConfig(BaseModelWithConfig): this field to guide the user through the OAuth2 flow and fill auth response in this field""" + credential_key: Optional[str] = None + """A user specified key used to load and save this credential in a credential + service. + """ + + def __init__(self, **data): + super().__init__(**data) + if self.credential_key: + return + self.credential_key = self.get_credential_key() + + @deprecated("This method is deprecated. Use credential_key instead.") def get_credential_key(self): - """Generates a hash key based on auth_scheme and raw_auth_credential. This - hash key can be used to store / retrieve exchanged_auth_credential in a - credentials store. + """Builds a hash key based on auth_scheme and raw_auth_credential used to + save / load this credential to / from a credentials service. """ + auth_scheme = self.auth_scheme if auth_scheme.model_extra: @@ -62,7 +78,7 @@ class AuthConfig(BaseModelWithConfig): ) auth_credential = self.raw_auth_credential - if auth_credential.model_extra: + if auth_credential and auth_credential.model_extra: auth_credential = auth_credential.model_copy(deep=True) auth_credential.model_extra.clear() credential_name = ( diff --git a/src/google/adk/auth/credential_service/__init__.py b/src/google/adk/auth/credential_service/__init__.py new file mode 100644 index 00000000..0a2669d7 --- /dev/null +++ b/src/google/adk/auth/credential_service/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/src/google/adk/auth/credential_service/base_credential_service.py b/src/google/adk/auth/credential_service/base_credential_service.py new file mode 100644 index 00000000..7416ccc6 --- /dev/null +++ b/src/google/adk/auth/credential_service/base_credential_service.py @@ -0,0 +1,75 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from abc import ABC +from abc import abstractmethod +from typing import Optional + +from ...tools.tool_context import ToolContext +from ...utils.feature_decorator import working_in_progress +from ..auth_credential import AuthCredential +from ..auth_tool import AuthConfig + + +@working_in_progress("Implementation are in progress. Don't use it for now.") +class BaseCredentialService(ABC): + """Abstract class for Service that loads / saves tool credentials from / to + the backend credential store.""" + + @abstractmethod + async def load_credential( + self, + auth_config: AuthConfig, + tool_context: ToolContext, + ) -> Optional[AuthCredential]: + """ + Loads the credential by auth config and current tool context from the + backend credential store. + + Args: + auth_config: The auth config which contains the auth scheme and auth + credential information. auth_config.get_credential_key will be used to + build the key to load the credential. + + tool_context: The context of the current invocation when the tool is + trying to load the credential. + + Returns: + Optional[AuthCredential]: the credential saved in the store. + + """ + + @abstractmethod + async def save_credential( + self, + auth_config: AuthConfig, + tool_context: ToolContext, + ) -> None: + """ + Saves the exchanged_auth_credential in auth config to the backend credential + store. + + Args: + auth_config: The auth config which contains the auth scheme and auth + credential information. auth_config.get_credential_key will be used to + build the key to save the credential. + + tool_context: The context of the current invocation when the tool is + trying to save the credential. + + Returns: + None + """ diff --git a/tests/unittests/auth/test_auth_config.py b/tests/unittests/auth/test_auth_config.py index 7a20cfc6..a398ef32 100644 --- a/tests/unittests/auth/test_auth_config.py +++ b/tests/unittests/auth/test_auth_config.py @@ -68,10 +68,28 @@ def auth_config(oauth2_auth_scheme, oauth2_credentials): ) -def test_get_credential_key(auth_config): +@pytest.fixture +def auth_config_with_key(oauth2_auth_scheme, oauth2_credentials): + """Create an AuthConfig for testing.""" + + return AuthConfig( + auth_scheme=oauth2_auth_scheme, + raw_auth_credential=oauth2_credentials, + credential_key="test_key", + ) + + +def test_custom_credential_key(auth_config_with_key): + """Test using custom credential key.""" + + key = auth_config_with_key.credential_key + assert key == "test_key" + + +def test_credential_key(auth_config): """Test generating a unique credential key.""" - key = auth_config.get_credential_key() + key = auth_config.credential_key assert key.startswith("adk_oauth2_") assert "_oauth2_" in key @@ -80,8 +98,8 @@ def test_get_credential_key_with_extras(auth_config): """Test generating a key when model_extra exists.""" # Add model_extra to test cleanup - original_key = auth_config.get_credential_key() - key = auth_config.get_credential_key() + original_key = auth_config.credential_key + key = auth_config.credential_key auth_config.auth_scheme.model_extra["extra_field"] = "value" auth_config.raw_auth_credential.model_extra["extra_field"] = "value" diff --git a/tests/unittests/auth/test_auth_handler.py b/tests/unittests/auth/test_auth_handler.py index 402e4d0c..aaed35a1 100644 --- a/tests/unittests/auth/test_auth_handler.py +++ b/tests/unittests/auth/test_auth_handler.py @@ -387,7 +387,7 @@ class TestGetAuthResponse: state = MockState() # Store a credential in the state - credential_key = auth_config.get_credential_key() + credential_key = auth_config.credential_key state["temp:" + credential_key] = oauth2_credentials_with_auth_uri result = handler.get_auth_response(state) @@ -418,7 +418,7 @@ class TestParseAndStoreAuthResponse: handler.parse_and_store_auth_response(state) - credential_key = auth_config.get_credential_key() + credential_key = auth_config.credential_key assert ( state["temp:" + credential_key] == auth_config.exchanged_auth_credential ) @@ -436,7 +436,7 @@ class TestParseAndStoreAuthResponse: handler.parse_and_store_auth_response(state) - credential_key = auth_config_with_exchanged.get_credential_key() + credential_key = auth_config_with_exchanged.credential_key assert state["temp:" + credential_key] == mock_exchange_token.return_value assert mock_exchange_token.called