From 63cabe8f1f1f4a88baf2d4eaf1d5c2f922c23d72 Mon Sep 17 00:00:00 2001 From: Sam Partee Date: Tue, 8 Oct 2024 17:04:16 -0700 Subject: [PATCH] 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 --- arcade/arcade/sdk/eval/eval.py | 6 +- .../google/arcade_google/tools/__init__.py | 2 +- .../google/arcade_google/tools/calendar.py | 244 ++++++-------- toolkits/google/arcade_google/tools/docs.py | 151 +++++++++ toolkits/google/arcade_google/tools/drive.py | 80 +++++ toolkits/google/arcade_google/tools/gmail.py | 312 ++++++------------ toolkits/google/arcade_google/tools/models.py | 62 ++++ toolkits/google/arcade_google/tools/utils.py | 27 ++ toolkits/google/evals/eval_google_docs.py | 186 +++++++++++ toolkits/google/evals/eval_google_drive.py | 105 ++++++ toolkits/google/tests/test_docs.py | 164 +++++++++ toolkits/google/tests/test_drive.py | 117 +++++++ 12 files changed, 1115 insertions(+), 341 deletions(-) create mode 100644 toolkits/google/arcade_google/tools/docs.py create mode 100644 toolkits/google/arcade_google/tools/drive.py create mode 100644 toolkits/google/evals/eval_google_docs.py create mode 100644 toolkits/google/evals/eval_google_drive.py create mode 100644 toolkits/google/tests/test_docs.py create mode 100644 toolkits/google/tests/test_drive.py diff --git a/arcade/arcade/sdk/eval/eval.py b/arcade/arcade/sdk/eval/eval.py index e11039fb..ab6d5c56 100644 --- a/arcade/arcade/sdk/eval/eval.py +++ b/arcade/arcade/sdk/eval/eval.py @@ -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 diff --git a/toolkits/google/arcade_google/tools/__init__.py b/toolkits/google/arcade_google/tools/__init__.py index dc57b3dd..60174475 100644 --- a/toolkits/google/arcade_google/tools/__init__.py +++ b/toolkits/google/arcade_google/tools/__init__.py @@ -1 +1 @@ -__all__ = ["gmail"] +__all__ = ["gmail", "calendar", "drive", "docs"] diff --git a/toolkits/google/arcade_google/tools/calendar.py b/toolkits/google/arcade_google/tools/calendar.py index 19f3ab51..8d708084 100644 --- a/toolkits/google/arcade_google/tools/calendar.py +++ b/toolkits/google/arcade_google/tools/calendar.py @@ -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}" diff --git a/toolkits/google/arcade_google/tools/docs.py b/toolkits/google/arcade_google/tools/docs.py new file mode 100644 index 00000000..f2f16973 --- /dev/null +++ b/toolkits/google/arcade_google/tools/docs.py @@ -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", + } diff --git a/toolkits/google/arcade_google/tools/drive.py b/toolkits/google/arcade_google/tools/drive.py new file mode 100644 index 00000000..3a775686 --- /dev/null +++ b/toolkits/google/arcade_google/tools/drive.py @@ -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} diff --git a/toolkits/google/arcade_google/tools/gmail.py b/toolkits/google/arcade_google/tools/gmail.py index 7cea305d..1df5cff0 100644 --- a/toolkits/google/arcade_google/tools/gmail.py +++ b/toolkits/google/arcade_google/tools/gmail.py @@ -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}) diff --git a/toolkits/google/arcade_google/tools/models.py b/toolkits/google/arcade_google/tools/models.py index f47e3afb..2130e5d4 100644 --- a/toolkits/google/arcade_google/tools/models.py +++ b/toolkits/google/arcade_google/tools/models.py @@ -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) + ) diff --git a/toolkits/google/arcade_google/tools/utils.py b/toolkits/google/arcade_google/tools/utils.py index 33f55b49..44997a7c 100644 --- a/toolkits/google/arcade_google/tools/utils.py +++ b/toolkits/google/arcade_google/tools/utils.py @@ -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)) diff --git a/toolkits/google/evals/eval_google_docs.py b/toolkits/google/evals/eval_google_docs.py new file mode 100644 index 00000000..1a262bf8 --- /dev/null +++ b/toolkits/google/evals/eval_google_docs.py @@ -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 diff --git a/toolkits/google/evals/eval_google_drive.py b/toolkits/google/evals/eval_google_drive.py new file mode 100644 index 00000000..a5d592b5 --- /dev/null +++ b/toolkits/google/evals/eval_google_drive.py @@ -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 diff --git a/toolkits/google/tests/test_docs.py b/toolkits/google/tests/test_docs.py new file mode 100644 index 00000000..4fdfc27a --- /dev/null +++ b/toolkits/google/tests/test_docs.py @@ -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!") diff --git a/toolkits/google/tests/test_drive.py b/toolkits/google/tests/test_drive.py new file mode 100644 index 00000000..c536b7ae --- /dev/null +++ b/toolkits/google/tests/test_drive.py @@ -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, + )