Skip to content

Commit

Permalink
Tools: Add Gmail support (#838)
Browse files Browse the repository at this point in the history
  • Loading branch information
scottmx81 authored Nov 27, 2024
1 parent 68442e6 commit 81a6f32
Show file tree
Hide file tree
Showing 20 changed files with 578 additions and 6 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ Toolkit is a deployable all-in-one RAG application that enables users to quickly
- [How to add tools](/docs/custom_tool_guides/tool_guide.md)
- [How to add auth to your tools](/docs/custom_tool_guides/tool_auth_guide.md)
- [How to setup Google Drive](/docs/custom_tool_guides/google_drive.md)
- [How to setup Gmail](/docs/custom_tool_guides/gmail.md)
- [How to setup Slack Tool](/docs/custom_tool_guides/slack.md)
- [How to setup Google Text-to-Speech](/docs/text_to_speech.md)
- [How to add authentication](/docs/auth_guide.md)
Expand Down
82 changes: 82 additions & 0 deletions docs/custom_tool_guides/gmail.md
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
```
1 change: 0 additions & 1 deletion src/backend/chat/collate.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,6 @@ async def rerank_and_chunk(

if not text:
# If one doesn't have text, skip it
chunked_outputs = None
reranked_results[tool_call_hashable] = tool_result
continue

Expand Down
3 changes: 3 additions & 0 deletions src/backend/config/configuration.template.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ tools:
- tavily_web_search
python_interpreter:
url: http://terrarium:8080
gmail:
user_scopes:
- https://www.googleapis.com/auth/gmail.readonly
slack:
user_scopes:
- search:read
Expand Down
3 changes: 3 additions & 0 deletions src/backend/config/secrets.template.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,9 @@ tools:
slack:
client_id:
client_secret:
gmail:
client_id:
client_secret:
auth:
secret_key:
google_oauth:
Expand Down
21 changes: 20 additions & 1 deletion src/backend/config/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,23 @@ class GoogleWebSearchSettings(BaseSettings, BaseModel):
default=None, validation_alias=AliasChoices("GOOGLE_SEARCH_CSE_ID", "cse_id")
)

class GmailSettings(BaseSettings, BaseModel):
model_config = SETTINGS_CONFIG
client_id: Optional[str] = Field(
default=None,
validation_alias=AliasChoices("GMAIL_CLIENT_ID", "client_id"),
)
client_secret: Optional[str] = Field(
default=None,
validation_alias=AliasChoices("GMAIL_CLIENT_SECRET", "client_secret"),
)
user_scopes: Optional[str] = Field(
default=None,
validation_alias=AliasChoices(
"GMAIL_USER_SCOPES", "scopes"
),
)


class BraveWebSearchSettings(BaseSettings, BaseModel):
model_config = SETTINGS_CONFIG
Expand Down Expand Up @@ -218,7 +235,9 @@ class ToolSettings(BaseSettings, BaseModel):
slack: Optional[SlackSettings] = Field(
default=SlackSettings()
)

gmail: Optional[GmailSettings] = Field(
default=GmailSettings()
)


class DatabaseSettings(BaseSettings, BaseModel):
Expand Down
2 changes: 2 additions & 0 deletions src/backend/config/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from backend.tools import (
BraveWebSearch,
Calculator,
GmailTool,
GoogleDrive,
GoogleWebSearch,
HybridWebSearch,
Expand Down Expand Up @@ -36,6 +37,7 @@ class Tool(Enum):
Brave_Web_Search = BraveWebSearch
Hybrid_Web_Search = HybridWebSearch
Slack = SlackTool
Gmail = GmailTool


def get_available_tools() -> dict[str, ToolDefinition]:
Expand Down
2 changes: 1 addition & 1 deletion src/backend/routers/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -357,7 +357,7 @@ async def delete_tool_auth(
if user_id is None or user_id == "" or user_id == "default":
logger.error_and_raise_http_exception(event="User ID not found.")

if tool_id not in [tool_name.value for tool_name in Tool]:
if tool_id not in [tool_name.value.ID for tool_name in Tool]:
logger.error_and_raise_http_exception(
event="tool_id must be present in the path of the request and must be a member of the Tool string enum class.",
)
Expand Down
5 changes: 4 additions & 1 deletion src/backend/tools/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from backend.tools.brave_search import BraveWebSearch
from backend.tools.calculator import Calculator
from backend.tools.files import ReadFileTool, SearchFileTool
from backend.tools.gmail import GmailAuth, GmailTool
from backend.tools.google_drive import GoogleDrive, GoogleDriveAuth
from backend.tools.google_search import GoogleWebSearch
from backend.tools.hybrid_search import HybridWebSearch
Expand All @@ -25,5 +26,7 @@
"GoogleWebSearch",
"HybridWebSearch",
"SlackTool",
"SlackAuth"
"SlackAuth",
"GmailTool",
"GmailAuth",
]
11 changes: 11 additions & 0 deletions src/backend/tools/gmail/__init__.py
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",
]
137 changes: 137 additions & 0 deletions src/backend/tools/gmail/auth.py
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
31 changes: 31 additions & 0 deletions src/backend/tools/gmail/client.py
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
2 changes: 2 additions & 0 deletions src/backend/tools/gmail/constants.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
SEARCH_LIMIT = 10
GMAIL_TOOL_ID = "gmail"
Loading

0 comments on commit 81a6f32

Please sign in to comment.