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