admin-valentin/arcade-mcp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
arcade-mcp/toolkits/google/evals/eval_google_drive.py
Renato Byrro ac0f5aa10c
Search Google Drive documents and retrieve contents (#265) 
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>
2025-03-07 18:42:12 -03:00

329 lines
12 KiB
Python
Raw Blame History

from arcade.sdk import ToolCatalog
from arcade.sdk.eval import (
BinaryCritic,
EvalRubric,
EvalSuite,
ExpectedToolCall,
tool_eval,
)
import arcade_google
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(
fail_threshold=0.9,
warn_threshold=0.95,
)
catalog = ToolCatalog()
catalog.add_module(arcade_google)
@tool_eval()
def get_file_tree_structure_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="get my google drive's file tree structure including shared drives",
user_message="get my google drive's file tree structure including shared drives",
expected_tool_calls=[
ExpectedToolCall(
func=get_file_tree_structure,
args={
"restrict_to_shared_drive_id": None,
"include_shared_drives": True,
"include_organization_domain_documents": False,
"order_by": None,
"limit": None,
},
)
],
critics=[
BinaryCritic(critic_field="include_shared_drives", weight=0.5),
BinaryCritic(critic_field="restrict_to_shared_drive_id", weight=0.5 / 4),
BinaryCritic(critic_field="include_organization_domain_documents", weight=0.5 / 4),
BinaryCritic(critic_field="order_by", weight=0.5 / 4),
BinaryCritic(critic_field="limit", weight=0.5 / 4),
],
)
suite.add_case(
name="get my google drive's file tree structure without shared drives",
user_message="get my google drive's file tree structure without shared drives",
expected_tool_calls=[
ExpectedToolCall(
func=get_file_tree_structure,
args={
"restrict_to_shared_drive_id": None,
"include_shared_drives": False,
"include_organization_domain_documents": False,
"order_by": None,
"limit": None,
},
)
],
critics=[
BinaryCritic(critic_field="include_shared_drives", weight=0.5),
BinaryCritic(critic_field="restrict_to_shared_drive_id", weight=0.5 / 4),
BinaryCritic(critic_field="include_organization_domain_documents", weight=0.5 / 4),
BinaryCritic(critic_field="order_by", weight=0.5 / 4),
BinaryCritic(critic_field="limit", weight=0.5 / 4),
],
)
suite.add_case(
name="what are the files in the folder 'hello world' in my google drive?",
user_message="what are the files in the folder 'hello world' in my google drive?",
expected_tool_calls=[
ExpectedToolCall(
func=get_file_tree_structure,
args={
"restrict_to_shared_drive_id": None,
"include_shared_drives": False,
"include_organization_domain_documents": False,
"order_by": None,
"limit": None,
},
)
],
critics=[
BinaryCritic(critic_field="include_shared_drives", weight=0.5),
BinaryCritic(critic_field="restrict_to_shared_drive_id", weight=0.5 / 4),
BinaryCritic(critic_field="include_organization_domain_documents", weight=0.5 / 4),
BinaryCritic(critic_field="order_by", weight=0.5 / 4),
BinaryCritic(critic_field="limit", weight=0.5 / 4),
],
)
suite.add_case(
name="how many files are there in all my google drives, including shared ones?",
user_message="how many files are there in all my google drives, including shared ones?",
expected_tool_calls=[
ExpectedToolCall(
func=get_file_tree_structure,
args={
"restrict_to_shared_drive_id": None,
"include_shared_drives": True,
"include_organization_domain_documents": False,
"order_by": None,
"limit": None,
},
)
],
critics=[
BinaryCritic(critic_field="include_shared_drives", weight=0.5),
BinaryCritic(critic_field="restrict_to_shared_drive_id", weight=0.5 / 4),
BinaryCritic(critic_field="include_organization_domain_documents", weight=0.5 / 4),
BinaryCritic(critic_field="order_by", weight=0.5 / 4),
BinaryCritic(critic_field="limit", weight=0.5 / 4),
],
)
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
Reference in a new issue View git blame Copy permalink