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:
parent
f997ca9449
commit
7b1110f2b7
7 changed files with 1401 additions and 248 deletions
18
toolkits/google/arcade_google/tools/constants.py
Normal file
18
toolkits/google/arcade_google/tools/constants.py
Normal 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
|
||||
43
toolkits/google/arcade_google/tools/exceptions.py
Normal file
43
toolkits/google/arcade_google/tools/exceptions.py
Normal 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
|
||||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue