Tools: Add Gmail support

README.md

@@ -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)

docs/custom_tool_guides/gmail.md

@@ -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
+```

src/backend/chat/collate.py

@@ -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
 

src/backend/config/configuration.template.yaml

@@ -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

src/backend/config/secrets.template.yaml

@@ -34,6 +34,9 @@ tools:
   slack:
     client_id:
     client_secret:
+  gmail:
+    client_id:
+    client_secret:
 auth:
   secret_key:
   google_oauth:

src/backend/config/settings.py

@@ -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
@@ -218,7 +235,9 @@ class ToolSettings(BaseSettings, BaseModel):
     slack: Optional[SlackSettings] = Field(
         default=SlackSettings()
     )
-
+    gmail: Optional[GmailSettings] = Field(
+        default=GmailSettings()
+    )
 
 
 class DatabaseSettings(BaseSettings, BaseModel):

src/backend/config/tools.py

@@ -6,6 +6,7 @@
 from backend.tools import (
     BraveWebSearch,
     Calculator,
+    GmailTool,
     GoogleDrive,
     GoogleWebSearch,
     HybridWebSearch,
@@ -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]:

src/backend/tools/__init__.py

@@ -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
@@ -25,5 +26,7 @@
     "GoogleWebSearch",
     "HybridWebSearch",
     "SlackTool",
-    "SlackAuth"
+    "SlackAuth",
+    "GmailTool",
+    "GmailAuth",
 ]

src/backend/tools/gmail/__init__.py

@@ -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",
+]

src/backend/tools/gmail/auth.py

@@ -0,0 +1,136 @@
+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,
+            "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

src/backend/tools/gmail/client.py

@@ -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

src/backend/tools/gmail/constants.py

@@ -0,0 +1,2 @@
+SEARCH_LIMIT = 10
+GMAIL_TOOL_ID = "gmail"