diff --git a/toolkits/google/arcade_google/tools/constants.py b/toolkits/google/arcade_google/tools/constants.py new file mode 100644 index 00000000..2ef97c93 --- /dev/null +++ b/toolkits/google/arcade_google/tools/constants.py @@ -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 diff --git a/toolkits/google/arcade_google/tools/exceptions.py b/toolkits/google/arcade_google/tools/exceptions.py new file mode 100644 index 00000000..88937f1f --- /dev/null +++ b/toolkits/google/arcade_google/tools/exceptions.py @@ -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 diff --git a/toolkits/google/arcade_google/tools/gmail.py b/toolkits/google/arcade_google/tools/gmail.py index ccf846e2..25b7d2b2 100644 --- a/toolkits/google/arcade_google/tools/gmail.py +++ b/toolkits/google/arcade_google/tools/gmail.py @@ -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} diff --git a/toolkits/google/arcade_google/tools/models.py b/toolkits/google/arcade_google/tools/models.py index 11540e2d..41431c21 100644 --- a/toolkits/google/arcade_google/tools/models.py +++ b/toolkits/google/arcade_google/tools/models.py @@ -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" diff --git a/toolkits/google/arcade_google/tools/utils.py b/toolkits/google/arcade_google/tools/utils.py index c9bd42e6..330590ef 100644 --- a/toolkits/google/arcade_google/tools/utils.py +++ b/toolkits/google/arcade_google/tools/utils.py @@ -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. diff --git a/toolkits/google/evals/eval_google_gmail.py b/toolkits/google/evals/eval_google_gmail.py index 93f62c58..a15db333 100644 --- a/toolkits/google/evals/eval_google_gmail.py +++ b/toolkits/google/evals/eval_google_gmail.py @@ -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 ", + "header_message_id": "", + "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 ", + "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 diff --git a/toolkits/google/tests/test_gmail.py b/toolkits/google/tests/test_gmail.py index 85f64a6d..95850705 100644 --- a/toolkits/google/tests/test_gmail.py +++ b/toolkits/google/tests/test_gmail.py @@ -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 = "

This is a test email.

" + 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"] == "

This is a test email.

" + + +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 = "This is HTML email." + 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"] == "This is HTML email." + + +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 = "

Timeel

" + 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"] == "

Timeel

" + + +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)