-
Notifications
You must be signed in to change notification settings - Fork 375
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
20 changed files
with
578 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,82 @@ | ||
# Gmail Tool Setup | ||
|
||
To set up the Gmail tool you will need to configure API access in Google Cloud Console. | ||
|
||
Follow the steps below to set it up: | ||
|
||
## 1. Create Project in Google Cloud Console | ||
|
||
Head to the [Google Cloud Console](https://console.cloud.google.com/) and create a new project. | ||
After creating the app, you will see the `APIs & Services` section. Under `Enabled APIs & services`, enable the | ||
Gmail API. | ||
|
||
## 2. Configure OAuth Consent Screen | ||
|
||
Before you can generate the API credentials, you must first configure the OAuth consent screen. | ||
|
||
You will need to configure the `Application home page` and `Authorized domain 1` fields, with a URL and domain that | ||
point to where you are running the Toolkit. If you are running it in a local environment, Ngrok can be used as a proxy | ||
to access the to the Toolkit in local. Using `localhost` is not accepted value for these fields. | ||
|
||
If you choose to use Ngrok, you can start it with: | ||
|
||
`ngrok http -domain <your_custom_domain>.ngrok.dev 8000` | ||
|
||
And then use the domain you used here in the OAuth Consent Screen configuration. | ||
|
||
## 3. Generate Credentials | ||
|
||
Once the OAuth consent screen has been configured, choose the `Credentials` menu option. Click `+ CREATE CREDENTIALS` | ||
at the top, and choose the OAuth client ID option. | ||
|
||
If running the Toolkit in your local environment, you can use `http://localhost` as the Authorized Javascript origin. | ||
|
||
For the Authorized redirect URI, it must point to the Toolkit backend. The path should be `/v1/tool/auth`. For example: | ||
|
||
```bash | ||
https://<your_backend_url>/v1/tool/auth | ||
``` | ||
|
||
## 3. Set Up Environment Variables | ||
|
||
Then set the following environment variables. You can either set the values in your `secrets.yaml` file: | ||
```bash | ||
Gmail: | ||
client_id: <your_client_id from the previous step> | ||
client_secret: <your_client_secret from the previous step> | ||
``` | ||
or update your `.env` configuration to contain: | ||
```bash | ||
GMAIL_CLIENT_ID=<your_client_id from the previous step> | ||
GMAIL_CLIENT_SECRET=<your_client_secret from the previous step> | ||
``` | ||
|
||
## 4. Enable the Gmail Tool in the Frontend | ||
|
||
To enable the Gmail tool in the frontend, you will need to modify the `src/interfaces/assistants_web/src/constants/tools.ts` | ||
file. Add the `TOOL_GMAIL_ID` to the `AGENT_SETTINGS_TOOLS` list. | ||
|
||
```typescript | ||
export const AGENT_SETTINGS_TOOLS = [ | ||
TOOL_HYBRID_WEB_SEARCH_ID, | ||
TOOL_PYTHON_INTERPRETER_ID, | ||
TOOL_WEB_SCRAPE_ID, | ||
TOOL_GMAIL_ID, | ||
]; | ||
``` | ||
|
||
To enable the Gmail tool in the frontend for the base agent, you will need to modify the | ||
`src/interfaces/assistants_web/src/constants/tools.ts` file. Remove `TOOL_GMAIL_ID` from the | ||
`BASE_AGENT_EXCLUDED_TOOLS` list. By default, the Gmail Tool is disabled for the Base Agent. | ||
|
||
```typescript | ||
export const BASE_AGENT_EXCLUDED_TOOLS = []; | ||
``` | ||
|
||
## 5. Run the Backend and Frontend | ||
|
||
run next command to start the backend and frontend: | ||
|
||
```bash | ||
make dev | ||
``` |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,11 @@ | ||
from backend.tools.gmail.auth import GmailAuth | ||
from backend.tools.gmail.constants import ( | ||
GMAIL_TOOL_ID, | ||
) | ||
from backend.tools.gmail.tool import GmailTool | ||
|
||
__all__ = [ | ||
"GmailAuth", | ||
"GmailTool", | ||
"GMAIL_TOOL_ID", | ||
] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,137 @@ | ||
import datetime | ||
import json | ||
import urllib.parse | ||
|
||
import requests | ||
from fastapi import Request | ||
|
||
from backend.config.settings import Settings | ||
from backend.crud import tool_auth as tool_auth_crud | ||
from backend.database_models.database import DBSessionDep | ||
from backend.database_models.tool_auth import ToolAuth as ToolAuthModel | ||
from backend.schemas.tool_auth import UpdateToolAuth | ||
from backend.services.auth.crypto import encrypt | ||
from backend.services.logger.utils import LoggerFactory | ||
from backend.tools.base import BaseToolAuthentication | ||
from backend.tools.gmail.constants import GMAIL_TOOL_ID | ||
from backend.tools.utils.mixins import ToolAuthenticationCacheMixin | ||
|
||
logger = LoggerFactory().get_logger() | ||
|
||
|
||
class GmailAuth(BaseToolAuthentication, ToolAuthenticationCacheMixin): | ||
TOOL_ID = GMAIL_TOOL_ID | ||
AUTH_ENDPOINT = "https://accounts.google.com/o/oauth2/v2/auth" | ||
TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token" | ||
DEFAULT_USER_SCOPES = ['https://www.googleapis.com/auth/gmail.readonly'] | ||
|
||
def __init__(self): | ||
super().__init__() | ||
|
||
self.GMAIL_CLIENT_ID = Settings().get("tools.gmail.client_id") | ||
self.GMAIL_CLIENT_SECRET = Settings().get("tools.gmail.client_secret") | ||
self.REDIRECT_URL = f"{self.BACKEND_HOST}/v1/tool/auth" | ||
self.USER_SCOPES = Settings().get("tools.gmail.user_scopes") or self.DEFAULT_USER_SCOPES | ||
|
||
if any([ | ||
self.GMAIL_CLIENT_ID is None, | ||
self.GMAIL_CLIENT_SECRET is None | ||
]): | ||
raise ValueError( | ||
"GMAIL_CLIENT_ID and GMAIL_CLIENT_SECRET must be set to use Gmail Tool Auth." | ||
) | ||
|
||
def get_auth_url(self, user_id: str) -> str: | ||
key = self.insert_tool_auth_cache(user_id, self.TOOL_ID) | ||
state = {"key": key} | ||
|
||
params = { | ||
"response_type": "code", | ||
"client_id": self.GMAIL_CLIENT_ID, | ||
"scope": " ".join(self.USER_SCOPES or []), | ||
"redirect_uri": self.REDIRECT_URL, | ||
"prompt": "select_account consent", | ||
"state": json.dumps(state), | ||
"access_type": "offline", | ||
"include_granted_scopes": "true", | ||
} | ||
|
||
return f"{self.AUTH_ENDPOINT}?{urllib.parse.urlencode(params)}" | ||
|
||
def retrieve_auth_token( | ||
self, request: Request, session: DBSessionDep, user_id: str | ||
) -> str: | ||
if request.query_params.get("error"): | ||
error = request.query_params.get("error") or "Unknown error" | ||
logger.error(event=f"[Gmail Tool] Auth token error: {error}.") | ||
return error | ||
|
||
body = { | ||
"code": request.query_params.get("code"), | ||
"client_id": self.GMAIL_CLIENT_ID, | ||
"client_secret": self.GMAIL_CLIENT_SECRET, | ||
"redirect_uri": self.REDIRECT_URL, | ||
"grant_type": "authorization_code", | ||
} | ||
|
||
response = requests.post(self.TOKEN_ENDPOINT, json=body) | ||
response_body = response.json() | ||
|
||
if response.status_code != 200: | ||
logger.error( | ||
event=f"[Gmail] Error retrieving auth token: {response_body}" | ||
) | ||
return "" | ||
|
||
tool_auth_crud.create_tool_auth( | ||
session, | ||
ToolAuthModel( | ||
user_id=user_id, | ||
tool_id=self.TOOL_ID, | ||
token_type=response_body["token_type"], | ||
encrypted_access_token=encrypt(response_body["access_token"]), | ||
encrypted_refresh_token=encrypt(response_body["refresh_token"]), | ||
expires_at=datetime.datetime.now() | ||
+ datetime.timedelta(seconds=response_body["expires_in"]), | ||
), | ||
) | ||
|
||
return "" | ||
|
||
def try_refresh_token( | ||
self, session: DBSessionDep, user_id: str, tool_auth: ToolAuthModel | ||
) -> bool: | ||
body = { | ||
"client_id": self.GMAIL_CLIENT_ID, | ||
"client_secret": self.GMAIL_CLIENT_SECRET, | ||
"refresh_token": tool_auth.refresh_token, | ||
"grant_type": "refresh_token", | ||
} | ||
|
||
response = requests.post(self.TOKEN_ENDPOINT, json=body) | ||
response_body = response.json() | ||
|
||
if response.status_code != 200: | ||
logger.error( | ||
event=f"[Gmail] Error refreshing token: {response_body}" | ||
) | ||
return False | ||
|
||
existing_tool_auth = tool_auth_crud.get_tool_auth( | ||
session, self.TOOL_ID, user_id | ||
) | ||
tool_auth_crud.update_tool_auth( | ||
session, | ||
existing_tool_auth, | ||
UpdateToolAuth( | ||
user_id=user_id, | ||
tool_id=self.TOOL_ID, | ||
token_type=response_body["token_type"], | ||
encrypted_access_token=encrypt(response_body["access_token"]), | ||
encrypted_refresh_token=tool_auth.encrypted_refresh_token, | ||
expires_at=datetime.datetime.now() | ||
+ datetime.timedelta(seconds=response_body["expires_in"]), | ||
), | ||
) | ||
|
||
return True |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
from google.oauth2.credentials import Credentials | ||
from googleapiclient.discovery import build | ||
|
||
|
||
class GmailClient: | ||
def __init__(self, auth_token, search_limit=20): | ||
creds = Credentials(auth_token) | ||
self.service = build("gmail", "v1", credentials=creds, cache_discovery=False) | ||
self.search_limit = search_limit | ||
|
||
def search_all(self, query): | ||
return ( | ||
self.service.users() | ||
.messages() | ||
.list(userId="me", q=query, maxResults=self.search_limit) | ||
.execute() | ||
) | ||
|
||
def retrieve_messages(self, message_ids): | ||
messages = [] | ||
|
||
for message_id in message_ids: | ||
message = ( | ||
self.service.users() | ||
.messages() | ||
.get(userId="me", id=message_id) | ||
.execute() | ||
) | ||
messages.append(message) | ||
|
||
return messages |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
SEARCH_LIMIT = 10 | ||
GMAIL_TOOL_ID = "gmail" |
Oops, something went wrong.