diff --git a/toolkits/google/arcade_google/tools/constants.py b/toolkits/google/arcade_google/constants.py similarity index 93% rename from toolkits/google/arcade_google/tools/constants.py rename to toolkits/google/arcade_google/constants.py index 98390598..46b3ef2d 100644 --- a/toolkits/google/arcade_google/tools/constants.py +++ b/toolkits/google/arcade_google/constants.py @@ -1,6 +1,6 @@ import os -from arcade_google.tools.models import GmailReplyToWhom +from arcade_google.models import GmailReplyToWhom # The default reply in Gmail is to only the sender. Since Gmail also offers the possibility of # changing the default to 'reply to all', we support both options through an env variable. diff --git a/toolkits/google/arcade_google/doc_to_html.py b/toolkits/google/arcade_google/doc_to_html.py new file mode 100644 index 00000000..d54fcef0 --- /dev/null +++ b/toolkits/google/arcade_google/doc_to_html.py @@ -0,0 +1,99 @@ +def convert_document_to_html(document: dict) -> str: + html = ( + "" + f"{document['title']}" + f'' + "" + ) + for element in document["body"]["content"]: + html += convert_structural_element(element) + html += "" + return html + + +def convert_structural_element(element: dict, wrap_paragraphs: bool = True) -> str: + if "sectionBreak" in element or "tableOfContents" in element: + return "" + + elif "paragraph" in element: + paragraph_content = "" + + prepend, append = get_paragraph_style_tags( + style=element["paragraph"]["paragraphStyle"], + wrap_paragraphs=wrap_paragraphs, + ) + + for item in element["paragraph"]["elements"]: + if "textRun" not in item: + continue + paragraph_content += extract_paragraph_content(item["textRun"]) + + if not paragraph_content: + return "" + + return f"{prepend}{paragraph_content.strip()}{append}" + + elif "table" in element: + table = [ + [ + "".join([ + convert_structural_element(element=cell_element, wrap_paragraphs=False) + for cell_element in cell["content"] + ]) + for cell in row["tableCells"] + ] + for row in element["table"]["tableRows"] + ] + return table_list_to_html(table) + + else: + raise ValueError(f"Unknown document body element type: {element}") + + +def extract_paragraph_content(text_run: dict) -> str: + content = text_run["content"] + style = text_run["textStyle"] + return apply_text_style(content, style) + + +def apply_text_style(content: str, style: dict) -> str: + content = content.rstrip("\n") + content = content.replace("\n", "
") + italic = style.get("italic", False) + bold = style.get("bold", False) + if italic: + content = f"{content}" + if bold: + content = f"{content}" + return content + + +def get_paragraph_style_tags(style: dict, wrap_paragraphs: bool = True) -> tuple[str, str]: + named_style = style["namedStyleType"] + if named_style == "NORMAL_TEXT": + return ("

", "

") if wrap_paragraphs else ("", "") + elif named_style == "TITLE": + return "

", "

" + elif named_style == "SUBTITLE": + return "

", "

" + elif named_style.startswith("HEADING_"): + try: + heading_level = int(named_style.split("_")[1]) + except ValueError: + return ("

", "

") if wrap_paragraphs else ("", "") + else: + return f"", f"" + return ("

", "

") if wrap_paragraphs else ("", "") + + +def table_list_to_html(table: list[list[str]]) -> str: + html = "" + for row in table: + html += "" + for cell in row: + if cell.endswith("
"): + cell = cell[:-4] + html += f"" + html += "" + html += "
{cell}
" + return html diff --git a/toolkits/google/arcade_google/doc_to_markdown.py b/toolkits/google/arcade_google/doc_to_markdown.py new file mode 100644 index 00000000..9595de56 --- /dev/null +++ b/toolkits/google/arcade_google/doc_to_markdown.py @@ -0,0 +1,64 @@ +import arcade_google.doc_to_html as doc_to_html + + +def convert_document_to_markdown(document: dict) -> str: + md = f"---\ntitle: {document['title']}\ndocumentId: {document['documentId']}\n---\n" + for element in document["body"]["content"]: + md += convert_structural_element(element) + return md + + +def convert_structural_element(element: dict) -> str: + if "sectionBreak" in element or "tableOfContents" in element: + return "" + + elif "paragraph" in element: + md = "" + prepend = get_paragraph_style_prepend_str(element["paragraph"]["paragraphStyle"]) + for item in element["paragraph"]["elements"]: + if "textRun" not in item: + continue + content = extract_paragraph_content(item["textRun"]) + md += f"{prepend}{content}" + return md + + elif "table" in element: + return doc_to_html.convert_structural_element(element) + + else: + raise ValueError(f"Unknown document body element type: {element}") + + +def extract_paragraph_content(text_run: dict) -> str: + content = text_run["content"] + style = text_run["textStyle"] + return apply_text_style(content, style) + + +def apply_text_style(content: str, style: dict) -> str: + append = "\n" if content.endswith("\n") else "" + content = content.rstrip("\n") + italic = style.get("italic", False) + bold = style.get("bold", False) + if italic: + content = f"_{content}_" + if bold: + content = f"**{content}**" + return f"{content}{append}" + + +def get_paragraph_style_prepend_str(style: dict) -> str: + named_style = style["namedStyleType"] + if named_style == "NORMAL_TEXT": + return "" + elif named_style == "TITLE": + return "# " + elif named_style == "SUBTITLE": + return "## " + elif named_style.startswith("HEADING_"): + try: + heading_level = int(named_style.split("_")[1]) + return f"{'#' * heading_level} " + except ValueError: + return "" + return "" diff --git a/toolkits/google/arcade_google/tools/exceptions.py b/toolkits/google/arcade_google/exceptions.py similarity index 100% rename from toolkits/google/arcade_google/tools/exceptions.py rename to toolkits/google/arcade_google/exceptions.py diff --git a/toolkits/google/arcade_google/tools/models.py b/toolkits/google/arcade_google/models.py similarity index 98% rename from toolkits/google/arcade_google/tools/models.py rename to toolkits/google/arcade_google/models.py index 41431c21..c4e20d9f 100644 --- a/toolkits/google/arcade_google/tools/models.py +++ b/toolkits/google/arcade_google/models.py @@ -344,6 +344,12 @@ class OrderBy(str, Enum): ) +class DocumentFormat(str, Enum): + MARKDOWN = "markdown" + HTML = "html" + GOOGLE_API_JSON = "google_api_json" + + # ---------------------------------------------------------------------------- # # Google Gmail Models and Enums # ---------------------------------------------------------------------------- # diff --git a/toolkits/google/arcade_google/tools/calendar.py b/toolkits/google/arcade_google/tools/calendar.py index 691884c6..1efd5bce 100644 --- a/toolkits/google/arcade_google/tools/calendar.py +++ b/toolkits/google/arcade_google/tools/calendar.py @@ -8,8 +8,8 @@ from google.oauth2.credentials import Credentials from googleapiclient.discovery import build from googleapiclient.errors import HttpError -from arcade_google.tools.models import EventVisibility, SendUpdatesOptions -from arcade_google.tools.utils import parse_datetime +from arcade_google.models import EventVisibility, SendUpdatesOptions +from arcade_google.utils import parse_datetime @tool( diff --git a/toolkits/google/arcade_google/tools/contacts.py b/toolkits/google/arcade_google/tools/contacts.py index ee037124..37cf478c 100644 --- a/toolkits/google/arcade_google/tools/contacts.py +++ b/toolkits/google/arcade_google/tools/contacts.py @@ -6,8 +6,8 @@ from arcade.sdk.auth import Google from google.oauth2.credentials import Credentials from googleapiclient.discovery import build -from arcade_google.tools.constants import DEFAULT_SEARCH_CONTACTS_LIMIT -from arcade_google.tools.utils import build_people_service, search_contacts +from arcade_google.constants import DEFAULT_SEARCH_CONTACTS_LIMIT +from arcade_google.utils import build_people_service, search_contacts async def _warmup_cache(service) -> None: # type: ignore[no-untyped-def] diff --git a/toolkits/google/arcade_google/tools/docs.py b/toolkits/google/arcade_google/tools/docs.py index a99788aa..a1596616 100644 --- a/toolkits/google/arcade_google/tools/docs.py +++ b/toolkits/google/arcade_google/tools/docs.py @@ -3,7 +3,7 @@ from typing import Annotated from arcade.sdk import ToolContext, tool from arcade.sdk.auth import Google -from arcade_google.tools.utils import build_docs_service +from arcade_google.utils import build_docs_service # Uses https://developers.google.com/docs/api/reference/rest/v1/documents/get diff --git a/toolkits/google/arcade_google/tools/drive.py b/toolkits/google/arcade_google/tools/drive.py index dc5aa628..c7e88b7a 100644 --- a/toolkits/google/arcade_google/tools/drive.py +++ b/toolkits/google/arcade_google/tools/drive.py @@ -4,91 +4,17 @@ from arcade.sdk import ToolContext, tool from arcade.sdk.auth import Google from googleapiclient.errors import HttpError -from arcade_google.tools.utils import ( +from arcade_google.doc_to_html import convert_document_to_html +from arcade_google.doc_to_markdown import convert_document_to_markdown +from arcade_google.models import DocumentFormat, OrderBy +from arcade_google.tools.docs import get_document_by_id +from arcade_google.utils import ( build_drive_service, build_file_tree, build_file_tree_request_params, - remove_none_values, + build_files_list_params, ) -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.file"], - ) -) -async def list_documents( - context: ToolContext, - corpora: Annotated[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[ - OrderBy, - "Sort order. Defaults to listing the most recently modified documents first", - ] = OrderBy.MODIFIED_TIME_DESC, - supports_all_drives: Annotated[ - bool, - "Whether the requesting application supports both My Drives and shared drives", - ] = False, - limit: Annotated[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: list[dict[str, Any]] = [] - - service = build_drive_service( - context.authorization.token if context.authorization and context.authorization.token else "" - ) - - 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} - @tool( requires_auth=Google( @@ -181,3 +107,181 @@ async def get_file_tree_structure( drives.append(drive) return {"drives": drives} + + +# 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.file"], + ) +) +async def search_documents( + context: ToolContext, + document_contains: Annotated[ + Optional[list[str]], + "Keywords or phrases that must be in the document title or body. Provide a list of " + "keywords or phrases if needed.", + ] = None, + document_not_contains: Annotated[ + Optional[list[str]], + "Keywords or phrases that must NOT be in the document title or body. Provide a list of " + "keywords or phrases if needed.", + ] = None, + search_only_in_shared_drive_id: Annotated[ + Optional[str], + "The ID of the shared drive to restrict the search to. If provided, the search will only " + "return documents from this drive. Defaults to None, which searches across all drives.", + ] = None, + include_shared_drives: Annotated[ + bool, + "Whether to include documents from shared drives. Defaults to False (searches only in " + "the user's 'My Drive').", + ] = False, + include_organization_domain_documents: Annotated[ + bool, + "Whether to include documents from the organization's domain. This is applicable to admin " + "users who have permissions to view organization-wide documents in a Google Workspace " + "account. Defaults to False.", + ] = False, + order_by: Annotated[ + Optional[list[OrderBy]], + "Sort order. Defaults to listing the most recently modified documents first", + ] = None, + limit: Annotated[int, "The number of documents to list"] = 50, + pagination_token: Annotated[ + Optional[str], "The pagination token to continue a previous request" + ] = None, +) -> 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)", +]: + """ + Searches for documents in the user's Google Drive. Excludes documents that are in the trash. + """ + if order_by is None: + order_by = [OrderBy.MODIFIED_TIME_DESC] + elif isinstance(order_by, OrderBy): + order_by = [order_by] + + page_size = min(10, limit) + files: list[dict[str, Any]] = [] + + service = build_drive_service( + context.authorization.token if context.authorization and context.authorization.token else "" + ) + + params = build_files_list_params( + mime_type="application/vnd.google-apps.document", + document_contains=document_contains, + document_not_contains=document_not_contains, + page_size=page_size, + order_by=order_by, + pagination_token=pagination_token, + include_shared_drives=include_shared_drives, + search_only_in_shared_drive_id=search_only_in_shared_drive_id, + include_organization_domain_documents=include_organization_domain_documents, + ) + + while len(files) < limit: + if pagination_token: + params["pageToken"] = pagination_token + else: + params.pop("pageToken", None) + + results = service.files().list(**params).execute() + batch = results.get("files", []) + files.extend(batch[: limit - len(files)]) + + pagination_token = results.get("nextPageToken") + if not pagination_token or len(batch) < page_size: + break + + return {"documents_count": len(files), "documents": files} + + +@tool( + requires_auth=Google( + scopes=["https://www.googleapis.com/auth/drive.readonly"], + ) +) +async def search_and_retrieve_documents( + context: ToolContext, + return_format: Annotated[ + DocumentFormat, + "The format of the document to return. Defaults to Markdown.", + ] = DocumentFormat.MARKDOWN, + document_contains: Annotated[ + Optional[list[str]], + "Keywords or phrases that must be in the document title or body. Provide a list of " + "keywords or phrases if needed.", + ] = None, + document_not_contains: Annotated[ + Optional[list[str]], + "Keywords or phrases that must NOT be in the document title or body. Provide a list of " + "keywords or phrases if needed.", + ] = None, + search_only_in_shared_drive_id: Annotated[ + Optional[str], + "The ID of the shared drive to restrict the search to. If provided, the search will only " + "return documents from this drive. Defaults to None, which searches across all drives.", + ] = None, + include_shared_drives: Annotated[ + bool, + "Whether to include documents from shared drives. Defaults to False (searches only in " + "the user's 'My Drive').", + ] = False, + include_organization_domain_documents: Annotated[ + bool, + "Whether to include documents from the organization's domain. This is applicable to admin " + "users who have permissions to view organization-wide documents in a Google Workspace " + "account. Defaults to False.", + ] = False, + order_by: Annotated[ + Optional[list[OrderBy]], + "Sort order. Defaults to listing the most recently modified documents first", + ] = None, + limit: Annotated[int, "The number of documents to list"] = 50, + pagination_token: Annotated[ + Optional[str], "The pagination token to continue a previous request" + ] = None, +) -> Annotated[ + dict, + "A dictionary containing 'documents_count' (number of documents returned) and 'documents' " + "(a list of documents with their content).", +]: + """ + Searches for documents in the user's Google Drive and returns a list of documents (with text + content) matching the search criteria. Excludes documents that are in the trash. + + Note: use this tool only when the user prompt requires the documents' content. If the user only + needs a list of documents, use the `search_documents` tool instead. + """ + response = await search_documents( + context=context, + document_contains=document_contains, + document_not_contains=document_not_contains, + search_only_in_shared_drive_id=search_only_in_shared_drive_id, + include_shared_drives=include_shared_drives, + include_organization_domain_documents=include_organization_domain_documents, + order_by=order_by, + limit=limit, + pagination_token=pagination_token, + ) + + documents = [] + + for item in response["documents"]: + document = await get_document_by_id(context, document_id=item["id"]) + + if return_format == DocumentFormat.MARKDOWN: + document = convert_document_to_markdown(document) + elif return_format == DocumentFormat.HTML: + document = convert_document_to_html(document) + + documents.append(document) + + return {"documents_count": len(documents), "documents": documents} diff --git a/toolkits/google/arcade_google/tools/gmail.py b/toolkits/google/arcade_google/tools/gmail.py index 25b7d2b2..87c86df7 100644 --- a/toolkits/google/arcade_google/tools/gmail.py +++ b/toolkits/google/arcade_google/tools/gmail.py @@ -7,10 +7,10 @@ from arcade.sdk.auth import Google from arcade.sdk.errors import RetryableToolError from googleapiclient.errors import HttpError -from arcade_google.tools.constants import GMAIL_DEFAULT_REPLY_TO -from arcade_google.tools.exceptions import GmailToolError -from arcade_google.tools.models import GmailAction, GmailReplyToWhom -from arcade_google.tools.utils import ( +from arcade_google.constants import GMAIL_DEFAULT_REPLY_TO +from arcade_google.exceptions import GmailToolError +from arcade_google.models import GmailAction, GmailReplyToWhom +from arcade_google.utils import ( DateRange, _build_gmail_service, build_email_message, diff --git a/toolkits/google/arcade_google/tools/utils.py b/toolkits/google/arcade_google/utils.py similarity index 89% rename from toolkits/google/arcade_google/tools/utils.py rename to toolkits/google/arcade_google/utils.py index 9a3fc100..f996282c 100644 --- a/toolkits/google/arcade_google/tools/utils.py +++ b/toolkits/google/arcade_google/utils.py @@ -13,16 +13,9 @@ from bs4 import BeautifulSoup from google.oauth2.credentials import Credentials from googleapiclient.discovery import Resource, build -from arcade_google.tools.constants import DEFAULT_SEARCH_CONTACTS_LIMIT -from arcade_google.tools.exceptions import GmailToolError, GoogleServiceError -from arcade_google.tools.models import ( - Corpora, - Day, - GmailAction, - GmailReplyToWhom, - OrderBy, - TimeSlot, -) +from arcade_google.constants import DEFAULT_SEARCH_CONTACTS_LIMIT +from arcade_google.exceptions import GmailToolError, GoogleServiceError +from arcade_google.models import Corpora, Day, GmailAction, GmailReplyToWhom, OrderBy, TimeSlot ## Set up basic configuration for logging to the console with DEBUG level and a specific format. logging.basicConfig( @@ -599,6 +592,85 @@ def build_drive_service(auth_token: Optional[str]) -> Resource: # type: ignore[ return build("drive", "v3", credentials=Credentials(auth_token)) +def build_files_list_query( + mime_type: str, + document_contains: Optional[list[str]] = None, + document_not_contains: Optional[list[str]] = None, +) -> str: + query = [f"(mimeType = '{mime_type}' and trashed = false)"] + + if isinstance(document_contains, str): + document_contains = [document_contains] + + if isinstance(document_not_contains, str): + document_not_contains = [document_not_contains] + + if document_contains: + for keyword in document_contains: + name_contains = keyword.replace("'", "\\'") + full_text_contains = keyword.replace("'", "\\'") + keyword_query = ( + f"(name contains '{name_contains}' or fullText contains '{full_text_contains}')" + ) + query.append(keyword_query) + + if document_not_contains: + for keyword in document_not_contains: + name_not_contains = keyword.replace("'", "\\'") + full_text_not_contains = keyword.replace("'", "\\'") + keyword_query = ( + f"(name not contains '{name_not_contains}' and " + f"fullText not contains '{full_text_not_contains}')" + ) + query.append(keyword_query) + + return " and ".join(query) + + +def build_files_list_params( + mime_type: str, + page_size: int, + order_by: list[OrderBy], + pagination_token: Optional[str], + include_shared_drives: bool, + search_only_in_shared_drive_id: Optional[str], + include_organization_domain_documents: bool, + document_contains: Optional[list[str]] = None, + document_not_contains: Optional[list[str]] = None, +) -> dict[str, Any]: + query = build_files_list_query( + mime_type=mime_type, + document_contains=document_contains, + document_not_contains=document_not_contains, + ) + + params = { + "q": query, + "pageSize": page_size, + "orderBy": ",".join([item.value for item in order_by]), + "pageToken": pagination_token, + } + + if ( + include_shared_drives + or search_only_in_shared_drive_id + or include_organization_domain_documents + ): + params["includeItemsFromAllDrives"] = "true" + params["supportsAllDrives"] = "true" + + if search_only_in_shared_drive_id: + params["driveId"] = search_only_in_shared_drive_id + params["corpora"] = Corpora.DRIVE.value + + if include_organization_domain_documents: + params["corpora"] = Corpora.DOMAIN.value + + params = remove_none_values(params) + + return params + + def build_file_tree_request_params( order_by: Optional[list[OrderBy]], page_token: Optional[str], diff --git a/toolkits/google/conftest.py b/toolkits/google/conftest.py index c579763d..c848402a 100644 --- a/toolkits/google/conftest.py +++ b/toolkits/google/conftest.py @@ -195,3 +195,969 @@ def sample_drive_file_tree_request_responses() -> tuple[dict, list]: ] return files_list, drives_get + + +@pytest.fixture +def sample_document_and_expected_formats(): + document = { + "title": "The Birth of Machine Experience Engineering", + "documentId": "1234567890", + "body": { + "content": [ + { + "endIndex": 1, + "sectionBreak": { + "sectionStyle": { + "columnSeparatorStyle": "NONE", + "contentDirection": "LEFT_TO_RIGHT", + "sectionType": "CONTINUOUS", + } + }, + }, + { + "startIndex": 1, + "endIndex": 45, + "paragraph": { + "elements": [ + { + "endIndex": 45, + "startIndex": 1, + "textRun": { + "content": "The Birth of Machine Experience Engineering\n", + "textStyle": { + "bold": True, + "fontSize": {"magnitude": 23, "unit": "PT"}, + }, + }, + } + ], + "paragraphStyle": { + "direction": "LEFT_TO_RIGHT", + "headingId": "h.wwd7ec37bh6k", + "keepLinesTogether": False, + "keepWithNext": False, + "namedStyleType": "HEADING_1", + "spaceAbove": {"magnitude": 24, "unit": "PT"}, + }, + }, + }, + { + "startIndex": 45, + "endIndex": 46, + "paragraph": { + "elements": [ + { + "startIndex": 304, + "endIndex": 305, + "inlineObjectElement": { + "inlineObjectId": "kix.2s5wy5oiaf79", + "textStyle": {}, + }, + }, + { + "endIndex": 46, + "startIndex": 45, + "textRun": {"content": "\n", "textStyle": {}}, + }, + ], + "paragraphStyle": { + "direction": "LEFT_TO_RIGHT", + "namedStyleType": "NORMAL_TEXT", + "spaceAbove": {"magnitude": 12, "unit": "PT"}, + "spaceBelow": {"magnitude": 12, "unit": "PT"}, + }, + }, + }, + { + "startIndex": 46, + "endIndex": 297, + "paragraph": { + "elements": [ + { + "startIndex": 46, + "endIndex": 146, + "textRun": { + "content": ( + "LLMs acting on behalf of humans and interacting with real-" + "world systems isn't theoretical anymore - " + ), + "textStyle": {}, + }, + }, + { + "startIndex": 146, + "endIndex": 175, + "textRun": { + "content": "Arcade has made it a reality.", + "textStyle": { + "bold": True, + "italic": True, + }, + }, + }, + { + "startIndex": 175, + "endIndex": 248, + "textRun": { + "content": ( + " With this shift, we're seeing the emergence of a new " + "software practice: " + ), + "textStyle": {}, + }, + }, + { + "startIndex": 248, + "endIndex": 295, + "textRun": { + "content": "Machine Experience Engineering (MX Engineering)", + "textStyle": { + "italic": True, + }, + }, + }, + { + "startIndex": 295, + "endIndex": 297, + "textRun": { + "content": ".\n", + "textStyle": {}, + }, + }, + ], + "paragraphStyle": { + "direction": "LEFT_TO_RIGHT", + "namedStyleType": "NORMAL_TEXT", + "spaceAbove": {"magnitude": 12, "unit": "PT"}, + "spaceBelow": {"magnitude": 12, "unit": "PT"}, + }, + }, + }, + { + "endIndex": 407, + "startIndex": 297, + "table": { + "columns": 3, + "rows": 3, + "tableRows": [ + { + "endIndex": 338, + "startIndex": 297, + "tableCells": [ + { + "content": [ + { + "endIndex": 318, + "paragraph": { + "elements": [ + { + "endIndex": 318, + "startIndex": 309, + "textRun": { + "content": "Column 1\n", + "textStyle": {"bold": True}, + }, + } + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 309, + } + ], + "endIndex": 318, + "startIndex": 308, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + { + "content": [ + { + "endIndex": 334, + "paragraph": { + "elements": [ + { + "endIndex": 326, + "startIndex": 319, + "textRun": { + "content": "Another", + "textStyle": {"italic": True}, + }, + }, + { + "endIndex": 334, + "startIndex": 326, + "textRun": { + "content": " column\n", + "textStyle": {}, + }, + }, + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 319, + } + ], + "endIndex": 334, + "startIndex": 318, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + { + "content": [ + { + "endIndex": 348, + "paragraph": { + "elements": [ + { + "endIndex": 348, + "startIndex": 335, + "textRun": { + "content": "Third column\n", + "textStyle": {}, + }, + } + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 335, + } + ], + "endIndex": 348, + "startIndex": 334, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + ], + "tableRowStyle": {"minRowHeight": {"unit": "PT"}}, + }, + { + "endIndex": 366, + "startIndex": 348, + "tableCells": [ + { + "content": [ + { + "endIndex": 356, + "paragraph": { + "elements": [ + { + "endIndex": 356, + "startIndex": 350, + "textRun": { + "content": "Hello\n", + "textStyle": {}, + }, + } + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 350, + } + ], + "endIndex": 356, + "startIndex": 349, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + { + "content": [ + { + "endIndex": 364, + "paragraph": { + "elements": [ + { + "endIndex": 364, + "startIndex": 357, + "textRun": { + "content": "world!\n", + "textStyle": {}, + }, + } + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 357, + } + ], + "endIndex": 364, + "startIndex": 356, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + { + "content": [ + { + "endIndex": 366, + "paragraph": { + "elements": [ + { + "endIndex": 366, + "startIndex": 365, + "textRun": { + "content": "\n", + "textStyle": {}, + }, + } + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 365, + } + ], + "endIndex": 366, + "startIndex": 364, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + ], + "tableRowStyle": {"minRowHeight": {"unit": "PT"}}, + }, + { + "endIndex": 415, + "startIndex": 366, + "tableCells": [ + { + "content": [ + { + "endIndex": 388, + "paragraph": { + "elements": [ + { + "endIndex": 388, + "startIndex": 368, + "textRun": { + "content": "The quick brown fox\n", + "textStyle": {}, + }, + } + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 368, + } + ], + "endIndex": 388, + "startIndex": 367, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + { + "content": [ + { + "endIndex": 401, + "paragraph": { + "elements": [ + { + "endIndex": 395, + "startIndex": 389, + "textRun": { + "content": "jumped", + "textStyle": {"italic": True}, + }, + }, + { + "endIndex": 401, + "startIndex": 395, + "textRun": { + "content": " over\n", + "textStyle": {}, + }, + }, + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 389, + } + ], + "endIndex": 401, + "startIndex": 388, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + { + "content": [ + { + "endIndex": 415, + "paragraph": { + "elements": [ + { + "endIndex": 415, + "startIndex": 402, + "textRun": { + "content": "the lazy dog\n", + "textStyle": {}, + }, + } + ], + "paragraphStyle": { + "alignment": "START", + "avoidWidowAndOrphan": False, + "borderBetween": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderBottom": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderLeft": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderRight": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "borderTop": { + "color": {}, + "dashStyle": "SOLID", + "padding": {"unit": "PT"}, + "width": {"unit": "PT"}, + }, + "direction": "LEFT_TO_RIGHT", + "indentEnd": {"unit": "PT"}, + "indentFirstLine": {"unit": "PT"}, + "indentStart": {"unit": "PT"}, + "keepLinesTogether": False, + "keepWithNext": False, + "lineSpacing": 100, + "namedStyleType": "NORMAL_TEXT", + "pageBreakBefore": False, + "shading": {"backgroundColor": {}}, + "spaceAbove": {"unit": "PT"}, + "spaceBelow": {"unit": "PT"}, + "spacingMode": "COLLAPSE_LISTS", + }, + }, + "startIndex": 402, + } + ], + "endIndex": 415, + "startIndex": 401, + "tableCellStyle": { + "backgroundColor": {}, + "columnSpan": 1, + "contentAlignment": "TOP", + "paddingBottom": {"magnitude": 5, "unit": "PT"}, + "paddingLeft": {"magnitude": 5, "unit": "PT"}, + "paddingRight": {"magnitude": 5, "unit": "PT"}, + "paddingTop": {"magnitude": 5, "unit": "PT"}, + "rowSpan": 1, + }, + }, + ], + "tableRowStyle": {"minRowHeight": {"unit": "PT"}}, + }, + ], + "tableStyle": { + "tableColumnProperties": [ + {"widthType": "EVENLY_DISTRIBUTED"}, + {"widthType": "EVENLY_DISTRIBUTED"}, + {"widthType": "EVENLY_DISTRIBUTED"}, + ] + }, + }, + }, + ] + }, + } + + expected_markdown = ( + "---\ntitle: The Birth of Machine Experience Engineering\ndocumentId: 1234567890\n---\n" + "# **The Birth of Machine Experience Engineering**\n" + "\n" + "LLMs acting on behalf of humans and interacting with real-world systems isn't theoretical " + "anymore - " + "**_Arcade has made it a reality._** With this shift, we're seeing the emergence of a new " + "software practice: " + "_Machine Experience Engineering (MX Engineering)_.\n" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "
Column 1Another columnThird column
Helloworld!
The quick brown foxjumped overthe lazy dog
" + ) + + expected_html = ( + "" + "The Birth of Machine Experience Engineering" + '' + "" + "

The Birth of Machine Experience Engineering

" + "

LLMs acting on behalf of humans and interacting with real-world systems isn't " + "theoretical anymore - " + "Arcade has made it a reality. With this shift, we're seeing the emergence " + "of a new software practice: Machine Experience Engineering (MX Engineering).

" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "" + "
Column 1Another columnThird column
Helloworld!
The quick brown foxjumped overthe lazy dog
" + "" + ) + + return document, expected_markdown, expected_html diff --git a/toolkits/google/evals/eval_google_drive.py b/toolkits/google/evals/eval_google_drive.py index c3f92e73..742f0fac 100644 --- a/toolkits/google/evals/eval_google_drive.py +++ b/toolkits/google/evals/eval_google_drive.py @@ -8,8 +8,12 @@ from arcade.sdk.eval import ( ) import arcade_google -from arcade_google.tools.drive import get_file_tree_structure, list_documents -from arcade_google.tools.models import Corpora, OrderBy +from arcade_google.models import DocumentFormat, OrderBy +from arcade_google.tools.drive import ( + get_file_tree_structure, + search_and_retrieve_documents, + search_documents, +) # Evaluation rubric rubric = EvalRubric( @@ -21,91 +25,6 @@ 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=[ - ExpectedToolCall( - func=list_documents, - args={ - "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=[ - ExpectedToolCall( - func=list_documents, - args={ - "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=[ - ExpectedToolCall( - func=list_documents, - args={ - "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 - - @tool_eval() def get_file_tree_structure_eval_suite() -> EvalSuite: """Create an evaluation suite for Google Drive tools.""" @@ -213,3 +132,198 @@ def get_file_tree_structure_eval_suite() -> EvalSuite: ) return suite + + +@tool_eval() +def search_documents_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="Search documents in Google Drive", + user_message="get my 49 most recently created documents, list the ones created most recently first.", + expected_tool_calls=[ + ExpectedToolCall( + func=search_documents, + args={ + "order_by": [OrderBy.CREATED_TIME_DESC.value], + "limit": 49, + }, + ) + ], + critics=[ + BinaryCritic(critic_field="order_by", weight=0.5), + BinaryCritic(critic_field="limit", weight=0.5), + ], + ) + + suite.add_case( + name="Search documents in Google Drive based on document keywords", + user_message="Search the documents that contain the word 'greedy' and the phrase 'hello, world'", + expected_tool_calls=[ + ExpectedToolCall( + func=search_documents, + args={ + "document_contains": ["greedy", "hello, world"], + }, + ) + ], + critics=[ + BinaryCritic(critic_field="document_contains", weight=1.0), + ], + ) + + suite.add_case( + name="Search documents in a specific Google Drive based on document keywords", + user_message="Search the documents that contain the word 'greedy' and the phrase 'hello, world' in the drive with id 'abc123'", + expected_tool_calls=[ + ExpectedToolCall( + func=search_documents, + args={ + "document_contains": ["greedy", "hello, world"], + "search_only_in_shared_drive_id": "abc123", + }, + ) + ], + critics=[ + BinaryCritic(critic_field="search_only_in_shared_drive_id", weight=0.5), + BinaryCritic(critic_field="document_contains", weight=0.5), + ], + ) + + suite.add_case( + name="Search documents in a Google Drive Workspace organization domain based on document keywords", + user_message="Search the documents that contain the phrase 'hello, world' in the organization domain", + expected_tool_calls=[ + ExpectedToolCall( + func=search_documents, + args={ + "document_contains": ["hello, world"], + "include_organization_domain_documents": True, + }, + ) + ], + critics=[ + BinaryCritic(critic_field="include_organization_domain_documents", weight=0.5), + BinaryCritic(critic_field="document_contains", weight=0.5), + ], + ) + + suite.add_case( + name="Search documents in shared drives", + user_message="Search the 5 documents from all drives corpora that nobody has touched in forever, excluding shared drives.", + expected_tool_calls=[ + ExpectedToolCall( + func=search_documents, + args={ + "limit": 5, + "include_shared_drives": False, + }, + ) + ], + critics=[ + BinaryCritic(critic_field="include_shared_drives", weight=0.5), + BinaryCritic(critic_field="limit", weight=0.5), + ], + ) + + 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 + + +@tool_eval() +def search_and_retrieve_documents_eval_suite() -> EvalSuite: + """Create an evaluation suite for Google Drive search and retrieve 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="Search and retrieve (write summary)", + user_message="Write a summary of the documents in my Google Drive about 'MX Engineering'", + expected_tool_calls=[ + ExpectedToolCall( + func=search_and_retrieve_documents, + args={ + "document_contains": ["MX Engineering"], + "return_format": DocumentFormat.MARKDOWN, + }, + ) + ], + critics=[ + BinaryCritic(critic_field="document_contains", weight=0.5), + BinaryCritic(critic_field="return_format", weight=0.5), + ], + ) + + suite.add_case( + name="Search and retrieve (project proposal)", + user_message="Display the document contents in HTML format from my Google Drive that contain the phrase 'project proposal'.", + expected_tool_calls=[ + ExpectedToolCall( + func=search_and_retrieve_documents, + args={ + "document_contains": ["project proposal"], + "return_format": DocumentFormat.HTML, + }, + ) + ], + critics=[ + BinaryCritic(critic_field="document_contains", weight=0.5), + BinaryCritic(critic_field="return_format", weight=0.5), + ], + ) + + suite.add_case( + name="Search and retrieve (meeting notes)", + user_message="Retrieve documents that contain both 'meeting notes' and 'budget' in JSON format.", + expected_tool_calls=[ + ExpectedToolCall( + func=search_and_retrieve_documents, + args={ + "document_contains": ["meeting notes", "budget"], + "return_format": DocumentFormat.GOOGLE_API_JSON, + }, + ) + ], + critics=[ + BinaryCritic(critic_field="document_contains", weight=0.5), + BinaryCritic(critic_field="return_format", weight=0.5), + ], + ) + + suite.add_case( + name="Search and retrieve (Q1 report)", + user_message="Show me the content of the documents that mention 'Q1 report' but do not include the expression 'Project XYZ'.", + expected_tool_calls=[ + ExpectedToolCall( + func=search_and_retrieve_documents, + args={ + "document_contains": ["Q1 report"], + "document_not_contains": ["Project XYZ"], + "return_format": DocumentFormat.MARKDOWN, + }, + ) + ], + critics=[ + BinaryCritic(critic_field="document_contains", weight=1 / 3), + BinaryCritic(critic_field="document_not_contains", weight=1 / 3), + BinaryCritic(critic_field="return_format", weight=1 / 3), + ], + ) + + return suite diff --git a/toolkits/google/evals/eval_google_gmail.py b/toolkits/google/evals/eval_google_gmail.py index a15db333..a142e08d 100644 --- a/toolkits/google/evals/eval_google_gmail.py +++ b/toolkits/google/evals/eval_google_gmail.py @@ -11,6 +11,7 @@ from arcade.sdk.eval import ( ) import arcade_google +from arcade_google.models import GmailReplyToWhom from arcade_google.tools.gmail import ( get_thread, list_emails_by_header, @@ -20,8 +21,7 @@ from arcade_google.tools.gmail import ( send_email, write_draft_reply_email, ) -from arcade_google.tools.models import GmailReplyToWhom -from arcade_google.tools.utils import DateRange +from arcade_google.utils import DateRange # Evaluation rubric rubric = EvalRubric( diff --git a/toolkits/google/tests/test_calendar.py b/toolkits/google/tests/test_calendar.py index 19377b26..d2ab62e4 100644 --- a/toolkits/google/tests/test_calendar.py +++ b/toolkits/google/tests/test_calendar.py @@ -5,8 +5,8 @@ from arcade.sdk import ToolAuthorizationContext, ToolContext from arcade.sdk.errors import ToolExecutionError from googleapiclient.errors import HttpError +from arcade_google.models import EventVisibility, SendUpdatesOptions from arcade_google.tools.calendar import create_event, delete_event, list_events, update_event -from arcade_google.tools.models import EventVisibility, SendUpdatesOptions @pytest.fixture diff --git a/toolkits/google/tests/test_doc_to_markdown.py b/toolkits/google/tests/test_doc_to_markdown.py new file mode 100644 index 00000000..d68a74ab --- /dev/null +++ b/toolkits/google/tests/test_doc_to_markdown.py @@ -0,0 +1,10 @@ +import pytest + +from arcade_google.doc_to_markdown import convert_document_to_markdown + + +@pytest.mark.asyncio +async def test_convert_document_to_markdown(sample_document_and_expected_formats): + (sample_document, expected_markdown, _) = sample_document_and_expected_formats + markdown = convert_document_to_markdown(sample_document) + assert markdown == expected_markdown diff --git a/toolkits/google/tests/test_docs.py b/toolkits/google/tests/test_docs.py index 6bc48af1..eec95994 100644 --- a/toolkits/google/tests/test_docs.py +++ b/toolkits/google/tests/test_docs.py @@ -10,7 +10,7 @@ from arcade_google.tools.docs import ( get_document_by_id, insert_text_at_end_of_document, ) -from arcade_google.tools.utils import build_docs_service +from arcade_google.utils import build_docs_service @pytest.fixture diff --git a/toolkits/google/tests/test_drive.py b/toolkits/google/tests/test_drive.py index ac6ce454..fce06062 100644 --- a/toolkits/google/tests/test_drive.py +++ b/toolkits/google/tests/test_drive.py @@ -4,9 +4,13 @@ import pytest from arcade.sdk.errors import ToolExecutionError from googleapiclient.errors import HttpError -from arcade_google.tools.drive import get_file_tree_structure, list_documents -from arcade_google.tools.models import Corpora, OrderBy -from arcade_google.tools.utils import build_drive_service +from arcade_google.models import Corpora, DocumentFormat, OrderBy +from arcade_google.tools.drive import ( + get_file_tree_structure, + search_and_retrieve_documents, + search_documents, +) +from arcade_google.utils import build_drive_service @pytest.fixture @@ -23,7 +27,7 @@ def mock_service(): @pytest.mark.asyncio -async def test_list_documents_success(mock_context, mock_service): +async def test_search_documents_success(mock_context, mock_service): # Mock the service.files().list().execute() method mock_service.files.return_value.list.return_value.execute.side_effect = [ { @@ -35,7 +39,7 @@ async def test_list_documents_success(mock_context, mock_service): } ] - result = await list_documents(mock_context, limit=2) + result = await search_documents(mock_context, limit=2) assert result["documents_count"] == 2 assert len(result["documents"]) == 2 @@ -44,7 +48,7 @@ async def test_list_documents_success(mock_context, mock_service): @pytest.mark.asyncio -async def test_list_documents_pagination(mock_context, mock_service): +async def test_search_documents_pagination(mock_context, mock_service): # Simulate multiple pages mock_service.files.return_value.list.return_value.execute.side_effect = [ { @@ -57,7 +61,7 @@ async def test_list_documents_pagination(mock_context, mock_service): }, ] - result = await list_documents(mock_context, limit=15) + result = await search_documents(mock_context, limit=15) assert result["documents_count"] == 15 assert len(result["documents"]) == 15 @@ -66,29 +70,33 @@ async def test_list_documents_pagination(mock_context, mock_service): @pytest.mark.asyncio -async def test_list_documents_http_error(mock_context, mock_service): +async def test_search_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) + with pytest.raises( + ToolExecutionError, match=f"Error in execution of {search_documents.__tool_name__}" + ): + await search_documents(mock_context) @pytest.mark.asyncio -async def test_list_documents_unexpected_error(mock_context, mock_service): +async def test_search_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) + with pytest.raises( + ToolExecutionError, match=f"Error in execution of {search_documents.__tool_name__}" + ): + await search_documents(mock_context) @pytest.mark.asyncio -async def test_list_documents_with_parameters(mock_context, mock_service): +async def test_search_documents_in_organization_domains(mock_context, mock_service): # Mock the service.files().list().execute() method mock_service.files.return_value.list.return_value.execute.side_effect = [ { @@ -99,24 +107,97 @@ async def test_list_documents_with_parameters(mock_context, mock_service): } ] - result = await list_documents( + result = await search_documents( mock_context, - corpora=Corpora.USER, order_by=OrderBy.MODIFIED_TIME_DESC, - supports_all_drives=False, + include_shared_drives=False, + include_organization_domain_documents=True, 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", + q="(mimeType = 'application/vnd.google-apps.document' and trashed = false)", + corpora=Corpora.DOMAIN.value, pageSize=1, - orderBy="modifiedTime desc", - corpora="user", - supportsAllDrives=False, + orderBy=OrderBy.MODIFIED_TIME_DESC.value, + includeItemsFromAllDrives="true", + supportsAllDrives="true", ) +@pytest.mark.asyncio +@patch("arcade_google.tools.drive.search_documents") +@patch("arcade_google.tools.drive.get_document_by_id") +async def test_search_and_retrieve_documents_in_markdown_format( + mock_get_document_by_id, + mock_search_documents, + mock_context, + sample_document_and_expected_formats, +): + (sample_document, expected_markdown, _) = sample_document_and_expected_formats + mock_search_documents.return_value = { + "documents_count": 1, + "documents": [{"id": sample_document["documentId"], "title": sample_document["title"]}], + } + mock_get_document_by_id.return_value = sample_document + result = await search_and_retrieve_documents( + mock_context, + document_contains=[sample_document["title"]], + return_format=DocumentFormat.MARKDOWN, + ) + assert result["documents_count"] == 1 + assert result["documents"][0] == expected_markdown + + +@pytest.mark.asyncio +@patch("arcade_google.tools.drive.search_documents") +@patch("arcade_google.tools.drive.get_document_by_id") +async def test_search_and_retrieve_documents_in_html_format( + mock_get_document_by_id, + mock_search_documents, + mock_context, + sample_document_and_expected_formats, +): + (sample_document, _, expected_html) = sample_document_and_expected_formats + mock_search_documents.return_value = { + "documents_count": 1, + "documents": [{"id": sample_document["documentId"], "title": sample_document["title"]}], + } + mock_get_document_by_id.return_value = sample_document + result = await search_and_retrieve_documents( + mock_context, + document_contains=[sample_document["title"]], + return_format=DocumentFormat.HTML, + ) + assert result["documents_count"] == 1 + assert result["documents"][0] == expected_html + + +@pytest.mark.asyncio +@patch("arcade_google.tools.drive.search_documents") +@patch("arcade_google.tools.drive.get_document_by_id") +async def test_search_and_retrieve_documents_in_google_json_format( + mock_get_document_by_id, + mock_search_documents, + mock_context, + sample_document_and_expected_formats, +): + (sample_document, _, _) = sample_document_and_expected_formats + mock_search_documents.return_value = { + "documents_count": 1, + "documents": [{"id": sample_document["documentId"], "title": sample_document["title"]}], + } + mock_get_document_by_id.return_value = sample_document + result = await search_and_retrieve_documents( + mock_context, + document_contains=[sample_document["title"]], + return_format=DocumentFormat.GOOGLE_API_JSON, + ) + assert result["documents_count"] == 1 + assert result["documents"][0] == sample_document + + @pytest.mark.asyncio async def test_get_file_tree_structure( mock_context, mock_service, sample_drive_file_tree_request_responses diff --git a/toolkits/google/tests/test_gmail.py b/toolkits/google/tests/test_gmail.py index 95850705..31627fe7 100644 --- a/toolkits/google/tests/test_gmail.py +++ b/toolkits/google/tests/test_gmail.py @@ -7,6 +7,7 @@ from arcade.sdk import ToolAuthorizationContext, ToolContext from arcade.sdk.errors import ToolExecutionError from googleapiclient.errors import HttpError +from arcade_google.models import GmailReplyToWhom from arcade_google.tools.gmail import ( delete_draft_email, get_thread, @@ -22,8 +23,7 @@ from arcade_google.tools.gmail import ( update_draft_email, write_draft_email, ) -from arcade_google.tools.models import GmailReplyToWhom -from arcade_google.tools.utils import ( +from arcade_google.utils import ( build_reply_body, parse_draft_email, parse_multipart_email, @@ -38,7 +38,7 @@ def mock_context(): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_send_email(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -73,7 +73,7 @@ async def test_send_email(mock_build, mock_context): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_write_draft_email(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -108,7 +108,7 @@ async def test_write_draft_email(mock_build, mock_context): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_update_draft_email(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -145,7 +145,7 @@ async def test_update_draft_email(mock_build, mock_context): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_send_draft_email(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -170,7 +170,7 @@ async def test_send_draft_email(mock_build, mock_context): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_delete_draft_email(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -192,7 +192,7 @@ async def test_delete_draft_email(mock_build, mock_context): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") @patch("arcade_google.tools.gmail.parse_draft_email") async def test_get_draft_emails(mock_parse_draft_email, mock_build, mock_context): # Setup test data @@ -265,7 +265,7 @@ async def test_get_draft_emails(mock_parse_draft_email, mock_build, mock_context @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") @patch("arcade_google.tools.gmail.parse_plain_text_email") async def test_search_emails_by_header(mock_parse_plain_text_email, mock_build, mock_context): # Setup test data @@ -342,7 +342,7 @@ async def test_search_emails_by_header(mock_parse_plain_text_email, mock_build, @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") @patch("arcade_google.tools.gmail.parse_plain_text_email") async def test_get_emails(mock_parse_plain_text_email, mock_build, mock_context): # Setup test data @@ -417,7 +417,7 @@ async def test_get_emails(mock_parse_plain_text_email, mock_build, mock_context) @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_trash_email(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -443,7 +443,7 @@ async def test_trash_email(mock_build, mock_context): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_search_threads(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -495,7 +495,7 @@ async def test_search_threads(mock_build, mock_context): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_list_threads(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -545,7 +545,7 @@ async def test_list_threads(mock_build, mock_context): @pytest.mark.asyncio -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_get_thread(mock_build, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -610,7 +610,7 @@ async def test_get_thread(mock_build, mock_context): ), ], ) -@patch("arcade_google.tools.utils.build") +@patch("arcade_google.tools.gmail._build_gmail_service") async def test_reply_to_email(mock_build, reply_to_whom, expected_to, expected_cc, mock_context): mock_service = MagicMock() mock_build.return_value = mock_service @@ -686,9 +686,9 @@ def test_parse_multipart_email_full(): } with ( - patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain, - patch("arcade_google.tools.utils._get_email_html_body") as mock_html, - patch("arcade_google.tools.utils._clean_email_body") as mock_clean, + patch("arcade_google.utils._get_email_plain_text_body") as mock_plain, + patch("arcade_google.utils._get_email_html_body") as mock_html, + patch("arcade_google.utils._clean_email_body") as mock_clean, ): # Mock the helper functions mock_plain.return_value = "This is a test email." @@ -731,9 +731,9 @@ def test_parse_multipart_email_plain_only(): } with ( - patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain, - patch("arcade_google.tools.utils._get_email_html_body") as mock_html, - patch("arcade_google.tools.utils._clean_email_body") as mock_clean, + patch("arcade_google.utils._get_email_plain_text_body") as mock_plain, + patch("arcade_google.utils._get_email_html_body") as mock_html, + patch("arcade_google.utils._clean_email_body") as mock_clean, ): # Mock the helper functions mock_plain.return_value = "Plain text only email." @@ -776,9 +776,9 @@ def test_parse_multipart_email_html_only(): } with ( - patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain, - patch("arcade_google.tools.utils._get_email_html_body") as mock_html, - patch("arcade_google.tools.utils._clean_email_body") as mock_clean, + patch("arcade_google.utils._get_email_plain_text_body") as mock_plain, + patch("arcade_google.utils._get_email_html_body") as mock_html, + patch("arcade_google.utils._clean_email_body") as mock_clean, ): # Mock the helper functions mock_plain.return_value = None @@ -844,9 +844,9 @@ def test_parse_multipart_email_missing_headers(): } with ( - patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain, - patch("arcade_google.tools.utils._get_email_html_body") as mock_html, - patch("arcade_google.tools.utils._clean_email_body") as mock_clean, + patch("arcade_google.utils._get_email_plain_text_body") as mock_plain, + patch("arcade_google.utils._get_email_html_body") as mock_html, + patch("arcade_google.utils._clean_email_body") as mock_clean, ): # Mock the helper functions mock_plain.return_value = "Timeel" @@ -888,9 +888,9 @@ def test_parse_multipart_email_missing_fields(): } with ( - patch("arcade_google.tools.utils._get_email_plain_text_body") as mock_plain, - patch("arcade_google.tools.utils._get_email_html_body") as mock_html, - patch("arcade_google.tools.utils._clean_email_body") as mock_clean, + patch("arcade_google.utils._get_email_plain_text_body") as mock_plain, + patch("arcade_google.utils._get_email_html_body") as mock_html, + patch("arcade_google.utils._clean_email_body") as mock_clean, ): # Mock the helper functions mock_plain.return_value = "Missing some headers."