admin-valentin/arcade-mcp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

gmail: search_emails_by_header tool (#28)

Browse source 
Adds a new gmail tool to search email, plus some cleanup.
This commit is contained in:
Nate Barbettini 2024-08-30 15:29:02 -07:00 committed by GitHub
parent aa1b59497b
commit 408f2e6300
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 259 additions and 100 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
arcade/arcade/core/catalog.py
View file

@ -3,6 +3,7 @@ import inspect
from collections.abc import Iterator
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from importlib import import_module
from types import ModuleType
from typing import (
@ -239,6 +240,9 @@ def create_input_definition(func: Callable) -> ToolInputs:
if is_string_literal(tool_field_info.field_type):
is_enum = True
enum_values = [str(e) for e in get_args(tool_field_info.field_type)]
elif issubclass(tool_field_info.field_type, Enum):
is_enum = True
enum_values = [e.value for e in tool_field_info.field_type]
# If the field has a default value, it is not required
# If the field is optional, it is not required
@ -479,6 +483,8 @@ def get_wire_type(
origin = _type.__origin__
if origin in [list, dict]:
return "json"
elif issubclass(_type, Enum):
return "string"
elif issubclass(_type, BaseModel):
return "json"
raise ToolDefinitionError(f"Unsupported parameter type: {_type}")

34
arcade/tests/tool/test_create_tool_definition.py
View file

@ -1,3 +1,4 @@
from enum import Enum
from typing import Annotated, Literal, Optional
import pytest
@ -88,13 +89,23 @@ def func_with_every_param(
pass
class TestEnum(Enum):
FOO_BAR = "foo bar"
BAZ = "baz"
@tool(desc="A function that takes an enum")
def func_with_enum_param(param1: Annotated[TestEnum, "an enum"]):
pass
@tool(desc="A function that takes a dictionary")
def func_with_dict_param(param1: Annotated[dict, "a cool dictionary"]):
pass
@tool(desc="A function that takes a string enum")
def func_with_string_enum_param(param1: Annotated[Literal["value1", "value2"], "a few choices"]):
@tool(desc="A function that takes a string literal")
def func_with_string_literal_param(param1: Annotated[Literal["value1", "value2"], "a few choices"]):
pass
@ -329,6 +340,23 @@ def func_with_complex_return() -> list[dict[str, str]]:
},
id="func_with_every_param",
),
pytest.param(
func_with_enum_param,
{
"inputs": ToolInputs(
parameters=[
InputParameter(
name="param1",
description="an enum",
inferrable=True,
required=True,
value_schema=ValueSchema(val_type="string", enum=["foo bar", "baz"]),
)
]
),
},
id="func_with_enum_param",
),
pytest.param(
func_with_dict_param,
{
@ -347,7 +375,7 @@ def func_with_complex_return() -> list[dict[str, str]]:
id="func_with_dict_param",
),
pytest.param(
func_with_string_enum_param,
func_with_string_literal_param,
{
"inputs": ToolInputs(
parameters=[

2
examples/fastapi/arcade_example_fastapi/main.py
View file

@ -21,6 +21,7 @@ actor = FastAPIActor(app)
# actor.register_tool(arithmetic.divide)
# actor.register_tool(arithmetic.sqrt)
actor.register_tool(gmail.get_emails)
actor.register_tool(gmail.search_emails_by_header)
actor.register_tool(gmail.write_draft)
actor.register_tool(repo.count_stargazers)
actor.register_tool(repo.search_issues)
@ -45,6 +46,7 @@ async def postChat(request: ChatRequest, tool_choice: str = "execute"):
max_tokens=500,
tools=[
"GetEmails",
"SearchEmailsByHeader",
"WriteDraft",
"CountStargazers",
"SetStarred",

219
toolkits/gmail/arcade_gmail/tools/gmail.py
View file

@ -1,11 +1,13 @@
import base64
import datetime
from enum import Enum
import json
import re
from base64 import urlsafe_b64decode
from email.mime.text import MIMEText
from typing import Annotated, Any, Dict, Optional
from typing import Annotated, Optional
from arcade.core.errors import ToolExecutionError
from googleapiclient.errors import HttpError
from bs4 import BeautifulSoup
from arcade_gmail.tools.utils import parse_email
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
@ -48,6 +50,122 @@ def get_draft_url(draft_id):
return f"https://mail.google.com/mail/u/0/#drafts/{draft_id}"
class DateRange(Enum):
TODAY = "today"
YESTERDAY = "yesterday"
LAST_7_DAYS = "last_7_days"
LAST_30_DAYS = "last_30_days"
THIS_MONTH = "this_month"
LAST_MONTH = "last_month"
THIS_YEAR = "this_year"
def to_date_query(self):
today = datetime.datetime.now()
result = "after:"
comparison_date = today
if self == DateRange.YESTERDAY:
comparison_date = today - datetime.timedelta(days=1)
elif self == DateRange.LAST_7_DAYS:
comparison_date = today - datetime.timedelta(days=7)
elif self == DateRange.LAST_30_DAYS:
comparison_date = today - datetime.timedelta(days=30)
elif self == DateRange.THIS_MONTH:
comparison_date = today.replace(day=1)
elif self == DateRange.LAST_MONTH:
comparison_date = (
today.replace(day=1) - datetime.timedelta(days=1)
).replace(day=1)
elif self == DateRange.THIS_YEAR:
comparison_date = today.replace(month=1, day=1)
elif self == DateRange.LAST_MONTH:
comparison_date = (
today.replace(month=1, day=1) - datetime.timedelta(days=1)
).replace(month=1, day=1)
return result + comparison_date.strftime("%Y/%m/%d")
@tool(
requires_auth=Google(
scope=["https://www.googleapis.com/auth/gmail.readonly"],
)
)
async def search_emails_by_header(
context: ToolContext,
sender: Annotated[
Optional[str], "The name or email address of the sender of the email"
] = None,
recipient: Annotated[
Optional[str], "The name or email address of the recipient"
] = None,
subject: Annotated[
Optional[str], "Words to find in the subject of the email"
] = None,
body: Annotated[Optional[str], "Words to find in the body of the email"] = None,
date_range: Annotated[Optional[DateRange], "The date range of the email"] = None,
limit: Annotated[Optional[int], "The maximum number of emails to return"] = 25,
) -> Annotated[str, "A list of email details in JSON format"]:
"""Search for emails by header.
One of the following MUST be provided: sender, recipient, subject, body."""
if not any([sender, recipient, subject, body]):
raise ValueError(
"At least one of sender, recipient, subject, or body must be provided."
)
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Build the query string
query = []
if sender:
query.append(f"from:{sender}")
if recipient:
query.append(f"to:{recipient}")
if subject:
query.append(f"subject:{subject}")
if body:
query.append(body)
if date_range:
query.append(date_range.to_date_query())
query_string = " ".join(query)
try:
# Perform the search
response = (
service.users()
.messages()
.list(userId="me", q=query_string, maxResults=limit or 100)
.execute()
)
messages = response.get("messages", [])
if not messages:
return json.dumps({"emails": []})
emails = []
for msg in messages:
try:
email_data = (
service.users().messages().get(userId="me", id=msg["id"]).execute()
)
email_details = parse_email(email_data)
if email_details:
emails.append(email_details)
except HttpError as e:
print(f"Error reading email {msg['id']}: {e}")
return json.dumps({"emails": emails})
except HttpError as e:
raise ToolExecutionError(
"Error searching emails",
developer_message=f"Gmail API Error: {e}",
)
@tool(
requires_auth=Google(
scope=["https://www.googleapis.com/auth/gmail.readonly"],
@ -94,96 +212,3 @@ async def get_emails(
except Exception as e:
print(f"Error reading emails: {e}")
return "Error reading emails"
def parse_email(email_data: Dict[str, Any]) -> Optional[Dict[str, str]]:
"""
Parse email data and extract relevant information.
Args:
email_data (Dict[str, Any]): Raw email data from Gmail API.
Returns:
Optional[Dict[str, str]]: Parsed email details or None if parsing fails.
"""
try:
payload = email_data["payload"]
headers = {d["name"].lower(): d["value"] for d in payload["headers"]}
body_data = get_email_body(payload)
return {
"from": headers.get("from", ""),
"date": headers.get("date", ""),
"subject": headers.get("subject", "No subject"),
"body": clean_email_body(body_data) if body_data else "",
}
except Exception as e:
print(f"Error parsing email {email_data.get('id', 'unknown')}: {e}")
return None
def get_email_body(payload: Dict[str, Any]) -> Optional[str]:
"""
Extract email body from payload.
Args:
payload (Dict[str, Any]): Email payload data.
Returns:
Optional[str]: Decoded email body or None if not found.
"""
if "body" in payload and payload["body"].get("data"):
return urlsafe_b64decode(payload["body"]["data"]).decode()
for part in payload.get("parts", []):
if part.get("mimeType") == "text/plain" and "data" in part["body"]:
return urlsafe_b64decode(part["body"]["data"]).decode()
return None
def clean_email_body(body: str) -> str:
"""
Remove HTML tags and clean up email body text while preserving most content.
Args:
body (str): The raw email body text.
Returns:
str: Cleaned email body text.
"""
try:
# Remove HTML tags using BeautifulSoup
soup = BeautifulSoup(body, "html.parser")
text = soup.get_text(separator=" ")
# Clean up the text
text = clean_text(text)
return text.strip()
except Exception as e:
print(f"Error cleaning email body: {e}")
return body
def clean_text(text: str) -> str:
"""
Clean up the text while preserving most content.
Args:
text (str): The input text.
Returns:
str: Cleaned text.
"""
# Replace multiple newlines with a single newline
text = re.sub(r"\n+", "\n", text)
# Replace multiple spaces with a single space
text = re.sub(r"\s+", " ", text)
# Remove leading/trailing whitespace from each line
text = "\n".join(line.strip() for line in text.split("\n"))
return text

98
toolkits/gmail/arcade_gmail/tools/utils.py Normal file
View file

@ -0,0 +1,98 @@
from base64 import urlsafe_b64decode
import re
from typing import Any, Dict, Optional
from bs4 import BeautifulSoup
def parse_email(email_data: Dict[str, Any]) -> Optional[Dict[str, str]]:
"""
Parse email data and extract relevant information.
Args:
email_data (Dict[str, Any]): Raw email data from Gmail API.
Returns:
Optional[Dict[str, str]]: Parsed email details or None if parsing fails.
"""
try:
payload = email_data["payload"]
headers = {d["name"].lower(): d["value"] for d in payload["headers"]}
body_data = _get_email_body(payload)
return {
"from": headers.get("from", ""),
"date": headers.get("date", ""),
"subject": headers.get("subject", "No subject"),
"body": _clean_email_body(body_data) if body_data else "",
}
except Exception as e:
print(f"Error parsing email {email_data.get('id', 'unknown')}: {e}")
return None
def _get_email_body(payload: Dict[str, Any]) -> Optional[str]:
"""
Extract email body from payload.
Args:
payload (Dict[str, Any]): Email payload data.
Returns:
Optional[str]: Decoded email body or None if not found.
"""
if "body" in payload and payload["body"].get("data"):
return urlsafe_b64decode(payload["body"]["data"]).decode()
for part in payload.get("parts", []):
if part.get("mimeType") == "text/plain" and "data" in part["body"]:
return urlsafe_b64decode(part["body"]["data"]).decode()
return None
def _clean_email_body(body: str) -> str:
"""
Remove HTML tags and clean up email body text while preserving most content.
Args:
body (str): The raw email body text.
Returns:
str: Cleaned email body text.
"""
try:
# Remove HTML tags using BeautifulSoup
soup = BeautifulSoup(body, "html.parser")
text = soup.get_text(separator=" ")
# Clean up the text
text = _clean_text(text)
return text.strip()
except Exception as e:
print(f"Error cleaning email body: {e}")
return body
def _clean_text(text: str) -> str:
"""
Clean up the text while preserving most content.
Args:
text (str): The input text.
Returns:
str: Cleaned text.
"""
# Replace multiple newlines with a single newline
text = re.sub(r"\n+", "\n", text)
# Replace multiple spaces with a single space
text = re.sub(r"\s+", " ", text)
# Remove leading/trailing whitespace from each line
text = "\n".join(line.strip() for line in text.split("\n"))
return text