Alex gmail improvements (#188)

Improved gmail toolkit. Added support for threading in draft replies,
multipart email parsing, and label management. Fixed the DateRange
parameter issue in list_emails_by_headers. Added logging and removed
print statements. Created custom exceptions for each specific google
toolkit.

-----
Summary of changes by @byrro:

- Fixed minor bug related to the `date_range` argument of
`list_emails_by_header`
- A few utility functions (`build_email_message`,
`build_reply_recipients`, `build_reply_body`) to centralize logic and
remove repeated code from email-sending tools
- New `reply_to_email` tool (apart from `write_draft_reply_email`,
implemented by Alex) to keep the toolkit consistent
- Evals and unit tests
- Handling of reply-to (only sender) and reply-to-all recipients
- Removed some unnecessary debug messages, which Alex had added to
replace print statements
- Removed HTML handling implemented by Alex in `write_draft_reply_email`
> I think we should either support HTML across all applicable tools or
not at all; I decided to remove it and leave this feature for a future
PR.

---------

Co-authored-by: Renato Byrro <rmbyrro@gmail.com>
This commit is contained in:
Alex Salazar 2025-02-27 06:56:32 -08:00 committed by GitHub
parent f997ca9449
commit 7b1110f2b7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 1401 additions and 248 deletions

View file

@ -0,0 +1,18 @@
import os
from arcade_google.tools.models import GmailReplyToWhom
# The default reply in Gmail is to only the sender. Since Gmail also offers the possibility of
# changing the default to 'reply to all', we support both options through an env variable.
# https://support.google.com/mail/answer/6585?hl=en&sjid=15399867888091633568-SA#null
try:
GMAIL_DEFAULT_REPLY_TO = GmailReplyToWhom(
# Values accepted are defined in the arcade_google.tools.models.GmailReplyToWhom Enum
os.getenv("ARCADE_GMAIL_DEFAULT_REPLY_TO", GmailReplyToWhom.ONLY_THE_SENDER.value).lower()
)
except ValueError as e:
raise ValueError(
"Invalid value for ARCADE_GMAIL_DEFAULT_REPLY_TO: "
f"'{os.getenv('ARCADE_GMAIL_DEFAULT_REPLY_TO')}'. Expected one of "
f"{list(GmailReplyToWhom.__members__.keys())}"
) from e

View file

@ -0,0 +1,43 @@
class GoogleToolError(Exception):
"""Base exception for Google tool errors."""
def __init__(self, message: str, developer_message: str | None = None):
self.message = message
self.developer_message = developer_message
super().__init__(self.message)
def __str__(self) -> str:
base_message = self.message
if self.developer_message:
return f"{base_message} (Developer: {self.developer_message})"
return base_message
class GoogleServiceError(GoogleToolError):
"""Raised when there's an error building or using the Google service."""
pass
class GmailToolError(GoogleToolError):
"""Raised when there's an error in the Gmail tools."""
pass
class GoogleCalendarToolError(GoogleToolError):
"""Raised when there's an error in the Google Calendar tools."""
pass
class GoogleDriveToolError(GoogleToolError):
"""Raised when there's an error in the Google Drive tools."""
pass
class GoogleDocsToolError(GoogleToolError):
"""Raised when there's an error in the Google Docs tools."""
pass

View file

@ -1,24 +1,30 @@
import base64
from email.message import EmailMessage
from email.mime.text import MIMEText
from typing import Annotated, Any, Optional
from arcade.sdk import ToolContext, tool
from arcade.sdk.auth import Google
from arcade.sdk.errors import RetryableToolError
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from arcade_google.tools.constants import GMAIL_DEFAULT_REPLY_TO
from arcade_google.tools.exceptions import GmailToolError
from arcade_google.tools.models import GmailAction, GmailReplyToWhom
from arcade_google.tools.utils import (
DateRange,
build_query_string,
_build_gmail_service,
build_email_message,
build_gmail_query_string,
build_reply_recipients,
fetch_messages,
get_draft_url,
get_email_details,
get_email_in_trash_url,
get_label_ids,
get_sent_email_url,
parse_draft_email,
parse_email,
process_email_messages,
parse_multipart_email,
parse_plain_text_email,
remove_none_values,
)
@ -40,37 +46,12 @@ async def send_email(
"""
Send an email using the Gmail API.
"""
service = _build_gmail_service(context)
email = build_email_message(recipient, subject, body, cc, bcc)
# Set up the Gmail API client
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
message = EmailMessage()
message.set_content(body)
message["To"] = recipient
message["Subject"] = subject
if cc:
message["Cc"] = ", ".join(cc)
if bcc:
message["Bcc"] = ", ".join(bcc)
# Encode the message in base64
encoded_message = base64.urlsafe_b64encode(message.as_bytes()).decode()
# Create the email
email = {"raw": encoded_message}
# Send the email
sent_message = service.users().messages().send(userId="me", body=email).execute()
email = parse_email(sent_message)
email = parse_plain_text_email(sent_message)
email["url"] = get_sent_email_url(sent_message["id"])
return email
@ -87,21 +68,79 @@ async def send_draft_email(
Send a draft email using the Gmail API.
"""
# Set up the Gmail API client
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
# Send the draft email
sent_message = service.users().drafts().send(userId="me", body={"id": email_id}).execute()
email = parse_email(sent_message)
email = parse_plain_text_email(sent_message)
email["url"] = get_sent_email_url(sent_message["id"])
return email
# Note: in the Gmail UI, a user can customize the recipient and cc fields before replying.
# We decided not to support this feature, since we'd need a way for LLMs to tell apart between
# adding or removing recipients/cc, or replacing with an entirely new list of addresses,
# which would make the tool more complex to call.
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/gmail.send"],
)
)
async def reply_to_email(
context: ToolContext,
body: Annotated[str, "The body of the email"],
reply_to_message_id: Annotated[str, "The ID of the message to reply to"],
reply_to_whom: Annotated[
GmailReplyToWhom,
"Whether to reply to every recipient (including cc) or only to the original sender. "
f"Defaults to '{GMAIL_DEFAULT_REPLY_TO}'.",
] = GMAIL_DEFAULT_REPLY_TO,
bcc: Annotated[Optional[list[str]], "BCC recipients of the email"] = None,
) -> Annotated[dict, "A dictionary containing the sent email details"]:
"""
Send a reply to an email message.
"""
if isinstance(reply_to_whom, str):
reply_to_whom = GmailReplyToWhom(reply_to_whom)
service = _build_gmail_service(context)
current_user = service.users().getProfile(userId="me").execute()
try:
replying_to_email = (
service.users().messages().get(userId="me", id=reply_to_message_id).execute()
)
except HttpError as e:
raise RetryableToolError(
message=f"Could not retrieve the message with id {reply_to_message_id}.",
developer_message=(
f"Could not retrieve the message with id {reply_to_message_id}. "
f"Reason: '{e.reason}'. Error details: '{e.error_details}'"
),
) from e
replying_to_email = parse_multipart_email(replying_to_email)
recipients = build_reply_recipients(
replying_to_email, current_user["emailAddress"], reply_to_whom
)
email = build_email_message(
recipient=recipients,
subject=f"Re: {replying_to_email['subject']}",
body=body,
cc=None
if reply_to_whom == GmailReplyToWhom.ONLY_THE_SENDER
else replying_to_email["cc"].split(","),
bcc=bcc,
replying_to=replying_to_email,
)
sent_message = service.users().messages().send(userId="me", body=email).execute()
email = parse_plain_text_email(sent_message)
email["url"] = get_sent_email_url(sent_message["id"])
return email
@ -124,29 +163,11 @@ async def write_draft_email(
Compose a new email draft using the Gmail API.
"""
# Set up the Gmail API client
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
message = MIMEText(body)
message["to"] = recipient
message["subject"] = subject
if cc:
message["Cc"] = ", ".join(cc)
if bcc:
message["Bcc"] = ", ".join(bcc)
# Encode the message in base64
raw_message = base64.urlsafe_b64encode(message.as_bytes()).decode()
# Create the draft
draft = {"message": {"raw": raw_message}}
draft = {
"message": build_email_message(recipient, subject, body, cc, bcc, action=GmailAction.DRAFT)
}
draft_message = service.users().drafts().create(userId="me", body=draft).execute()
email = parse_draft_email(draft_message)
@ -154,6 +175,76 @@ async def write_draft_email(
return email
# Note: in the Gmail UI, a user can customize the recipient and cc fields before replying.
# We decided not to support this feature, since we'd need a way for LLMs to tell apart between
# adding or removing recipients/cc, or replacing with an entirely new list of addresses,
# which would make the tool more complex to call.
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/gmail.compose"],
)
)
async def write_draft_reply_email(
context: ToolContext,
body: Annotated[str, "The body of the draft reply email"],
reply_to_message_id: Annotated[str, "The Gmail message ID of the message to draft a reply to"],
reply_to_whom: Annotated[
GmailReplyToWhom,
"Whether to reply to every recipient (including cc) or only to the original sender. "
f"Defaults to '{GMAIL_DEFAULT_REPLY_TO}'.",
] = GMAIL_DEFAULT_REPLY_TO,
bcc: Annotated[Optional[list[str]], "BCC recipients of the draft reply email"] = None,
) -> Annotated[dict, "A dictionary containing the created draft reply email details"]:
"""
Compose a draft reply to an email message.
"""
if isinstance(reply_to_whom, str):
reply_to_whom = GmailReplyToWhom(reply_to_whom)
service = _build_gmail_service(context)
current_user = service.users().getProfile(userId="me").execute()
try:
replying_to_email = (
service.users().messages().get(userId="me", id=reply_to_message_id).execute()
)
except HttpError as e:
raise RetryableToolError(
message="Could not retrieve the message to respond to.",
developer_message=(
"Could not retrieve the message to respond to. "
f"Reason: '{e.reason}'. Error details: '{e.error_details}'"
),
)
replying_to_email = parse_multipart_email(replying_to_email)
recipients = build_reply_recipients(
replying_to_email, current_user["emailAddress"], reply_to_whom
)
draft_message = {
"message": build_email_message(
recipient=recipients,
subject=f"Re: {replying_to_email['subject']}",
body=body,
cc=None
if reply_to_whom == GmailReplyToWhom.ONLY_THE_SENDER
else replying_to_email["cc"].split(","),
bcc=bcc,
replying_to=replying_to_email,
action=GmailAction.DRAFT,
),
}
draft = service.users().drafts().create(userId="me", body=draft_message).execute()
email = parse_draft_email(draft)
email["url"] = get_draft_url(draft["id"])
return email
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/gmail.compose"],
@ -171,17 +262,7 @@ async def update_draft_email(
"""
Update an existing email draft using the Gmail API.
"""
# Set up the Gmail API client
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
message = MIMEText(body)
message["to"] = recipient
@ -203,6 +284,7 @@ async def update_draft_email(
email = parse_draft_email(updated_draft_message)
email["url"] = get_draft_url(updated_draft_message["id"])
return email
@ -218,17 +300,7 @@ async def delete_draft_email(
"""
Delete a draft email using the Gmail API.
"""
# Set up the Gmail API client
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
# Delete the draft
service.users().drafts().delete(userId="me", id=draft_email_id).execute()
@ -248,21 +320,12 @@ async def trash_email(
Move an email to the trash folder using the Gmail API.
"""
# Set up the Gmail API client
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
# Trash the email
trashed_email = service.users().messages().trash(userId="me", id=email_id).execute()
email = parse_email(trashed_email)
email = parse_plain_text_email(trashed_email)
email["url"] = get_email_in_trash_url(trashed_email["id"])
return email
@ -280,16 +343,7 @@ async def list_draft_emails(
"""
Lists draft emails in the user's draft mailbox using the Gmail API.
"""
# Set up the Gmail API client
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
listed_drafts = service.users().drafts().list(userId="me").execute()
@ -306,12 +360,13 @@ async def list_draft_emails(
if draft_details:
emails.append(draft_details)
except Exception as e:
print(f"Error reading draft email {draft_id}: {e}")
raise GmailToolError(
message=f"Error reading draft email {draft_id}.", developer_message=str(e)
)
return {"emails": emails}
# Email Search Tools
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/gmail.readonly"],
@ -324,39 +379,58 @@ async def list_emails_by_header(
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[int, "The maximum number of emails to return"] = 25,
label: Annotated[Optional[str], "The label name to filter by"] = None,
max_results: Annotated[int, "The maximum number of emails to return"] = 25,
) -> Annotated[
dict, "A dictionary containing a list of email details matching the search criteria"
]:
"""
Search for emails by header using the Gmail API.
At least one of the following parameters MUST be provided: sender, recipient, subject, body.
At least one of the following parameters MUST be provided: sender, recipient,
subject, date_range, label, or body.
"""
if not any([sender, recipient, subject, body]):
service = _build_gmail_service(context)
# Ensure at least one search parameter is provided
if not any([sender, recipient, subject, body, label, date_range]):
raise RetryableToolError(
message="At least one of sender, recipient, subject, or body must be provided.",
message=(
"At least one of sender, recipient, subject, body, label, query, "
"or date_range must be provided."
),
developer_message=(
"At least one of sender, recipient, subject, or body must be provided."
"At least one of sender, recipient, subject, body, label, query, "
"or date_range must be provided."
),
)
query = build_query_string(sender, recipient, subject, body, date_range)
# Check if label is valid
if label:
label_ids = get_label_ids(service, [label])
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
messages = fetch_messages(service, query, limit)
if not label_ids:
labels = service.users().labels().list(userId="me").execute().get("labels", [])
label_names = [label["name"] for label in labels]
raise RetryableToolError(
message=f"Invalid label: {label}",
developer_message=f"Invalid label: {label}",
additional_prompt_content=f"List of valid labels: {label_names}",
)
# Build a Gmail-style query string based on the filters
query = build_gmail_query_string(sender, recipient, subject, body, date_range, label)
# Fetch matching messages. This fetches message metadata from Gmail
messages = fetch_messages(service, query, max_results)
# If no messages found, return an empty list
if not messages:
return {"emails": []}
emails = process_email_messages(service, messages)
# Process each message into a structured email object
emails = get_email_details(service, messages)
# Return the list of emails in a dictionary with key "emails"
return {"emails": emails}
@ -372,16 +446,7 @@ async def list_emails(
"""
Read emails from a Gmail account and extract plain text content.
"""
# Set up the Gmail API client
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
messages = service.users().messages().list(userId="me").execute().get("messages", [])
@ -392,12 +457,13 @@ async def list_emails(
for msg in messages[:n_emails]:
try:
email_data = service.users().messages().get(userId="me", id=msg["id"]).execute()
email_details = parse_email(email_data)
email_details = parse_plain_text_email(email_data)
if email_details:
emails.append(email_details)
except Exception as e:
print(f"Error reading email {msg['id']}: {e}")
raise GmailToolError(
message=f"Error reading email {msg['id']}.", developer_message=str(e)
)
return {"emails": emails}
@ -421,18 +487,10 @@ async def search_threads(
date_range: Annotated[Optional[DateRange], "The date range of the email"] = None,
) -> Annotated[dict, "A dictionary containing a list of thread details"]:
"""Search for threads in the user's mailbox"""
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
query = (
build_query_string(sender, recipient, subject, body, date_range)
build_gmail_query_string(sender, recipient, subject, body, date_range)
if any([sender, recipient, subject, body, date_range])
else None
)
@ -506,16 +564,101 @@ async def get_thread(
}
params = remove_none_values(params)
service = build(
"gmail",
"v1",
credentials=Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
),
)
service = _build_gmail_service(context)
thread = service.users().threads().get(**params).execute()
thread["messages"] = [parse_email(message) for message in thread.get("messages", [])]
thread["messages"] = [parse_plain_text_email(message) for message in thread.get("messages", [])]
return dict(thread)
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/gmail.modify"],
)
)
async def change_email_labels(
context: ToolContext,
email_id: Annotated[str, "The ID of the email to modify labels for"],
labels_to_add: Annotated[list[str], "List of label names to add"],
labels_to_remove: Annotated[list[str], "List of label names to remove"],
) -> Annotated[dict, "List of labels that were added, removed, and not found"]:
"""
Add and remove labels from an email using the Gmail API.
"""
service = _build_gmail_service(context)
add_labels = get_label_ids(service, labels_to_add)
remove_labels = get_label_ids(service, labels_to_remove)
invalid_labels = (
set(labels_to_add + labels_to_remove) - set(add_labels.keys()) - set(remove_labels.keys())
)
if invalid_labels:
# prepare the list of valid labels
labels = service.users().labels().list(userId="me").execute().get("labels", [])
label_names = [label["name"] for label in labels]
# raise a retryable error with the list of valid labels
raise RetryableToolError(
message=f"Invalid labels: {invalid_labels}",
developer_message=f"Invalid labels: {invalid_labels}",
additional_prompt_content=f"List of valid labels: {label_names}",
)
# Prepare the modification body with label IDs.
body = {
"addLabelIds": list(add_labels.values()),
"removeLabelIds": list(remove_labels.values()),
}
try: # Modify the email labels.
service.users().messages().modify(userId="me", id=email_id, body=body).execute()
except Exception as e:
raise GmailToolError(
message=f"Error modifying labels for email {email_id}", developer_message=str(e)
)
# Confirmation JSON with lists for added and removed labels.
confirmation = {
"addedLabels": list(add_labels.keys()),
"removedLabels": list(remove_labels.keys()),
}
return {"confirmation": dict(confirmation)}
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/gmail.readonly"],
)
)
async def list_labels(
context: ToolContext,
) -> Annotated[dict, "A dictionary containing a list of label details"]:
"""List all the labels in the user's mailbox."""
service = _build_gmail_service(context)
labels = service.users().labels().list(userId="me").execute().get("labels", [])
return {"labels": labels}
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/gmail.labels"],
)
)
async def create_label(
context: ToolContext,
label_name: Annotated[str, "The name of the label to create"],
) -> Annotated[dict, "The details of the created label"]:
"""Create a new label in the user's mailbox."""
service = _build_gmail_service(context)
label = service.users().labels().create(userId="me", body={"name": label_name}).execute()
return {"label": label}

View file

@ -342,3 +342,16 @@ class OrderBy(str, Enum):
# The last time the file was viewed by the user (descending)
"viewedByMeTime desc"
)
# ---------------------------------------------------------------------------- #
# Google Gmail Models and Enums
# ---------------------------------------------------------------------------- #
class GmailReplyToWhom(str, Enum):
EVERY_RECIPIENT = "every_recipient"
ONLY_THE_SENDER = "only_the_sender"
class GmailAction(str, Enum):
SEND = "send"
DRAFT = "draft"

View file

@ -1,16 +1,28 @@
import logging
import re
from base64 import urlsafe_b64decode
from base64 import urlsafe_b64decode, urlsafe_b64encode
from datetime import datetime, timedelta
from email.message import EmailMessage
from email.mime.text import MIMEText
from enum import Enum
from typing import Any, Optional
from typing import Any, Optional, Union
from zoneinfo import ZoneInfo
from arcade.sdk import ToolContext
from bs4 import BeautifulSoup
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import Resource, build
from googleapiclient.errors import HttpError
from arcade_google.tools.models import Day, TimeSlot
from arcade_google.tools.exceptions import GmailToolError, GoogleServiceError
from arcade_google.tools.models import Day, GmailAction, GmailReplyToWhom, TimeSlot
## Set up basic configuration for logging to the console with DEBUG level and a specific format.
logging.basicConfig(
level=logging.DEBUG,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
def parse_datetime(datetime_str: str, time_zone: str) -> datetime:
@ -73,21 +85,77 @@ class DateRange(Enum):
return result + comparison_date.strftime("%Y/%m/%d")
def process_email_messages(service: Any, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
emails = []
for msg in messages:
try:
email_data = service.users().messages().get(userId="me", id=msg["id"]).execute()
email_details = parse_email(email_data)
emails += [email_details] if email_details else []
except HttpError as e:
print(f"Error reading email {msg['id']}: {e}")
return emails
def build_email_message(
recipient: str,
subject: str,
body: str,
cc: Optional[list[str]] = None,
bcc: Optional[list[str]] = None,
replying_to: Optional[dict[str, Any]] = None,
action: GmailAction = GmailAction.SEND,
) -> dict[str, Any]:
if replying_to:
body = build_reply_body(body, replying_to)
message: Union[EmailMessage, MIMEText]
if action == GmailAction.SEND:
message = EmailMessage()
message.set_content(body)
elif action == GmailAction.DRAFT:
message = MIMEText(body)
message["To"] = recipient
message["Subject"] = subject
if cc:
message["Cc"] = ",".join(cc)
if bcc:
message["Bcc"] = ",".join(bcc)
if replying_to:
message["In-Reply-To"] = replying_to["header_message_id"]
message["References"] = f"{replying_to['header_message_id']}, {replying_to['references']}"
encoded_message = urlsafe_b64encode(message.as_bytes()).decode()
data = {"raw": encoded_message}
if replying_to:
data["threadId"] = replying_to["thread_id"]
return data
def parse_email(email_data: dict[str, Any]) -> dict[str, Any]:
def build_reply_body(body: str, replying_to: dict[str, Any]) -> str:
attribution = f"On {replying_to['date']}, {replying_to['from']} wrote:"
lines = replying_to["plain_text_body"].split("\n")
quoted_plain = "\n".join([f"> {line}" for line in lines])
return f"{body}\n\n{attribution}\n\n{quoted_plain}"
def build_reply_recipients(
replying_to: dict[str, Any], current_user_email_address: str, reply_to_whom: GmailReplyToWhom
) -> str:
if reply_to_whom == GmailReplyToWhom.ONLY_THE_SENDER:
recipients = [replying_to["from"]]
elif reply_to_whom == GmailReplyToWhom.EVERY_RECIPIENT:
recipients = [replying_to["from"], *replying_to["to"].split(",")]
else:
raise ValueError(f"Unsupported reply_to_whom value: {reply_to_whom}")
recipients = [
email_address.strip()
for email_address in recipients
if email_address.strip().lower() != current_user_email_address.lower().strip()
]
return ", ".join(recipients)
def parse_plain_text_email(email_data: dict[str, Any]) -> dict[str, Any]:
"""
Parse email data and extract relevant information.
Only returns the plain text body.
Args:
email_data (Dict[str, Any]): Raw email data from Gmail API.
@ -95,23 +163,71 @@ def parse_email(email_data: dict[str, Any]) -> dict[str, Any]:
Returns:
Optional[Dict[str, str]]: Parsed email details or None if parsing fails.
"""
try:
payload = email_data.get("payload", {})
headers = {d["name"].lower(): d["value"] for d in payload.get("headers", [])}
payload = email_data.get("payload", {})
headers = {d["name"].lower(): d["value"] for d in payload.get("headers", [])}
body_data = _get_email_body(payload)
body_data = _get_email_plain_text_body(payload)
return {
"id": email_data.get("id", ""),
"thread_id": email_data.get("threadId", ""),
"from": headers.get("from", ""),
"date": headers.get("date", ""),
"subject": headers.get("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 email_data
email_details = {
"id": email_data.get("id", ""),
"thread_id": email_data.get("threadId", ""),
"label_ids": email_data.get("labelIds", []),
"history_id": email_data.get("historyId", ""),
"snippet": email_data.get("snippet", ""),
"to": headers.get("to", ""),
"cc": headers.get("cc", ""),
"from": headers.get("from", ""),
"reply_to": headers.get("reply-to", ""),
"in_reply_to": headers.get("in-reply-to", ""),
"references": headers.get("references", ""),
"header_message_id": headers.get("message-id", ""),
"date": headers.get("date", ""),
"subject": headers.get("subject", ""),
"body": body_data or "",
}
return email_details
def parse_multipart_email(email_data: dict[str, Any]) -> dict[str, Any]:
"""
Parse email data and extract relevant information.
Returns the plain text and HTML body along with the images.
Args:
email_data (Dict[str, Any]): Raw email data from Gmail API.
Returns:
Optional[Dict[str, Any]]: Parsed email details or None if parsing fails.
"""
payload = email_data.get("payload", {})
headers = {d["name"].lower(): d["value"] for d in payload.get("headers", [])}
# Extract different parts of the email
plain_text_body = _get_email_plain_text_body(payload)
html_body = _get_email_html_body(payload)
email_details = {
"id": email_data.get("id", ""),
"thread_id": email_data.get("threadId", ""),
"label_ids": email_data.get("labelIds", []),
"history_id": email_data.get("historyId", ""),
"snippet": email_data.get("snippet", ""),
"to": headers.get("to", ""),
"cc": headers.get("cc", ""),
"from": headers.get("from", ""),
"reply_to": headers.get("reply-to", ""),
"in_reply_to": headers.get("in-reply-to", ""),
"references": headers.get("references", ""),
"header_message_id": headers.get("message-id", ""),
"date": headers.get("date", ""),
"subject": headers.get("subject", ""),
"plain_text_body": plain_text_body or _clean_email_body(html_body),
"html_body": html_body or "",
}
return email_details
def parse_draft_email(draft_email_data: dict[str, Any]) -> dict[str, str]:
@ -124,24 +240,20 @@ def parse_draft_email(draft_email_data: dict[str, Any]) -> dict[str, str]:
Returns:
Optional[Dict[str, str]]: Parsed draft email details or None if parsing fails.
"""
try:
message = draft_email_data.get("message", {})
payload = message.get("payload", {})
headers = {d["name"].lower(): d["value"] for d in payload.get("headers", [])}
message = draft_email_data.get("message", {})
payload = message.get("payload", {})
headers = {d["name"].lower(): d["value"] for d in payload.get("headers", [])}
body_data = _get_email_body(payload)
body_data = _get_email_plain_text_body(payload)
return {
"id": draft_email_data.get("id", ""),
"thread_id": draft_email_data.get("threadId", ""),
"from": headers.get("from", ""),
"date": headers.get("internaldate", ""),
"subject": headers.get("subject", ""),
"body": _clean_email_body(body_data) if body_data else "",
}
except Exception as e:
print(f"Error parsing draft email {draft_email_data.get('id', 'unknown')}: {e}")
return draft_email_data
return {
"id": draft_email_data.get("id", ""),
"thread_id": draft_email_data.get("threadId", ""),
"from": headers.get("from", ""),
"date": headers.get("internaldate", ""),
"subject": headers.get("subject", ""),
"body": _clean_email_body(body_data) if body_data else "",
}
def get_draft_url(draft_id: str) -> str:
@ -152,13 +264,143 @@ def get_sent_email_url(sent_email_id: str) -> str:
return f"https://mail.google.com/mail/u/0/#sent/{sent_email_id}"
def get_email_details(service: Any, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""
Retrieves full message data for each message ID in the given list and extracts email details.
:param service: Authenticated Gmail API service instance.
:param messages: A list of dictionaries, each representing a message with an 'id' key.
:return: A list of dictionaries, each containing parsed email details.
"""
emails = []
for msg in messages:
try:
# Fetch the full message data from Gmail using the message ID
email_data = service.users().messages().get(userId="me", id=msg["id"]).execute()
# Parse the raw email data into a structured form
email_details = parse_plain_text_email(email_data)
# Only add the details if parsing was successful
if email_details:
emails.append(email_details)
except Exception as e:
# Log any errors encountered while trying to fetch or parse a message
raise GmailToolError(
message=f"Error reading email {msg['id']}.", developer_message=str(e)
)
return emails
def get_email_in_trash_url(email_id: str) -> str:
return f"https://mail.google.com/mail/u/0/#trash/{email_id}"
def _get_email_body(payload: dict[str, Any]) -> Optional[str]:
def _build_gmail_service(context: ToolContext) -> Any:
"""
Extract email body from payload.
Private helper function to build and return the Gmail service client.
Args:
context (ToolContext): The context containing authorization details.
Returns:
googleapiclient.discovery.Resource: An authorized Gmail API service instance.
"""
try:
credentials = Credentials(
context.authorization.token
if context.authorization and context.authorization.token
else ""
)
except Exception as e:
raise GoogleServiceError(message="Failed to build Gmail service.", developer_message=str(e))
return build("gmail", "v1", credentials=credentials)
def _extract_plain_body(parts: list) -> Optional[str]:
"""
Recursively extract the email body from parts, handling both plain text and HTML.
Args:
parts (List[Dict[str, Any]]): List of email parts.
Returns:
Optional[str]: Decoded and cleaned email body or None if not found.
"""
for part in parts:
mime_type = part.get("mimeType")
if mime_type == "text/plain" and "data" in part.get("body", {}):
return urlsafe_b64decode(part["body"]["data"]).decode()
elif mime_type.startswith("multipart/"):
subparts = part.get("parts", [])
body = _extract_plain_body(subparts)
if body:
return body
return _extract_html_body(parts)
def _extract_html_body(parts: list) -> Optional[str]:
"""
Recursively extract the email body from parts, handling only HTML.
Args:
parts (List[Dict[str, Any]]): List of email parts.
Returns:
Optional[str]: Decoded and cleaned email body or None if not found.
"""
for part in parts:
mime_type = part.get("mimeType")
if mime_type == "text/html" and "data" in part.get("body", {}):
html_content = urlsafe_b64decode(part["body"]["data"]).decode()
return html_content
elif mime_type.startswith("multipart/"):
subparts = part.get("parts", [])
body = _extract_html_body(subparts)
if body:
return body
return None
def _get_email_images(payload: dict[str, Any]) -> Optional[list[str]]:
"""
Extract the email images from an email payload.
Args:
payload (Dict[str, Any]): Email payload data.
Returns:
Optional[List[str]]: List of decoded image contents or None if none found.
"""
images = []
for part in payload.get("parts", []):
mime_type = part.get("mimeType")
if mime_type.startswith("image/") and "data" in part.get("body", {}):
image_content = part["body"]["data"]
images.append(image_content)
elif mime_type.startswith("multipart/"):
subparts = part.get("parts", [])
subimages = _get_email_images(subparts)
if subimages:
images.extend(subimages)
if images:
return images
return None
def _get_email_plain_text_body(payload: dict[str, Any]) -> Optional[str]:
"""
Extract email body from payload, handling 'multipart/alternative' parts.
Args:
payload (Dict[str, Any]): Email payload data.
@ -166,17 +408,33 @@ def _get_email_body(payload: dict[str, Any]) -> Optional[str]:
Returns:
Optional[str]: Decoded email body or None if not found.
"""
# Direct body extraction
if "body" in payload and payload["body"].get("data"):
return _clean_email_body(urlsafe_b64decode(payload["body"]["data"]).decode())
# Handle multipart and alternative parts
return _clean_email_body(_extract_plain_body(payload.get("parts", [])))
def _get_email_html_body(payload: dict[str, Any]) -> Optional[str]:
"""
Extract email html body from payload, handling 'multipart/alternative' parts.
Args:
payload (Dict[str, Any]): Email payload data.
Returns:
Optional[str]: Decoded email body or None if not found.
"""
# Direct body extraction
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
# Handle multipart and alternative parts
return _extract_html_body(payload.get("parts", []))
def _clean_email_body(body: str) -> str:
def _clean_email_body(body: Optional[str]) -> str:
"""
Remove HTML tags and clean up email body text while preserving most content.
@ -186,6 +444,9 @@ def _clean_email_body(body: str) -> str:
Returns:
str: Cleaned email body text.
"""
if not body:
return ""
try:
# Remove HTML tags using BeautifulSoup
soup = BeautifulSoup(body, "html.parser")
@ -195,8 +456,8 @@ def _clean_email_body(body: str) -> str:
cleaned_text = _clean_text(text)
return cleaned_text.strip()
except Exception as e:
print(f"Error cleaning email body: {e}")
except Exception:
logger.exception("Error cleaning email body")
return body
@ -240,12 +501,13 @@ def _update_datetime(day: Day | None, time: TimeSlot | None, time_zone: str) ->
return None
def build_query_string(
sender: str | None,
recipient: str | None,
subject: str | None,
body: str | None,
date_range: DateRange | None,
def build_gmail_query_string(
sender: str | None = None,
recipient: str | None = None,
subject: str | None = None,
body: str | None = None,
date_range: DateRange | None = None,
label: str | None = None,
) -> str:
"""Helper function to build a query string
for Gmail list_emails_by_header and search_threads tools.
@ -261,9 +523,43 @@ def build_query_string(
query.append(body)
if date_range:
query.append(date_range.to_date_query())
if label:
query.append(f"label:{label}")
return " ".join(query)
def get_label_ids(service: Any, label_names: list[str]) -> dict[str, str]:
"""
Retrieve label IDs for given label names.
Returns a dictionary mapping label names to their IDs.
Args:
service: Authenticated Gmail API service instance.
label_names: List of label names to retrieve IDs for.
Returns:
A dictionary mapping found label names to their corresponding IDs.
"""
try:
# Fetch all existing labels from Gmail
labels = service.users().labels().list(userId="me").execute().get("labels", [])
except Exception as e:
raise GmailToolError(message="Failed to list labels.", developer_message=str(e)) from e
# Create a mapping from label names to their IDs
label_id_map = {label["name"]: label["id"] for label in labels}
found_labels = {}
for name in label_names:
label_id = label_id_map.get(name)
if label_id:
found_labels[name] = label_id
else:
logger.warning(f"Label '{name}' does not exist")
return found_labels
def fetch_messages(service: Any, query_string: str, limit: int) -> list[dict[str, Any]]:
"""
Helper function to fetch messages from Gmail API for the list_emails_by_header tool.

View file

@ -1,3 +1,5 @@
import json
from arcade.sdk import ToolCatalog
from arcade.sdk.eval import (
BinaryCritic,
@ -11,10 +13,14 @@ from arcade.sdk.eval import (
import arcade_google
from arcade_google.tools.gmail import (
get_thread,
list_emails_by_header,
list_threads,
reply_to_email,
search_threads,
send_email,
write_draft_reply_email,
)
from arcade_google.tools.models import GmailReplyToWhom
from arcade_google.tools.utils import DateRange
# Evaluation rubric
@ -159,3 +165,267 @@ def gmail_eval_suite() -> EvalSuite:
)
return suite
@tool_eval()
def gmail_reply_eval_suite() -> EvalSuite:
"""Create an evaluation suite for Gmail reply tools."""
suite = EvalSuite(
name="Gmail Reply Tools Evaluation",
system_message="You are an AI assistant that can send and manage emails using the provided tools.",
catalog=catalog,
rubric=rubric,
)
email_history = [
{"role": "user", "content": "get the latest emails I received from johndoe@gmail.com"},
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call_jowMD7aB9sVPClOfvNof7Llu",
"type": "function",
"function": {
"name": "Google_ListEmailsByHeader",
"arguments": json.dumps({
"sender": "johndoe@gmail.com",
"max_results": 5,
}),
},
}
],
},
{
"role": "tool",
"content": json.dumps({
"emails": [
{
"body": "test 1",
"cc": "",
"date": "Tue, 11 Feb 2025 11:33:08 -0300",
"from": "John Doe <johndoe@gmail.com>",
"header_message_id": "<cierty475cty7245yxq@mail.gmail.com>",
"history_id": "123456",
"id": "q34759q435nv",
"in_reply_to": "",
"label_ids": ["INBOX"],
"references": "",
"reply_to": "",
"snippet": "test 1",
"subject": "test 1",
"thread_id": "345y6v3596",
"to": "myself@gmail.com",
},
{
"body": "test 2",
"cc": "",
"date": "Mon, 20 Jan 2025 13:04:42 -0800",
"from": "John Doe <johndoe@gmail.com>",
"header_message_id": "<28745ytvw8745ct4@mail.gmail.com>",
"history_id": "3456758",
"id": "9475tvy24578yx",
"in_reply_to": "",
"label_ids": [],
"references": "",
"reply_to": "",
"snippet": "test 2",
"subject": "test 2",
"thread_id": "249576v3496",
"to": "myself@gmail.com",
},
]
}),
"tool_call_id": "call_jowMD7aB9sVPClOfvNof7Llu",
"name": "Google_ListEmailsByHeader",
},
{
"role": "assistant",
"content": "Here are the latest emails you received from johndoe@gmail.com:\n\n1. **Subject**: test 1\n - **Date**: Tue, 11 Feb 2025 11:33:08 -0300\n - **Snippet**: test 1\n\n2. **Subject**: test 2\n - **Date**: Mon, 20 Jan 2025 13:04:42 -0800\n - **Snippet**: test 2\n\nIf you need further details from any specific email, let me know!",
},
]
suite.add_case(
name="Reply to an email",
user_message="Reply to the email from johndoe@example.com about 'test 2' saying 'tested and working well'",
expected_tool_calls=[
ExpectedToolCall(
func=reply_to_email,
args={
"reply_to_message_id": "9475tvy24578yx",
"body": "tested and working well",
"reply_to_whom": GmailReplyToWhom.ONLY_THE_SENDER.value,
},
)
],
critics=[
SimilarityCritic(critic_field="subject", weight=1 / 7),
SimilarityCritic(critic_field="body", weight=1 / 7),
BinaryCritic(critic_field="recipient", weight=1 / 7),
BinaryCritic(critic_field="cc", weight=1 / 7),
BinaryCritic(critic_field="bcc", weight=1 / 7),
BinaryCritic(critic_field="reply_to_whom", weight=1 / 7),
BinaryCritic(critic_field="reply_to_message_id", weight=1 / 7),
],
additional_messages=email_history,
)
suite.add_case(
name="Reply to an email with every recipient",
user_message="Reply to every recipient in the email from johndoe@example.com about 'test 2' saying 'tested and working well'",
expected_tool_calls=[
ExpectedToolCall(
func=reply_to_email,
args={
"reply_to_message_id": "9475tvy24578yx",
"body": "tested and working well",
"reply_to_whom": GmailReplyToWhom.EVERY_RECIPIENT.value,
},
)
],
critics=[
SimilarityCritic(critic_field="subject", weight=1 / 7),
SimilarityCritic(critic_field="body", weight=1 / 7),
BinaryCritic(critic_field="recipient", weight=1 / 7),
BinaryCritic(critic_field="cc", weight=1 / 7),
BinaryCritic(critic_field="bcc", weight=1 / 7),
BinaryCritic(critic_field="reply_to_whom", weight=1 / 7),
BinaryCritic(critic_field="reply_to_message_id", weight=1 / 7),
],
additional_messages=email_history,
)
suite.add_case(
name="Reply to an email with bcc",
user_message="Reply to the email from johndoe@example.com about 'test 2' saying 'tested and working well' and send it to janedoe@example.com as bcc as well",
expected_tool_calls=[
ExpectedToolCall(
func=reply_to_email,
args={
"reply_to_message_id": "9475tvy24578yx",
"body": "tested and working well",
"bcc": ["janedoe@example.com"],
"reply_to_whom": GmailReplyToWhom.ONLY_THE_SENDER.value,
},
)
],
critics=[
SimilarityCritic(critic_field="subject", weight=1 / 7),
SimilarityCritic(critic_field="body", weight=1 / 7),
BinaryCritic(critic_field="recipient", weight=1 / 7),
BinaryCritic(critic_field="cc", weight=1 / 7),
BinaryCritic(critic_field="bcc", weight=1 / 7),
BinaryCritic(critic_field="reply_to_whom", weight=1 / 7),
BinaryCritic(critic_field="reply_to_message_id", weight=1 / 7),
],
additional_messages=email_history,
)
suite.add_case(
name="Write draft reply",
user_message="Write a draft reply to the email from johndoe@example.com about 'test 2' saying 'tested and working well'",
expected_tool_calls=[
ExpectedToolCall(
func=write_draft_reply_email,
args={
"reply_to_message_id": "9475tvy24578yx",
"body": "tested and working well",
"reply_to_whom": GmailReplyToWhom.ONLY_THE_SENDER.value,
},
)
],
critics=[
SimilarityCritic(critic_field="subject", weight=1 / 7),
SimilarityCritic(critic_field="body", weight=1 / 7),
BinaryCritic(critic_field="recipient", weight=1 / 7),
BinaryCritic(critic_field="cc", weight=1 / 7),
BinaryCritic(critic_field="bcc", weight=1 / 7),
BinaryCritic(critic_field="reply_to_message_id", weight=1 / 7),
BinaryCritic(critic_field="reply_to_whom", weight=1 / 7),
],
additional_messages=email_history,
)
suite.add_case(
name="Write draft reply to every recipient",
user_message="Write a draft reply to every recipient in the email from johndoe@example.com about 'test 2' saying 'tested and working well'",
expected_tool_calls=[
ExpectedToolCall(
func=write_draft_reply_email,
args={
"reply_to_message_id": "9475tvy24578yx",
"body": "tested and working well",
"reply_to_whom": GmailReplyToWhom.EVERY_RECIPIENT.value,
},
)
],
critics=[
SimilarityCritic(critic_field="subject", weight=1 / 7),
SimilarityCritic(critic_field="body", weight=1 / 7),
BinaryCritic(critic_field="recipient", weight=1 / 7),
BinaryCritic(critic_field="cc", weight=1 / 7),
BinaryCritic(critic_field="bcc", weight=1 / 7),
BinaryCritic(critic_field="reply_to_whom", weight=0.125),
BinaryCritic(critic_field="reply_to_message_id", weight=1 / 7),
],
additional_messages=email_history,
)
return suite
@tool_eval()
def gmail_list_emails_by_header_eval_suite() -> EvalSuite:
"""Create an evaluation suite for Gmail tools."""
suite = EvalSuite(
name="Gmail list_emails_by_header tool evaluation",
system_message="You are an AI assistant that can send and manage emails using the provided tools.",
catalog=catalog,
rubric=rubric,
)
suite.add_case(
name="List emails by header using date-range",
user_message="List all emails from johndoe@example.com to janedoe@example.com about 'Arcade AI' from yesterday",
expected_tool_calls=[
ExpectedToolCall(
func=list_emails_by_header,
args={
"sender": "johndoe@example.com",
"recipient": "janedoe@example.com",
"subject": "Arcade AI",
"date_range": DateRange.YESTERDAY.value,
},
)
],
critics=[
BinaryCritic(critic_field="sender", weight=1 / 4),
BinaryCritic(critic_field="recipient", weight=1 / 4),
SimilarityCritic(critic_field="subject", weight=1 / 4),
BinaryCritic(critic_field="date_range", weight=1 / 4),
],
)
suite.add_case(
name="List emails by header using date-range",
user_message="List all emails from johndoe@example.com to janedoe@example.com about 'Arcade AI' from the last month",
expected_tool_calls=[
ExpectedToolCall(
func=list_emails_by_header,
args={
"sender": "johndoe@example.com",
"recipient": "janedoe@example.com",
"subject": "Arcade AI",
"date_range": DateRange.LAST_MONTH.value,
},
)
],
critics=[
BinaryCritic(critic_field="sender", weight=1 / 4),
BinaryCritic(critic_field="recipient", weight=1 / 4),
SimilarityCritic(critic_field="subject", weight=1 / 4),
BinaryCritic(critic_field="date_range", weight=1 / 4),
],
)
return suite

View file

@ -1,3 +1,5 @@
from base64 import urlsafe_b64encode
from email.message import EmailMessage
from unittest.mock import MagicMock, patch
import pytest
@ -12,6 +14,7 @@ from arcade_google.tools.gmail import (
list_emails,
list_emails_by_header,
list_threads,
reply_to_email,
search_threads,
send_draft_email,
send_email,
@ -19,7 +22,13 @@ from arcade_google.tools.gmail import (
update_draft_email,
write_draft_email,
)
from arcade_google.tools.utils import parse_draft_email, parse_email
from arcade_google.tools.models import GmailReplyToWhom
from arcade_google.tools.utils import (
build_reply_body,
parse_draft_email,
parse_multipart_email,
parse_plain_text_email,
)
@pytest.fixture
@ -29,7 +38,7 @@ def mock_context():
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_send_email(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -64,7 +73,7 @@ async def test_send_email(mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_write_draft_email(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -99,7 +108,7 @@ async def test_write_draft_email(mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_update_draft_email(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -136,7 +145,7 @@ async def test_update_draft_email(mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_send_draft_email(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -161,7 +170,7 @@ async def test_send_draft_email(mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_delete_draft_email(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -183,7 +192,7 @@ async def test_delete_draft_email(mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
@patch("arcade_google.tools.gmail.parse_draft_email")
async def test_get_draft_emails(mock_parse_draft_email, mock_build, mock_context):
# Setup test data
@ -234,7 +243,7 @@ async def test_get_draft_emails(mock_parse_draft_email, mock_build, mock_context
# Mock the response from the Gmail get drafts API
mock_service.users().drafts().get().execute.return_value = mock_drafts_get_response
# Mock the parse_email function since parse_email doesn't accept object of type MagicMock
# Mock the parse_draft_email function since parse_draft_email doesn't accept object of type MagicMock
mock_parse_draft_email.return_value = parse_draft_email(mock_drafts_get_response)
# Test happy path
@ -256,9 +265,9 @@ async def test_get_draft_emails(mock_parse_draft_email, mock_build, mock_context
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.gmail.parse_email")
async def test_search_emails_by_header(mock_parse_email, mock_build, mock_context):
@patch("arcade_google.tools.utils.build")
@patch("arcade_google.tools.gmail.parse_plain_text_email")
async def test_search_emails_by_header(mock_parse_plain_text_email, mock_build, mock_context):
# Setup test data
mock_messages_list_response = {
"messages": [
@ -307,11 +316,13 @@ async def test_search_emails_by_header(mock_parse_email, mock_build, mock_contex
# Mock the response from the Gmail get messages API
mock_service.users().messages().get().execute.return_value = mock_messages_get_response
# Mock the parse_email function since parse_email doesn't accept object of type MagicMock
mock_parse_email.return_value = parse_email(mock_messages_get_response)
# Mock the parse_plain_text_email function since parse_plain_text_email doesn't accept object of type MagicMock
mock_parse_plain_text_email.return_value = parse_plain_text_email(mock_messages_get_response)
# Test happy path
result = await list_emails_by_header(context=mock_context, sender="noreply@github.com", limit=2)
result = await list_emails_by_header(
context=mock_context, sender="noreply@github.com", max_results=2
)
assert isinstance(result, dict)
assert "emails" in result
@ -325,13 +336,15 @@ async def test_search_emails_by_header(mock_parse_email, mock_build, mock_contex
)
with pytest.raises(ToolExecutionError):
await list_emails_by_header(context=mock_context, sender="noreply@github.com", limit=2)
await list_emails_by_header(
context=mock_context, sender="noreply@github.com", max_results=2
)
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.gmail.parse_email")
async def test_get_emails(mock_parse_email, mock_build, mock_context):
@patch("arcade_google.tools.utils.build")
@patch("arcade_google.tools.gmail.parse_plain_text_email")
async def test_get_emails(mock_parse_plain_text_email, mock_build, mock_context):
# Setup test data
mock_messages_list_response = {
"messages": [
@ -379,8 +392,8 @@ async def test_get_emails(mock_parse_email, mock_build, mock_context):
# Mock the Gmail get messages API
mock_service.users().messages().get().execute.return_value = mock_messages_get_response
# Mock the parse_email function since parse_email doesn't accept object of type MagicMock
mock_parse_email.return_value = parse_email(mock_messages_get_response)
# Mock the parse_plain_text_email function since parse_plain_text_email doesn't accept object of type MagicMock
mock_parse_plain_text_email.return_value = parse_plain_text_email(mock_messages_get_response)
# Test happy path
result = await list_emails(context=mock_context, n_emails=1)
@ -404,7 +417,7 @@ async def test_get_emails(mock_parse_email, mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_trash_email(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -430,7 +443,7 @@ async def test_trash_email(mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_search_threads(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -482,7 +495,7 @@ async def test_search_threads(mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_list_threads(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -532,7 +545,7 @@ async def test_list_threads(mock_build, mock_context):
@pytest.mark.asyncio
@patch("arcade_google.tools.gmail.build")
@patch("arcade_google.tools.utils.build")
async def test_get_thread(mock_build, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
@ -579,3 +592,360 @@ async def test_get_thread(mock_build, mock_context):
context=mock_context,
thread_id="invalid_thread",
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"reply_to_whom, expected_to, expected_cc",
[
(
GmailReplyToWhom.EVERY_RECIPIENT,
"sender@example.com, to1@example.com, to2@example.com",
"cc1@example.com, cc2@example.com",
),
(
GmailReplyToWhom.ONLY_THE_SENDER,
"sender@example.com",
"",
),
],
)
@patch("arcade_google.tools.utils.build")
async def test_reply_to_email(mock_build, reply_to_whom, expected_to, expected_cc, mock_context):
mock_service = MagicMock()
mock_build.return_value = mock_service
original_message = {
"id": "id123456",
"threadId": "thread123456",
"payload": {
"headers": [
{"name": "Message-ID", "value": "id123456"},
{"name": "Subject", "value": "test"},
{"name": "From", "value": "sender@example.com"},
{"name": "To", "value": "to1@example.com, to2@example.com, test@example.com"},
{"name": "Cc", "value": "cc1@example.com, cc2@example.com"},
{"name": "References", "value": "thread123456"},
],
},
}
mock_service.users().getProfile().execute.return_value = {"emailAddress": "test@example.com"}
mock_service.users().messages().get().execute.return_value = original_message
result = await reply_to_email(
context=mock_context,
body="test",
reply_to_message_id="id123456",
reply_to_whom=reply_to_whom,
)
assert isinstance(result, dict)
assert "url" in result
replying_to = parse_multipart_email(original_message)
expected_body = build_reply_body("test", replying_to)
expected_message = EmailMessage()
expected_message.set_content(expected_body)
expected_message["To"] = expected_to
expected_message["Subject"] = "Re: test"
if expected_cc:
expected_message["Cc"] = expected_cc
expected_message["In-Reply-To"] = "id123456"
expected_message["References"] = "id123456, thread123456"
mock_service.users().messages().send.assert_called_once_with(
userId="me",
body={
"raw": urlsafe_b64encode(expected_message.as_bytes()).decode(),
"threadId": "thread123456",
},
)
def test_parse_multipart_email_full():
"""
Test parsing a multipart email with both plain text and HTML bodies.
"""
email_data = {
"id": "email123",
"threadId": "thread123",
"labelIds": ["INBOX", "UNREAD"],
"historyId": "history123",
"snippet": "This is a test email.",
"payload": {
"headers": [
{"name": "To", "value": "recipient@example.com"},
{"name": "From", "value": "sender@example.com"},
{"name": "Subject", "value": "Test Email"},
{"name": "Date", "value": "Mon, 1 Jan 2024 10:00:00 -0000"},
],
"body": {"size": 100, "data": "VGhpcyBpcyBhIHRlc3QgZW1haWwu"},
},
}
with (
patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain,
patch("arcade_google.tools.utils._get_email_html_body") as mock_html,
patch("arcade_google.tools.utils._clean_email_body") as mock_clean,
):
# Mock the helper functions
mock_plain.return_value = "This is a test email."
mock_html.return_value = "<p>This is a test email.</p>"
mock_clean.return_value = "This is a test email."
result = parse_multipart_email(email_data)
assert result["id"] == "email123"
assert result["thread_id"] == "thread123"
assert result["label_ids"] == ["INBOX", "UNREAD"]
assert result["snippet"] == "This is a test email."
assert result["to"] == "recipient@example.com"
assert result["from"] == "sender@example.com"
assert result["subject"] == "Test Email"
assert result["date"] == "Mon, 1 Jan 2024 10:00:00 -0000"
assert result["plain_text_body"] == "This is a test email."
assert result["html_body"] == "<p>This is a test email.</p>"
def test_parse_multipart_email_plain_only():
"""
Test parsing an email with only a plain text body.
"""
email_data = {
"id": "email456",
"threadId": "thread456",
"labelIds": ["INBOX"],
"historyId": "history456",
"snippet": "Plain text only email.",
"payload": {
"headers": [
{"name": "To", "value": "recipient2@example.com"},
{"name": "From", "value": "sender2@example.com"},
{"name": "Subject", "value": "Plain Text Email"},
{"name": "Date", "value": "Tue, 2 Feb 2024 11:00:00 -0000"},
],
"body": {"size": 150, "data": "UGxhaW4gdGV4dCBvbmx5IGVtYWlsLg=="},
},
}
with (
patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain,
patch("arcade_google.tools.utils._get_email_html_body") as mock_html,
patch("arcade_google.tools.utils._clean_email_body") as mock_clean,
):
# Mock the helper functions
mock_plain.return_value = "Plain text only email."
mock_html.return_value = None
mock_clean.return_value = "Plain text only email."
result = parse_multipart_email(email_data)
assert result["id"] == "email456"
assert result["thread_id"] == "thread456"
assert result["label_ids"] == ["INBOX"]
assert result["snippet"] == "Plain text only email."
assert result["to"] == "recipient2@example.com"
assert result["from"] == "sender2@example.com"
assert result["subject"] == "Plain Text Email"
assert result["date"] == "Tue, 2 Feb 2024 11:00:00 -0000"
assert result["plain_text_body"] == "Plain text only email."
assert result["html_body"] == ""
def test_parse_multipart_email_html_only():
"""
Test parsing an email with only an HTML body.
"""
email_data = {
"id": "email789",
"threadId": "thread789",
"labelIds": ["SENT"],
"historyId": "history789",
"snippet": "HTML only email.",
"payload": {
"headers": [
{"name": "To", "value": "recipient3@example.com"},
{"name": "From", "value": "sender3@example.com"},
{"name": "Subject", "value": "HTML Email"},
{"name": "Date", "value": "Wed, 3 Mar 2024 12:00:00 -0000"},
],
"body": {"size": 200, "data": "PGh0bWw+VGhpcyBpcyBIVE1MIGVtYWlsLjwvaHRtbD4="},
},
}
with (
patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain,
patch("arcade_google.tools.utils._get_email_html_body") as mock_html,
patch("arcade_google.tools.utils._clean_email_body") as mock_clean,
):
# Mock the helper functions
mock_plain.return_value = None
mock_html.return_value = "<html>This is HTML email.</html>"
mock_clean.return_value = "This is HTML email."
result = parse_multipart_email(email_data)
assert result["id"] == "email789"
assert result["thread_id"] == "thread789"
assert result["label_ids"] == ["SENT"]
assert result["snippet"] == "HTML only email."
assert result["to"] == "recipient3@example.com"
assert result["from"] == "sender3@example.com"
assert result["subject"] == "HTML Email"
assert result["date"] == "Wed, 3 Mar 2024 12:00:00 -0000"
assert result["plain_text_body"] == "This is HTML email."
assert result["html_body"] == "<html>This is HTML email.</html>"
def test_parse_multipart_email_missing_payload():
"""
Test parsing an email with missing payload.
"""
email_data = {
"id": "email000",
"threadId": "thread000",
"labelIds": ["INBOX"],
"historyId": "history000",
"snippet": "Missing payload email.",
# 'payload' key is missing
}
result = parse_multipart_email(email_data)
# Since payload is missing, headers and bodies should be default or empty
assert result["id"] == "email000"
assert result["thread_id"] == "thread000"
assert result["label_ids"] == ["INBOX"]
assert result["snippet"] == "Missing payload email."
assert result["to"] == ""
assert result["from"] == ""
assert result["subject"] == ""
assert result["date"] == ""
assert result["plain_text_body"] == ""
assert result["html_body"] == ""
def test_parse_multipart_email_missing_headers():
"""
Test parsing an email with missing headers in the payload.
"""
email_data = {
"id": "email111",
"threadId": "thread111",
"labelIds": ["INBOX"],
"historyId": "history111",
"snippet": "Missing headers email.",
"payload": {
# 'headers' key is missing
"body": {"size": 100, "data": "VGltZWw="}
},
}
with (
patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain,
patch("arcade_google.tools.utils._get_email_html_body") as mock_html,
patch("arcade_google.tools.utils._clean_email_body") as mock_clean,
):
# Mock the helper functions
mock_plain.return_value = "Timeel"
mock_html.return_value = "<p>Timeel</p>"
mock_clean.return_value = "Timeel"
result = parse_multipart_email(email_data)
assert result["id"] == "email111"
assert result["thread_id"] == "thread111"
assert result["label_ids"] == ["INBOX"]
assert result["snippet"] == "Missing headers email."
assert result["to"] == ""
assert result["from"] == ""
assert result["subject"] == ""
assert result["date"] == ""
assert result["plain_text_body"] == "Timeel"
assert result["html_body"] == "<p>Timeel</p>"
def test_parse_multipart_email_missing_fields():
"""
Test parsing an email with some missing fields in headers.
"""
email_data = {
"id": "email222",
"threadId": "thread222",
"labelIds": ["INBOX"],
"historyId": "history222",
"snippet": "Missing some headers.",
"payload": {
"headers": [
{"name": "From", "value": "sender4@example.com"},
{"name": "Subject", "value": "Partial Headers"},
# 'To' and 'Date' headers are missing
],
"body": {"size": 100, "data": "TWlzc2luZyBzb21lIGhlYWRlcnMu"},
},
}
with (
patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain,
patch("arcade_google.tools.utils._get_email_html_body") as mock_html,
patch("arcade_google.tools.utils._clean_email_body") as mock_clean,
):
# Mock the helper functions
mock_plain.return_value = "Missing some headers."
mock_html.return_value = None
mock_clean.return_value = "Missing some headers."
result = parse_multipart_email(email_data)
assert result["id"] == "email222"
assert result["thread_id"] == "thread222"
assert result["label_ids"] == ["INBOX"]
assert result["snippet"] == "Missing some headers."
assert result["to"] == ""
assert result["from"] == "sender4@example.com"
assert result["subject"] == "Partial Headers"
assert result["date"] == ""
assert result["plain_text_body"] == "Missing some headers."
assert result["html_body"] == ""
def test_parse_multipart_email_empty():
"""
Test parsing an empty email data.
"""
email_data = {}
result = parse_multipart_email(email_data)
assert result["id"] == ""
assert result["thread_id"] == ""
assert result["label_ids"] == []
assert result["snippet"] == ""
assert result["to"] == ""
assert result["from"] == ""
assert result["subject"] == ""
assert result["date"] == ""
assert result["plain_text_body"] == ""
assert result["html_body"] == ""
def test_parse_multipart_email_invalid_payload_structure():
"""
Test parsing an email with an invalid payload structure.
"""
email_data = {
"id": "email333",
"threadId": "thread333",
"labelIds": ["INBOX"],
"historyId": "history333",
"snippet": "Invalid payload structure.",
"payload": {
"headers": "This should be a list, not a string",
"body": {"size": 100, "data": "SW52YWxpZCBwYXlsb2Fk"},
},
}
with pytest.raises(TypeError):
parse_multipart_email(email_data)