admin-valentin/arcade-mcp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Search Google Drive documents and retrieve contents (#265)

Browse source 
This tool will be useful in scenarios akin to RAG, where someone wants
to ask questions or request the production of a summary, for instance,
about a bunch of documents related to a particular topic. Currently, to
fulfill such requests, the LLM needs to first `list_documents`, then
`get_document_by_id` for each document.

We also implement a utility functions to return documents in Markdown
and HTML, since the Drive API JSON is verbose and would waste too many
tokens unnecessarily.

Limitations: the Markdown/HTML utilities do not handle table of contents
(which I think aren't really useful here), headers, footers, or
footnotes.

---
This PR deprecates `list_documents` and implements `search_documents`,
apart from `search_and_retrieve_documents`). This configuration makes it
easier for LLMs to understand when to call each tool.

Both tools had their interfaces refactored to remove Google API-specific
arguments that were confusing LLMs sometimes, such as "corpora" and
"support_all_drives". It now accepts arguments that better relate to
expected user requests.

---------

Co-authored-by: Eric Gustin <eric@arcade.dev>
This commit is contained in:
Renato Byrro 2025-03-07 18:42:12 -03:00 committed by GitHub
parent 2135101acd
commit ac0f5aa10c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 1758 additions and 242 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
toolkits/google/arcade_google/tools/constants.py → toolkits/google/arcade_google/constants.py
View file

@ -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.

99
toolkits/google/arcade_google/doc_to_html.py Normal file
View file

@ -0,0 +1,99 @@
def convert_document_to_html(document: dict) -> str:
html = (
"<html><head>"
f"<title>{document['title']}</title>"
f'<meta name="documentId" content="{document["documentId"]}">'
"</head><body>"
)
for element in document["body"]["content"]:
html += convert_structural_element(element)
html += "</body></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", "<br>")
italic = style.get("italic", False)
bold = style.get("bold", False)
if italic:
content = f"<i>{content}</i>"
if bold:
content = f"<b>{content}</b>"
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 ("<p>", "</p>") if wrap_paragraphs else ("", "")
elif named_style == "TITLE":
return "<h1>", "</h1>"
elif named_style == "SUBTITLE":
return "<h2>", "</h2>"
elif named_style.startswith("HEADING_"):
try:
heading_level = int(named_style.split("_")[1])
except ValueError:
return ("<p>", "</p>") if wrap_paragraphs else ("", "")
else:
return f"<h{heading_level}>", f"</h{heading_level}>"
return ("<p>", "</p>") if wrap_paragraphs else ("", "")
def table_list_to_html(table: list[list[str]]) -> str:
html = "<table>"
for row in table:
html += "<tr>"
for cell in row:
if cell.endswith("<br>"):
cell = cell[:-4]
html += f"<td>{cell}</td>"
html += "</tr>"
html += "</table>"
return html

64
toolkits/google/arcade_google/doc_to_markdown.py Normal file
View file

@ -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 ""

0
toolkits/google/arcade_google/tools/exceptions.py → toolkits/google/arcade_google/exceptions.py
View file

6
toolkits/google/arcade_google/tools/models.py → toolkits/google/arcade_google/models.py
View file

@ -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
# ---------------------------------------------------------------------------- #

4
toolkits/google/arcade_google/tools/calendar.py
View file

@ -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(

4
toolkits/google/arcade_google/tools/contacts.py
View file

@ -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]

2
toolkits/google/arcade_google/tools/docs.py
View file

@ -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

264
toolkits/google/arcade_google/tools/drive.py
View file

@ -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}

8
toolkits/google/arcade_google/tools/gmail.py
View file

@ -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,

92
toolkits/google/arcade_google/tools/utils.py → toolkits/google/arcade_google/utils.py
View file

@ -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],

966
toolkits/google/conftest.py
View file

@ -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"
"<table>"
"<tr>"
"<td><b>Column 1</b></td>"
"<td><i>Another</i> column</td>"
"<td>Third column</td>"
"</tr>"
"<tr>"
"<td>Hello</td>"
"<td>world!</td>"
"<td></td>"
"</tr>"
"<tr>"
"<td>The quick brown fox</td>"
"<td><i>jumped</i> over</td>"
"<td>the lazy dog</td>"
"</tr>"
"</table>"
)
expected_html = (
"<html><head>"
"<title>The Birth of Machine Experience Engineering</title>"
'<meta name="documentId" content="1234567890">'
"</head><body>"
"<h1><b>The Birth of Machine Experience Engineering</b></h1>"
"<p>LLMs acting on behalf of humans and interacting with real-world systems isn't "
"theoretical anymore - "
"<b><i>Arcade has made it a reality.</i></b> With this shift, we're seeing the emergence "
"of a new software practice: <i>Machine Experience Engineering (MX Engineering)</i>.</p>"
"<table>"
"<tr>"
"<td><b>Column 1</b></td>"
"<td><i>Another</i> column</td>"
"<td>Third column</td>"
"</tr>"
"<tr>"
"<td>Hello</td>"
"<td>world!</td>"
"<td></td>"
"</tr>"
"<tr>"
"<td>The quick brown fox</td>"
"<td><i>jumped</i> over</td>"
"<td>the lazy dog</td>"
"</tr>"
"</table>"
"</body></html>"
)
return document, expected_markdown, expected_html

288
toolkits/google/evals/eval_google_drive.py
View file

@ -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

4
toolkits/google/evals/eval_google_gmail.py
View file

@ -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(

2
toolkits/google/tests/test_calendar.py
View file

@ -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

10
toolkits/google/tests/test_doc_to_markdown.py Normal file
View file

@ -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

2
toolkits/google/tests/test_docs.py
View file

@ -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

123
toolkits/google/tests/test_drive.py
View file

@ -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

60
toolkits/google/tests/test_gmail.py
View file

@ -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."