Google Toolkit (Drive, Docs) (#97)

**New Tools Added**
- `docs.py`: Provides tools for Google Docs functionalities, including
creating documents and inserting text.
- `drive.py`: Introduces tools for Google Drive operations, such as
listing documents.

This PR also focuses on simplifying the error handling logic in the Google
toolkit, specifically within the Calendar and Gmail tools. The primary
change involves removing redundant `try-except` blocks that were
catching `HttpError` and general exceptions, and re-raising them as
`ToolExecutionError`. By removing these blocks, we allow exceptions to
propagate naturally, and be handled by the ``ToolExecutor``

---------

Co-authored-by: Eric Gustin <eric@arcade-ai.com>
This commit is contained in:
Sam Partee 2024-10-08 17:04:16 -07:00 committed by GitHub
parent 2c8a172aaa
commit 63cabe8f1f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 1115 additions and 341 deletions

View file

@ -237,8 +237,12 @@ class EvalCase:
evaluation_result.passed = False
evaluation_result.warning = False
expected_count = len(self.expected_tool_calls)
expected_tool_names = ", ".join(
tool_call.name for tool_call in self.expected_tool_calls
)
evaluation_result.failure_reason = (
f"Expected {expected_count} tool call(s), but got {actual_count}"
f"Expected {expected_count} tool call(s), but got {actual_count}. "
+ f"\nExpected tool calls: {expected_tool_names}.\nActual tool calls: {', '.join(actual_tools)}"
)
return evaluation_result

View file

@ -1 +1 @@
__all__ = ["gmail"]
__all__ = ["gmail", "calendar", "drive", "docs"]

View file

@ -119,66 +119,56 @@ async def list_events(
"""
service = build("calendar", "v3", credentials=Credentials(context.authorization.token))
try:
# Get the calendar's time zone
calendar = service.calendars().get(calendarId=calendar_id).execute()
time_zone = calendar["timeZone"]
# Get the calendar's time zone
calendar = service.calendars().get(calendarId=calendar_id).execute()
time_zone = calendar["timeZone"]
# Convert enum values to datetime with timezone offset
start_datetime = datetime.combine(
min_day.to_date(time_zone), min_time_slot.to_time()
).astimezone(ZoneInfo(time_zone))
end_datetime = datetime.combine(
max_day.to_date(time_zone), max_time_slot.to_time()
).astimezone(ZoneInfo(time_zone))
# Convert enum values to datetime with timezone offset
start_datetime = datetime.combine(
min_day.to_date(time_zone), min_time_slot.to_time()
).astimezone(ZoneInfo(time_zone))
end_datetime = datetime.combine(max_day.to_date(time_zone), max_time_slot.to_time()).astimezone(
ZoneInfo(time_zone)
)
if start_datetime > end_datetime:
start_datetime, end_datetime = end_datetime, start_datetime
if start_datetime > end_datetime:
start_datetime, end_datetime = end_datetime, start_datetime
events_result = (
service.events()
.list(
calendarId=calendar_id,
timeMin=start_datetime.isoformat(),
timeMax=end_datetime.isoformat(),
maxResults=max_results,
singleEvents=True,
orderBy="startTime",
)
.execute()
events_result = (
service.events()
.list(
calendarId=calendar_id,
timeMin=start_datetime.isoformat(),
timeMax=end_datetime.isoformat(),
maxResults=max_results,
singleEvents=True,
orderBy="startTime",
)
.execute()
)
items_keys = [
"attachments",
"attendees",
"creator",
"description",
"end",
"eventType",
"htmlLink",
"id",
"location",
"organizer",
"start",
"summary",
"visibility",
]
items_keys = [
"attachments",
"attendees",
"creator",
"description",
"end",
"eventType",
"htmlLink",
"id",
"location",
"organizer",
"start",
"summary",
"visibility",
]
events = [
{key: event[key] for key in items_keys if key in event}
for event in events_result.get("items", [])
]
events = [
{key: event[key] for key in items_keys if key in event}
for event in events_result.get("items", [])
]
return {"events_count": len(events), "events": events}
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{list_events.__name__}' tool.", str(e)
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{list_events.__name__}' tool.",
str(e),
)
return {"events_count": len(events), "events": events}
@tool(
@ -233,76 +223,66 @@ async def update_event(
"""
service = build("calendar", "v3", credentials=Credentials(context.authorization.token))
calendar = service.calendars().get(calendarId="primary").execute()
time_zone = calendar["timeZone"]
try:
calendar = service.calendars().get(calendarId="primary").execute()
time_zone = calendar["timeZone"]
try:
event = service.events().get(calendarId="primary", eventId=event_id).execute()
except HttpError: # TODO: This is a first pass. We should do better.
valid_events_with_id = (
service.events()
.list(
calendarId="primary",
timeMin=(datetime.now() - timedelta(days=2)).isoformat(),
timeMax=(datetime.now() - timedelta(days=2)).isoformat(),
maxResults=50,
singleEvents=True,
orderBy="startTime",
)
.execute()
)
raise RetryableToolError(
f"Event with ID {event_id} not found.",
additional_prompt_content=f"Here is list of valid events. The event_id parameter must match one of these: {valid_events_with_id}",
retry_after_ms=1000,
developer_message=f"Event with ID {event_id} not found. Please try again with a valid event ID.",
)
update_fields = {
"start": _update_datetime(updated_start_day, updated_start_time, time_zone),
"end": _update_datetime(updated_end_day, updated_end_time, time_zone),
"calendarId": updated_calendar_id,
"sendUpdates": send_updates.value if send_updates else None,
"summary": updated_summary,
"description": updated_description,
"location": updated_location,
"visibility": updated_visibility.value if updated_visibility else None,
}
event.update({k: v for k, v in update_fields.items() if v is not None})
if attendee_emails_to_remove:
event["attendees"] = [
attendee
for attendee in event.get("attendees", [])
if attendee.get("email", "") not in attendee_emails_to_remove
]
if attendee_emails_to_add:
event["attendees"] = event.get("attendees", []) + [
{"email": email} for email in attendee_emails_to_add
]
updated_event = (
event = service.events().get(calendarId="primary", eventId=event_id).execute()
except HttpError: # TODO: This is a first pass. We should do better.
valid_events_with_id = (
service.events()
.update(
.list(
calendarId="primary",
eventId=event_id,
sendUpdates=send_updates.value,
body=event,
timeMin=(datetime.now() - timedelta(days=2)).isoformat(),
timeMax=(datetime.now() - timedelta(days=2)).isoformat(),
maxResults=50,
singleEvents=True,
orderBy="startTime",
)
.execute()
)
return f"Event with ID {event_id} successfully updated at {updated_event['updated']}. View updated event at {updated_event['htmlLink']}"
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{update_event.__name__}' tool.", str(e)
raise RetryableToolError(
f"Event with ID {event_id} not found.",
additional_prompt_content=f"Here is list of valid events. The event_id parameter must match one of these: {valid_events_with_id}",
retry_after_ms=1000,
developer_message=f"Event with ID {event_id} not found. Please try again with a valid event ID.",
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{update_event.__name__}' tool.",
str(e),
update_fields = {
"start": _update_datetime(updated_start_day, updated_start_time, time_zone),
"end": _update_datetime(updated_end_day, updated_end_time, time_zone),
"calendarId": updated_calendar_id,
"sendUpdates": send_updates.value if send_updates else None,
"summary": updated_summary,
"description": updated_description,
"location": updated_location,
"visibility": updated_visibility.value if updated_visibility else None,
}
event.update({k: v for k, v in update_fields.items() if v is not None})
if attendee_emails_to_remove:
event["attendees"] = [
attendee
for attendee in event.get("attendees", [])
if attendee.get("email", "") not in attendee_emails_to_remove
]
if attendee_emails_to_add:
event["attendees"] = event.get("attendees", []) + [
{"email": email} for email in attendee_emails_to_add
]
updated_event = (
service.events()
.update(
calendarId="primary",
eventId=event_id,
sendUpdates=send_updates.value,
body=event,
)
.execute()
)
return f"Event with ID {event_id} successfully updated at {updated_event['updated']}. View updated event at {updated_event['htmlLink']}"
@tool(
@ -321,26 +301,16 @@ async def delete_event(
"""Delete an event from Google Calendar."""
service = build("calendar", "v3", credentials=Credentials(context.authorization.token))
try:
service.events().delete(
calendarId=calendar_id, eventId=event_id, sendUpdates=send_updates.value
).execute()
service.events().delete(
calendarId=calendar_id, eventId=event_id, sendUpdates=send_updates.value
).execute()
notification_message = ""
if send_updates == SendUpdatesOptions.ALL:
notification_message = "Notifications were sent to all attendees."
elif send_updates == SendUpdatesOptions.EXTERNAL_ONLY:
notification_message = "Notifications were sent to external attendees only."
elif send_updates == SendUpdatesOptions.NONE:
notification_message = "No notifications were sent to attendees."
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{delete_event.__name__}' tool.", str(e)
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{delete_event.__name__}' tool.",
str(e),
)
else:
return f"Event with ID '{event_id}' successfully deleted from calendar '{calendar_id}'. {notification_message}"
notification_message = ""
if send_updates == SendUpdatesOptions.ALL:
notification_message = "Notifications were sent to all attendees."
elif send_updates == SendUpdatesOptions.EXTERNAL_ONLY:
notification_message = "Notifications were sent to external attendees only."
elif send_updates == SendUpdatesOptions.NONE:
notification_message = "No notifications were sent to attendees."
return f"Event with ID '{event_id}' successfully deleted from calendar '{calendar_id}'. {notification_message}"

View file

@ -0,0 +1,151 @@
from typing import Annotated
from arcade.core.schema import ToolContext
from arcade.sdk import tool
from arcade.sdk.auth import Google
from arcade_google.tools.utils import build_docs_service
# Uses https://developers.google.com/docs/api/reference/rest/v1/documents/get
# Example `arcade chat` query: `get document with ID 1234567890`
# Note: Document IDs are returned in the response of the Google Drive's `list_documents` tool
@tool(
requires_auth=Google(
scopes=[
"https://www.googleapis.com/auth/documents.readonly",
],
)
)
async def get_document_by_id(
context: ToolContext,
document_id: Annotated[str, "The ID of the document to retrieve."],
) -> Annotated[dict, "The document contents as a dictionary"]:
"""
Get the latest version of the specified Google Docs document.
"""
service = build_docs_service(context.authorization.token)
# Execute the documents().get() method. Returns a Document object
# https://developers.google.com/docs/api/reference/rest/v1/documents#Document
request = service.documents().get(documentId=document_id)
response = request.execute()
return response
# Uses https://developers.google.com/docs/api/reference/rest/v1/documents/batchUpdate
# Example `arcade chat` query: `insert "The END" at the end of document with ID 1234567890`
@tool(
requires_auth=Google(
scopes=[
"https://www.googleapis.com/auth/documents",
],
)
)
async def insert_text_at_end_of_document(
context: ToolContext,
document_id: Annotated[str, "The ID of the document to update."],
text_content: Annotated[str, "The text content to insert into the document"],
) -> Annotated[dict, "The response from the batchUpdate API as a dict."]:
"""
Updates an existing Google Docs document using the batchUpdate API endpoint.
"""
document = await get_document_by_id(context, document_id)
end_index = document["body"]["content"][-1]["endIndex"]
service = build_docs_service(context.authorization.token)
requests = [
{
"insertText": {
"location": {
"index": int(end_index) - 1,
},
"text": text_content,
}
}
]
# Execute the documents().batchUpdate() method
response = (
service.documents()
.batchUpdate(documentId=document_id, body={"requests": requests})
.execute()
)
return response
# Uses https://developers.google.com/docs/api/reference/rest/v1/documents/create
# Example `arcade chat` query: `create blank document with title "My New Document"`
@tool(
requires_auth=Google(
scopes=[
"https://www.googleapis.com/auth/documents",
],
)
)
async def create_blank_document(
context: ToolContext, title: Annotated[str, "The title of the blank document to create"]
) -> Annotated[dict, "The created document's title, documentId, and documentUrl in a dictionary"]:
"""
Create a blank Google Docs document with the specified title.
"""
service = build_docs_service(context.authorization.token)
body = {"title": title}
# Execute the documents().create() method. Returns a Document object https://developers.google.com/docs/api/reference/rest/v1/documents#Document
request = service.documents().create(body=body)
response = request.execute()
return {
"title": response["title"],
"documentId": response["documentId"],
"documentUrl": f"https://docs.google.com/document/d/{response['documentId']}/edit",
}
# Uses https://developers.google.com/docs/api/reference/rest/v1/documents/batchUpdate
# Example `arcade chat` query: `create document with title "My New Document" and text content "Hello, World!"`
@tool(
requires_auth=Google(
scopes=[
"https://www.googleapis.com/auth/documents",
],
)
)
async def create_document_from_text(
context: ToolContext,
title: Annotated[str, "The title of the document to create"],
text_content: Annotated[str, "The text content to insert into the document"],
) -> Annotated[dict, "The created document's title, documentId, and documentUrl in a dictionary"]:
"""
Create a Google Docs document with the specified title and text content.
"""
# First, create a blank document
document = await create_blank_document(context, title)
service = build_docs_service(context.authorization.token)
requests = [
{
"insertText": {
"location": {
"index": 1,
},
"text": text_content,
}
}
]
# Execute the batchUpdate method to insert text
service.documents().batchUpdate(
documentId=document["documentId"], body={"requests": requests}
).execute()
return {
"title": document["title"],
"documentId": document["documentId"],
"documentUrl": f"https://docs.google.com/document/d/{document['documentId']}/edit",
}

View file

@ -0,0 +1,80 @@
from typing import Annotated, Optional
from arcade.core.schema import ToolContext
from arcade.sdk import tool
from arcade.sdk.auth import Google
from arcade_google.tools.utils import build_drive_service, remove_none_values
from .models import Corpora, OrderBy
# Implements: https://googleapis.github.io/google-api-python-client/docs/dyn/drive_v3.files.html#list
# Example `arcade chat` query: `list my 5 most recently modified documents`
# TODO: Support query with natural language. Currently, the tool expects a fully formed query string as input with the syntax defined here: https://developers.google.com/drive/api/guides/search-files
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/drive.readonly"],
)
)
async def list_documents(
context: ToolContext,
corpora: Annotated[Optional[Corpora], "The source of files to list"] = Corpora.USER,
title_keywords: Annotated[
Optional[list[str]], "Keywords or phrases that must be in the document title"
] = None,
order_by: Annotated[
Optional[OrderBy],
"Sort order. Defaults to listing the most recently modified documents first",
] = OrderBy.MODIFIED_TIME_DESC,
supports_all_drives: Annotated[
Optional[bool],
"Whether the requesting application supports both My Drives and shared drives",
] = False,
limit: Annotated[Optional[int], "The number of documents to list"] = 50,
) -> Annotated[
dict,
"A dictionary containing 'documents_count' (number of documents returned) and 'documents' (a list of document details including 'kind', 'mimeType', 'id', and 'name' for each document)",
]:
"""
List documents in the user's Google Drive. Excludes documents that are in the trash.
"""
page_size = min(10, limit)
page_token = None # The page token is used for continuing a previous request on the next page
files = []
service = build_drive_service(context.authorization.token)
query = "mimeType = 'application/vnd.google-apps.document' and trashed = false"
if title_keywords:
# Escape single quotes in title_keywords
title_keywords = [keyword.replace("'", "\\'") for keyword in title_keywords]
# Only support logically ANDed keywords in query for now
keyword_queries = [f"name contains '{keyword}'" for keyword in title_keywords]
query += " and " + " and ".join(keyword_queries)
# Prepare the request parameters
params = {
"q": query,
"pageSize": page_size,
"orderBy": order_by.value,
"corpora": corpora.value,
"supportsAllDrives": supports_all_drives,
}
params = remove_none_values(params)
# Paginate through the results until the limit is reached
while len(files) < limit:
if page_token:
params["pageToken"] = page_token
else:
params.pop("pageToken", None)
results = service.files().list(**params).execute()
batch = results.get("files", [])
files.extend(batch[: limit - len(files)])
page_token = results.get("nextPageToken")
if not page_token or len(batch) < page_size:
break
return {"documents_count": len(files), "documents": files}

View file

@ -8,10 +8,10 @@ from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from arcade.core.errors import ToolExecutionError, ToolInputError
from arcade.core.schema import ToolContext
from arcade.sdk import tool
from arcade.sdk.auth import Google
from arcade.sdk.error import RetryableToolError
from arcade_google.tools.utils import (
DateRange,
build_query_string,
@ -42,37 +42,27 @@ async def send_email(
Send an email using the Gmail API.
"""
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
message = EmailMessage()
message.set_content(body)
message["To"] = recipient
message["Subject"] = subject
if cc:
message["Cc"] = ", ".join(cc)
if bcc:
message["Bcc"] = ", ".join(bcc)
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()
# Encode the message in base64
encoded_message = base64.urlsafe_b64encode(message.as_bytes()).decode()
# Create the email
email = {"raw": encoded_message}
# Create the email
email = {"raw": encoded_message}
# Send the email
sent_message = service.users().messages().send(userId="me", body=email).execute()
return f"Email with ID {sent_message['id']} sent: {get_sent_email_url(sent_message['id'])}"
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{send_email.__name__}' tool.", str(e)
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{send_email.__name__}' tool.",
str(e),
)
# Send the email
sent_message = service.users().messages().send(userId="me", body=email).execute()
return f"Email with ID {sent_message['id']} sent: {get_sent_email_url(sent_message['id'])}"
@tool(
@ -87,24 +77,16 @@ async def send_draft_email(
Send a draft email using the Gmail API.
"""
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Send the draft email
sent_message = service.users().drafts().send(userId="me", body={"id": email_id}).execute()
# Send the draft email
sent_message = service.users().drafts().send(userId="me", body={"id": email_id}).execute()
# Construct the URL to the sent email
return f"Draft email with ID {sent_message['id']} sent: {get_sent_email_url(sent_message['id'])}"
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{send_draft_email.__name__}' tool.", str(e)
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{send_draft_email.__name__}' tool.",
str(e),
)
# Construct the URL to the sent email
return (
f"Draft email with ID {sent_message['id']} sent: {get_sent_email_url(sent_message['id'])}"
)
# Draft Management Tools
@ -124,37 +106,27 @@ 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))
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
message = MIMEText(body)
message["to"] = recipient
message["subject"] = subject
if cc:
message["Cc"] = ", ".join(cc)
if bcc:
message["Bcc"] = ", ".join(bcc)
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()
# Encode the message in base64
raw_message = base64.urlsafe_b64encode(message.as_bytes()).decode()
# Create the draft
draft = {"message": {"raw": raw_message}}
# Create the draft
draft = {"message": {"raw": raw_message}}
draft_message = service.users().drafts().create(userId="me", body=draft).execute()
return f"Draft email with ID {draft_message['id']} created: {get_draft_url(draft_message['id'])}"
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{write_draft_email.__name__}' tool.",
str(e),
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{write_draft_email.__name__}' tool.",
str(e),
)
draft_message = service.users().drafts().create(userId="me", body=draft).execute()
return (
f"Draft email with ID {draft_message['id']} created: {get_draft_url(draft_message['id'])}"
)
@tool(
@ -175,38 +147,27 @@ async def update_draft_email(
Update an existing email draft using the Gmail API.
"""
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
message = MIMEText(body)
message["to"] = recipient
message["subject"] = subject
if cc:
message["Cc"] = ", ".join(cc)
if bcc:
message["Bcc"] = ", ".join(bcc)
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()
# Encode the message in base64
raw_message = base64.urlsafe_b64encode(message.as_bytes()).decode()
# Update the draft
draft = {"id": draft_email_id, "message": {"raw": raw_message}}
# Update the draft
draft = {"id": draft_email_id, "message": {"raw": raw_message}}
updated_draft_message = (
service.users().drafts().update(userId="me", id=draft_email_id, body=draft).execute()
)
return f"Draft email with ID {updated_draft_message['id']} updated: {get_draft_url(updated_draft_message['id'])}"
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{update_draft_email.__name__}' tool.",
str(e),
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{update_draft_email.__name__}' tool.",
str(e),
)
updated_draft_message = (
service.users().drafts().update(userId="me", id=draft_email_id, body=draft).execute()
)
return f"Draft email with ID {updated_draft_message['id']} updated: {get_draft_url(updated_draft_message['id'])}"
@tool(
@ -222,24 +183,12 @@ async def delete_draft_email(
Delete a draft email using the Gmail API.
"""
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Delete the draft
service.users().drafts().delete(userId="me", id=draft_email_id).execute()
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{delete_draft_email.__name__}' tool.",
str(e),
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{delete_draft_email.__name__}' tool.",
str(e),
)
else:
return f"Draft email with ID {draft_email_id} deleted successfully."
# Delete the draft
service.users().drafts().delete(userId="me", id=draft_email_id).execute()
return f"Draft email with ID {draft_email_id} deleted successfully."
# Email Management Tools
@ -255,23 +204,13 @@ async def trash_email(
Move an email to the trash folder using the Gmail API.
"""
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Trash the email
service.users().messages().trash(userId="me", id=email_id).execute()
# Trash the email
service.users().messages().trash(userId="me", id=email_id).execute()
return f"Email with ID {email_id} trashed successfully: {get_email_in_trash_url(email_id)}"
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{trash_email.__name__}' tool.", str(e)
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{trash_email.__name__}' tool.",
str(e),
)
return f"Email with ID {email_id} trashed successfully: {get_email_in_trash_url(email_id)}"
# Draft Search Tools
@ -287,38 +226,27 @@ async def list_draft_emails(
"""
Lists draft emails in the user's draft mailbox using the Gmail API.
"""
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
listed_drafts = service.users().drafts().list(userId="me").execute()
listed_drafts = service.users().drafts().list(userId="me").execute()
if not listed_drafts:
return {"emails": []}
if not listed_drafts:
return {"emails": []}
draft_ids = [draft["id"] for draft in listed_drafts.get("drafts", [])][:n_drafts]
draft_ids = [draft["id"] for draft in listed_drafts.get("drafts", [])][:n_drafts]
emails = []
for draft_id in draft_ids:
try:
draft_data = service.users().drafts().get(userId="me", id=draft_id).execute()
draft_details = parse_draft_email(draft_data)
if draft_details:
emails.append(draft_details)
except Exception as e:
print(f"Error reading draft email {draft_id}: {e}")
emails = []
for draft_id in draft_ids:
try:
draft_data = service.users().drafts().get(userId="me", id=draft_id).execute()
draft_details = parse_draft_email(draft_data)
if draft_details:
emails.append(draft_details)
except Exception as e:
print(f"Error reading draft email {draft_id}: {e}")
return json.dumps({"emails": emails})
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{list_draft_emails.__name__}' tool.",
str(e),
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{list_draft_emails.__name__}' tool.",
str(e),
)
return json.dumps({"emails": emails})
# Email Search Tools
@ -343,31 +271,21 @@ async def list_emails_by_header(
At least one of the following parametersMUST be provided: sender, recipient, subject, body.
"""
if not any([sender, recipient, subject, body]):
raise ToolInputError(
"At least one of sender, recipient, subject, or body must be provided."
raise RetryableToolError(
message="At least one of sender, recipient, subject, or body must be provided.",
developer_message="At least one of sender, recipient, subject, or body must be provided.",
)
query = build_query_string(sender, recipient, subject, body, date_range)
try:
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
messages = fetch_messages(service, query, limit)
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
messages = fetch_messages(service, query, limit)
if not messages:
return json.dumps({"emails": []})
if not messages:
return json.dumps({"emails": []})
emails = process_messages(service, messages)
return json.dumps({"emails": emails})
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{list_emails_by_header.__name__}' tool.",
str(e),
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{list_emails_by_header.__name__}' tool.",
str(e),
)
emails = process_messages(service, messages)
return json.dumps({"emails": emails})
def process_messages(service, messages):
@ -394,32 +312,22 @@ async def list_emails(
"""
Read emails from a Gmail account and extract plain text content.
"""
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
messages = service.users().messages().list(userId="me").execute().get("messages", [])
messages = service.users().messages().list(userId="me").execute().get("messages", [])
if not messages:
return {"emails": []}
if not messages:
return {"emails": []}
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)
if email_details:
emails.append(email_details)
except Exception as e:
print(f"Error reading email {msg['id']}: {e}")
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)
if email_details:
emails.append(email_details)
except Exception as e:
print(f"Error reading email {msg['id']}: {e}")
return json.dumps({"emails": emails})
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{list_emails.__name__}' tool.", str(e)
)
except Exception as e:
raise ToolExecutionError(
f"Unexpected Error encountered during execution of '{list_emails.__name__}' tool.",
str(e),
)
return json.dumps({"emails": emails})

View file

@ -230,3 +230,65 @@ class SendUpdatesOptions(Enum):
NONE = "none" # No notifications are sent
ALL = "all" # Notifications are sent to all guests
EXTERNAL_ONLY = "externalOnly" # Notifications are sent to non-Google Calendar guests only.
# Utils for Google Drive tools
class Corpora(str, Enum):
"""
Bodies of items (files/documents) to which the query applies.
Prefer 'user' or 'drive' to 'allDrives' for efficiency.
By default, corpora is set to 'user'.
"""
USER = "user"
DOMAIN = "domain"
DRIVE = "drive"
ALL_DRIVES = "allDrives"
class OrderBy(str, Enum):
"""
Sort keys for ordering files in Google Drive.
Each key has both ascending and descending options.
"""
CREATED_TIME = "createdTime" # When the file was created (ascending)
CREATED_TIME_DESC = "createdTime desc" # When the file was created (descending)
FOLDER = "folder" # The folder ID, sorted using alphabetical ordering (ascending)
FOLDER_DESC = "folder desc" # The folder ID, sorted using alphabetical ordering (descending)
MODIFIED_BY_ME_TIME = (
"modifiedByMeTime" # The last time the file was modified by the user (ascending)
)
MODIFIED_BY_ME_TIME_DESC = (
"modifiedByMeTime desc" # The last time the file was modified by the user (descending)
)
MODIFIED_TIME = "modifiedTime" # The last time the file was modified by anyone (ascending)
MODIFIED_TIME_DESC = (
"modifiedTime desc" # The last time the file was modified by anyone (descending)
)
NAME = "name" # The name of the file, sorted using alphabetical ordering (e.g., 1, 12, 2, 22) (ascending)
NAME_DESC = "name desc" # The name of the file, sorted using alphabetical ordering (e.g., 1, 12, 2, 22) (descending)
NAME_NATURAL = "name_natural" # The name of the file, sorted using natural sort ordering (e.g., 1, 2, 12, 22) (ascending)
NAME_NATURAL_DESC = "name_natural desc" # The name of the file, sorted using natural sort ordering (e.g., 1, 2, 12, 22) (descending)
QUOTA_BYTES_USED = (
"quotaBytesUsed" # The number of storage quota bytes used by the file (ascending)
)
QUOTA_BYTES_USED_DESC = (
"quotaBytesUsed desc" # The number of storage quota bytes used by the file (descending)
)
RECENCY = "recency" # The most recent timestamp from the file's date-time fields (ascending)
RECENCY_DESC = (
"recency desc" # The most recent timestamp from the file's date-time fields (descending)
)
SHARED_WITH_ME_TIME = (
"sharedWithMeTime" # When the file was shared with the user, if applicable (ascending)
)
SHARED_WITH_ME_TIME_DESC = "sharedWithMeTime desc" # When the file was shared with the user, if applicable (descending)
STARRED = "starred" # Whether the user has starred the file (ascending)
STARRED_DESC = "starred desc" # Whether the user has starred the file (descending)
VIEWED_BY_ME_TIME = (
"viewedByMeTime" # The last time the file was viewed by the user (ascending)
)
VIEWED_BY_ME_TIME_DESC = (
"viewedByMeTime desc" # The last time the file was viewed by the user (descending)
)

View file

@ -5,6 +5,8 @@ from enum import Enum
from typing import Any, Optional
from bs4 import BeautifulSoup
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from arcade_google.tools.models import Day, TimeSlot
@ -225,3 +227,28 @@ def fetch_messages(service, query_string, limit):
.execute()
)
return response.get("messages", [])
def remove_none_values(params: dict) -> dict:
"""
Remove None values from a dictionary.
:param params: The dictionary to clean
:return: A new dictionary with None values removed
"""
return {k: v for k, v in params.items() if v is not None}
# Drive utils
def build_drive_service(token: str):
"""
Build a Drive service object.
"""
return build("drive", "v3", credentials=Credentials(token))
# Docs utils
def build_docs_service(token: str):
"""
Build a Drive service object.
"""
return build("docs", "v1", credentials=Credentials(token))

View file

@ -0,0 +1,186 @@
import arcade_google
from arcade_google.tools.docs import (
create_blank_document,
create_document_from_text,
get_document_by_id,
insert_text_at_end_of_document,
)
from arcade.core.catalog import ToolCatalog
from arcade.sdk.eval import (
BinaryCritic,
EvalRubric,
EvalSuite,
SimilarityCritic,
tool_eval,
)
# Evaluation rubric
rubric = EvalRubric(
fail_threshold=0.9,
warn_threshold=0.95,
)
catalog = ToolCatalog()
catalog.add_module(arcade_google)
@tool_eval()
def docs_eval_suite() -> EvalSuite:
"""Create an evaluation suite for Google Docs tools."""
suite = EvalSuite(
name="Google Docs Tools Evaluation",
system_message="You are an AI assistant that can create and manage Google Docs using the provided tools.",
catalog=catalog,
rubric=rubric,
)
# A previous tool call to list_documents
additional_messages = [
{"role": "user", "content": "list my 10 most recently created docs"},
{
"role": "assistant",
"content": "Please go to this URL and authorize the action: [Link](https://accounts.google.com/)",
},
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call_gegK723W2hXsORjBmq1Oexqk",
"type": "function",
"function": {
"name": "Google_ListDocuments",
"arguments": '{"limit":10,"order_by":"createdTime desc"}',
},
}
],
},
{
"role": "tool",
"content": '{"documents":[{"id":"1e0rCoT1Yd14WuuEvd3hSUcN_-VD3df4T3Q08uLm3TWc","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst10"},{"id":"1eTSWd-5zQds8K9OWYygwtCFMUyuuMize3bh3HaRsKts","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst9"},{"id":"19Dqugn0rVi89K0C__lpg1HbhQOTenccyZOhPgivTHMs","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst8"},{"id":"1RCibzx14eqP3vS9yI4nD13OKf8Vee56RiszS53OkR7I","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst7"},{"id":"1imFb04JQuBn8SiSsRFf6fEuYCyXkbII4KX8fsmnT0jo","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst6"},{"id":"1ZC3oypdfLWFgBd-emeSykJf9tZOae6USsFboygRCr-w","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst5"},{"id":"1-gFGNWmwLxEiKa6NNixLNq3X-phXRMORVZfVTfBg8Sc","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst4"},{"id":"1eQ8UBO_PY3Lem4R8OVdIc9ODXt0MrSUAnEu994Qz8P8","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst3"},{"id":"1TOxB0MLry-JzntDWDT1LFywTLdr3XDWPT5L5UsHMs5c","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst2"},{"id":"1a1UQ7C90s8kGfnO8k6wfAZz_Cy5nGN2MkCoRB5y2j3w","kind":"drive#file","mimeType":"application/vnd.google-apps.document","name":"Tst1"}],"documents_count":10}',
"tool_call_id": "call_gegK723W2hXsORjBmq1Oexqk",
"name": "Google_ListDocuments",
},
{
"role": "assistant",
"content": "Here are your 10 most recently created Google Docs:\n\n1. [Tst10](https://docs.google.com/document/d/1e0rCoT1Yd14WuuEvd3hSUcN_-VD3df4T3Q08uLm3TWc)\n2. [Tst9](https://docs.google.com/document/d/1eTSWd-5zQds8K9OWYygwtCFMUyuuMize3bh3HaRsKts)\n3. [Tst8](https://docs.google.com/document/d/19Dqugn0rVi89K0C__lpg1HbhQOTenccyZOhPgivTHMs)\n4. [Tst7](https://docs.google.com/document/d/1RCibzx14eqP3vS9yI4nD13OKf8Vee56RiszS53OkR7I)\n5. [Tst6](https://docs.google.com/document/d/1imFb04JQuBn8SiSsRFf6fEuYCyXkbII4KX8fsmnT0jo)\n6. [Tst5](https://docs.google.com/document/d/1ZC3oypdfLWFgBd-emeSykJf9tZOae6USsFboygRCr-w)\n7. [Tst4](https://docs.google.com/document/d/1-gFGNWmwLxEiKa6NNixLNq3X-phXRMORVZfVTfBg8Sc)\n8. [Tst3](https://docs.google.com/document/d/1eQ8UBO_PY3Lem4R8OVdIc9ODXt0MrSUAnEu994Qz8P8)\n9. [Tst2](https://docs.google.com/document/d/1TOxB0MLry-JzntDWDT1LFywTLdr3XDWPT5L5UsHMs5c)\n10. [Tst1](https://docs.google.com/document/d/1a1UQ7C90s8kGfnO8k6wfAZz_Cy5nGN2MkCoRB5y2j3w)\n\nYou can click the links to open each document.",
},
]
suite.add_case(
name="Get document content",
user_message="Can you read me the contents of Tst9 doc and also Tst10 doc please",
expected_tool_calls=[
(
get_document_by_id,
{
"document_id": "1eTSWd-5zQds8K9OWYygwtCFMUyuuMize3bh3HaRsKts",
},
),
(
get_document_by_id,
{
"document_id": "1e0rCoT1Yd14WuuEvd3hSUcN_-VD3df4T3Q08uLm3TWc",
},
),
],
critics=[
BinaryCritic(critic_field="document_id", weight=0.6),
],
additional_messages=additional_messages,
)
suite.add_case(
name="Insert text at end of document",
user_message="Please add the text 'This is a new paragraph.' to the end of Tst4.",
expected_tool_calls=[
(
insert_text_at_end_of_document,
{
"document_id": "1-gFGNWmwLxEiKa6NNixLNq3X-phXRMORVZfVTfBg8Sc",
"text_content": "This is a new paragraph.",
},
)
],
critics=[
BinaryCritic(critic_field="document_id", weight=0.5),
SimilarityCritic(critic_field="text_content", weight=0.5),
],
additional_messages=additional_messages,
)
suite.add_case(
name="Read the contents of two documents and then insert text at end of a different document.",
user_message="Can you read me the contents of Tst9 doc and also Tst10 doc please. Also, please add the text 'This is a new paragraph.' to the end of Tst4.",
expected_tool_calls=[
(
insert_text_at_end_of_document,
{
"document_id": "1-gFGNWmwLxEiKa6NNixLNq3X-phXRMORVZfVTfBg8Sc",
"text_content": "This is a new paragraph.",
},
),
(
get_document_by_id,
{
"document_id": "1eTSWd-5zQds8K9OWYygwtCFMUyuuMize3bh3HaRsKts",
},
),
(
get_document_by_id,
{
"document_id": "1e0rCoT1Yd14WuuEvd3hSUcN_-VD3df4T3Q08uLm3TWc",
},
),
],
critics=[
BinaryCritic(critic_field="document_id", weight=0.3),
SimilarityCritic(critic_field="text_content", weight=0.3),
BinaryCritic(critic_field="document_id", weight=0.3),
],
additional_messages=additional_messages,
)
suite.add_case(
name="Create blank document",
user_message="Create a new Doc titled 'Meeting Notes'.",
expected_tool_calls=[
(
create_blank_document,
{
"title": "Meeting Notes",
},
)
],
critics=[
SimilarityCritic(critic_field="title", weight=1.0),
],
)
suite.add_case(
name="Create document from text",
user_message="Create a new doc called To-Do List with the content 'Buy groceries, Call mom, Finish report'.",
expected_tool_calls=[
(
create_document_from_text,
{
"title": "To-Do List",
"text_content": "Buy groceries\nCall mom\nFinish report",
},
)
],
critics=[
SimilarityCritic(critic_field="title", weight=0.5),
SimilarityCritic(critic_field="text_content", weight=0.5),
],
)
suite.add_case(
name="No tool call case",
user_message="Create a new microsoft word document titled 'My Resume'.",
expected_tool_calls=[],
critics=[],
)
return suite

View file

@ -0,0 +1,105 @@
import arcade_google
from arcade_google.tools.drive import list_documents
from arcade_google.tools.models import Corpora, OrderBy
from arcade.core.catalog import ToolCatalog
from arcade.sdk.eval import (
BinaryCritic,
EvalRubric,
EvalSuite,
tool_eval,
)
# Evaluation rubric
rubric = EvalRubric(
fail_threshold=0.9,
warn_threshold=0.95,
)
catalog = ToolCatalog()
catalog.add_module(arcade_google)
@tool_eval()
def drive_eval_suite() -> EvalSuite:
"""Create an evaluation suite for Google Drive tools."""
suite = EvalSuite(
name="Google Drive Tools Evaluation",
system_message="You are an AI assistant that can manage Google Drive documents using the provided tools.",
catalog=catalog,
rubric=rubric,
)
suite.add_case(
name="List documents in Google Drive",
user_message="show me the titles of my 39 most recently created documents. Show me the newest ones first and the oldest ones last.",
expected_tool_calls=[
(
list_documents,
{
"corpora": Corpora.USER,
"order_by": OrderBy.CREATED_TIME_DESC,
"supports_all_drives": False,
"limit": 39,
},
)
],
critics=[
BinaryCritic(critic_field="corpora", weight=0.25),
BinaryCritic(critic_field="order_by", weight=0.25),
BinaryCritic(critic_field="supports_all_drives", weight=0.25),
BinaryCritic(critic_field="limit", weight=0.25),
],
)
suite.add_case(
name="List documents in Google Drive based on title keywords",
user_message="list all documents that have title that contains the word'greedy' and also the phrase 'Joe's algo'",
expected_tool_calls=[
(
list_documents,
{
"corpora": Corpora.USER,
"title_keywords": ["greedy", "Joe's algo"],
"order_by": OrderBy.MODIFIED_TIME_DESC,
"supports_all_drives": False,
"limit": 50,
},
)
],
critics=[
BinaryCritic(critic_field="order_by", weight=0.25),
BinaryCritic(critic_field="title_keywords", weight=0.75),
],
)
suite.add_case(
name="List documents in shared drives",
user_message="List the 5 documents from all drives that nobody has touched in forever, including shared ones.",
expected_tool_calls=[
(
list_documents,
{
"corpora": Corpora.ALL_DRIVES,
"order_by": OrderBy.MODIFIED_TIME,
"supports_all_drives": True,
"limit": 5,
},
)
],
critics=[
BinaryCritic(critic_field="corpora", weight=0.25),
BinaryCritic(critic_field="order_by", weight=0.25),
BinaryCritic(critic_field="supports_all_drives", weight=0.25),
BinaryCritic(critic_field="limit", weight=0.25),
],
)
suite.add_case(
name="No tool call case",
user_message="List my 10 most recently modified documents that are stored in my Microsoft OneDrive.",
expected_tool_calls=[],
critics=[],
)
return suite

View file

@ -0,0 +1,164 @@
from unittest.mock import AsyncMock, patch
import pytest
from arcade_google.tools.docs import (
create_blank_document,
create_document_from_text,
get_document_by_id,
insert_text_at_end_of_document,
)
from arcade_google.tools.utils import build_docs_service
from googleapiclient.errors import HttpError
from arcade.core.errors import ToolExecutionError
@pytest.fixture
def mock_context():
context = AsyncMock()
context.authorization.token = "mock_token" # noqa: S105
return context
@pytest.fixture
def mock_service():
with patch("arcade_google.tools.docs." + build_docs_service.__name__) as mock_build_service:
yield mock_build_service.return_value
@pytest.mark.asyncio
async def test_get_document_by_id_success(mock_context, mock_service):
# Mock the service.documents().get().execute() method
mock_service.documents.return_value.get.return_value.execute.return_value = {
"body": {"content": [{"endIndex": 1, "paragraph": {}}]},
"documentId": "test_document_id",
"title": "Test Document",
}
result = await get_document_by_id(mock_context, "test_document_id")
assert result["documentId"] == "test_document_id"
assert result["title"] == "Test Document"
@pytest.mark.asyncio
async def test_get_document_by_id_http_error(mock_context, mock_service):
# Simulate HttpError
mock_service.documents.return_value.get.return_value.execute.side_effect = HttpError(
resp=AsyncMock(status=404), content=b'{"error": {"message": "Not Found"}}'
)
with pytest.raises(ToolExecutionError, match="Error in execution of GetDocumentById"):
await get_document_by_id(mock_context, "invalid_document_id")
@pytest.mark.asyncio
async def test_insert_text_at_end_of_document_success(mock_context, mock_service):
# Mock get_document_by_id to return a document with endIndex
with patch(
"arcade_google.tools.docs.get_document_by_id",
return_value={"body": {"content": [{"endIndex": 1, "paragraph": {}}]}},
):
# Mock the service.documents().batchUpdate().execute() method
mock_service.documents.return_value.batchUpdate.return_value.execute.return_value = {
"documentId": "test_document_id",
"replies": [],
}
result = await insert_text_at_end_of_document(
mock_context, "test_document_id", "Sample text"
)
assert result["documentId"] == "test_document_id"
@pytest.mark.asyncio
@pytest.mark.asyncio
async def test_insert_text_at_end_of_document_http_error(mock_context, mock_service):
with patch(
"arcade_google.tools.docs.get_document_by_id",
return_value={"body": {"content": [{"endIndex": 1, "paragraph": {}}]}},
):
# Simulate HttpError during batchUpdate
mock_service.documents.return_value.batchUpdate.return_value.execute.side_effect = (
HttpError(resp=AsyncMock(status=400), content=b'{"error": {"message": "Bad Request"}}')
)
with pytest.raises(
ToolExecutionError, match="Error in execution of InsertTextAtEndOfDocument"
):
await insert_text_at_end_of_document(mock_context, "test_document_id", "Sample text")
@pytest.mark.asyncio
async def test_create_blank_document_success(mock_context, mock_service):
# Mock the service.documents().create().execute() method
mock_service.documents.return_value.create.return_value.execute.return_value = {
"documentId": "new_document_id",
"title": "New Document",
}
result = await create_blank_document(mock_context, "New Document")
assert result["documentId"] == "new_document_id"
assert result["title"] == "New Document"
assert "documentUrl" in result
@pytest.mark.asyncio
async def test_create_blank_document_http_error(mock_context, mock_service):
# Simulate HttpError during create
mock_service.documents.return_value.create.return_value.execute.side_effect = HttpError(
resp=AsyncMock(status=403), content=b'{"error": {"message": "Forbidden"}}'
)
with pytest.raises(ToolExecutionError, match="Error in execution of CreateBlankDocument"):
await create_blank_document(mock_context, "New Document")
@pytest.mark.asyncio
async def test_create_document_from_text_success(mock_context, mock_service):
with patch(
"arcade_google.tools.docs." + create_blank_document.__name__
) as mock_create_blank_document:
# Mock create_blank_document
mock_create_blank_document.return_value = {
"documentId": "new_document_id",
"title": "New Document",
}
# Mock the service.documents().batchUpdate().execute() method
mock_service.documents.return_value.batchUpdate.return_value.execute.return_value = {
"documentId": "new_document_id",
"replies": [],
}
result = await create_document_from_text(mock_context, "New Document", "Hello, World!")
assert result["documentId"] == "new_document_id"
assert result["title"] == "New Document"
assert "documentUrl" in result
@pytest.mark.asyncio
async def test_create_document_from_text_http_error(mock_context, mock_service):
with patch(
"arcade_google.tools.docs." + create_blank_document.__name__
) as mock_create_blank_document:
# Mock create_blank_document
mock_create_blank_document.return_value = {
"documentId": "new_document_id",
"title": "New Document",
}
# Simulate HttpError during batchUpdate
mock_service.documents.return_value.batchUpdate.return_value.execute.side_effect = (
HttpError(
resp=AsyncMock(status=500), content=b'{"error": {"message": "Internal Error"}}'
)
)
with pytest.raises(
ToolExecutionError, match="Error in execution of CreateDocumentFromText"
):
await create_document_from_text(mock_context, "New Document", "Hello, World!")

View file

@ -0,0 +1,117 @@
from unittest.mock import AsyncMock, patch
import pytest
from arcade_google.tools.drive import list_documents
from arcade_google.tools.models import Corpora, OrderBy
from arcade_google.tools.utils import build_drive_service
from googleapiclient.errors import HttpError
from arcade.core.errors import ToolExecutionError
@pytest.fixture
def mock_context():
context = AsyncMock()
context.authorization.token = "mock_token" # noqa: S105
return context
@pytest.fixture
def mock_service():
with patch("arcade_google.tools.drive." + build_drive_service.__name__) as mock_build_service:
yield mock_build_service.return_value
@pytest.mark.asyncio
async def test_list_documents_success(mock_context, mock_service):
# Mock the service.files().list().execute() method
mock_service.files.return_value.list.return_value.execute.side_effect = [
{
"files": [
{"id": "file1", "name": "Document 1"},
{"id": "file2", "name": "Document 2"},
],
"nextPageToken": None,
}
]
result = await list_documents(mock_context, limit=2)
assert result["documents_count"] == 2
assert len(result["documents"]) == 2
assert result["documents"][0]["id"] == "file1"
assert result["documents"][1]["id"] == "file2"
@pytest.mark.asyncio
async def test_list_documents_pagination(mock_context, mock_service):
# Simulate multiple pages
mock_service.files.return_value.list.return_value.execute.side_effect = [
{
"files": [{"id": f"file{i}", "name": f"Document {i}"} for i in range(1, 11)],
"nextPageToken": "token1",
},
{
"files": [{"id": f"file{i}", "name": f"Document {i}"} for i in range(11, 21)],
"nextPageToken": None,
},
]
result = await list_documents(mock_context, limit=15)
assert result["documents_count"] == 15
assert len(result["documents"]) == 15
assert result["documents"][0]["id"] == "file1"
assert result["documents"][-1]["id"] == "file15"
@pytest.mark.asyncio
async def test_list_documents_http_error(mock_context, mock_service):
# Simulate HttpError
mock_service.files.return_value.list.return_value.execute.side_effect = HttpError(
resp=AsyncMock(status=403), content=b'{"error": {"message": "Forbidden"}}'
)
with pytest.raises(ToolExecutionError, match="Error in execution of ListDocuments"):
await list_documents(mock_context)
@pytest.mark.asyncio
async def test_list_documents_unexpected_error(mock_context, mock_service):
# Simulate unexpected exception
mock_service.files.return_value.list.return_value.execute.side_effect = Exception(
"Unexpected error"
)
with pytest.raises(ToolExecutionError, match="Error in execution of ListDocuments"):
await list_documents(mock_context)
@pytest.mark.asyncio
async def test_list_documents_with_parameters(mock_context, mock_service):
# Mock the service.files().list().execute() method
mock_service.files.return_value.list.return_value.execute.side_effect = [
{
"files": [
{"id": "file1", "name": "Document 1"},
],
"nextPageToken": None,
}
]
result = await list_documents(
mock_context,
corpora=Corpora.USER,
order_by=OrderBy.MODIFIED_TIME_DESC,
supports_all_drives=False,
limit=1,
)
assert result["documents_count"] == 1
mock_service.files.return_value.list.assert_called_with(
q="mimeType = 'application/vnd.google-apps.document' and trashed = false",
pageSize=1,
orderBy="modifiedTime desc",
corpora="user",
supportsAllDrives=False,
)