ac0f5aa10c
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>
391 lines
16 KiB
Python
391 lines
16 KiB
Python
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
from arcade.sdk.errors import ToolExecutionError
|
|
from googleapiclient.errors import HttpError
|
|
|
|
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
|
|
def mock_context():
|
|
context = AsyncMock()
|
|
context.authorization.token = "mock_token" # noqa: S105
|
|
return context
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_service():
|
|
with patch("arcade_google.tools.drive." + build_drive_service.__name__) as mock_build_service:
|
|
yield mock_build_service.return_value
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_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 = [
|
|
{
|
|
"files": [
|
|
{"id": "file1", "name": "Document 1"},
|
|
{"id": "file2", "name": "Document 2"},
|
|
],
|
|
"nextPageToken": None,
|
|
}
|
|
]
|
|
|
|
result = await search_documents(mock_context, limit=2)
|
|
|
|
assert result["documents_count"] == 2
|
|
assert len(result["documents"]) == 2
|
|
assert result["documents"][0]["id"] == "file1"
|
|
assert result["documents"][1]["id"] == "file2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_documents_pagination(mock_context, mock_service):
|
|
# Simulate multiple pages
|
|
mock_service.files.return_value.list.return_value.execute.side_effect = [
|
|
{
|
|
"files": [{"id": f"file{i}", "name": f"Document {i}"} for i in range(1, 11)],
|
|
"nextPageToken": "token1",
|
|
},
|
|
{
|
|
"files": [{"id": f"file{i}", "name": f"Document {i}"} for i in range(11, 21)],
|
|
"nextPageToken": None,
|
|
},
|
|
]
|
|
|
|
result = await search_documents(mock_context, limit=15)
|
|
|
|
assert result["documents_count"] == 15
|
|
assert len(result["documents"]) == 15
|
|
assert result["documents"][0]["id"] == "file1"
|
|
assert result["documents"][-1]["id"] == "file15"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_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=f"Error in execution of {search_documents.__tool_name__}"
|
|
):
|
|
await search_documents(mock_context)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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=f"Error in execution of {search_documents.__tool_name__}"
|
|
):
|
|
await search_documents(mock_context)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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 = [
|
|
{
|
|
"files": [
|
|
{"id": "file1", "name": "Document 1"},
|
|
],
|
|
"nextPageToken": None,
|
|
}
|
|
]
|
|
|
|
result = await search_documents(
|
|
mock_context,
|
|
order_by=OrderBy.MODIFIED_TIME_DESC,
|
|
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)",
|
|
corpora=Corpora.DOMAIN.value,
|
|
pageSize=1,
|
|
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
|
|
):
|
|
files_list_sample, drives_get_sample = sample_drive_file_tree_request_responses
|
|
|
|
mock_service.files.return_value.list.return_value.execute.side_effect = [files_list_sample]
|
|
mock_service.drives.return_value.get.return_value.execute.side_effect = drives_get_sample
|
|
|
|
result = await get_file_tree_structure(mock_context, include_shared_drives=True)
|
|
|
|
expected_file_tree = {
|
|
"drives": [
|
|
{
|
|
"id": "0AFqcR6obkydtUk9PVA",
|
|
"name": "Shared Drive 1",
|
|
"children": [
|
|
{
|
|
"createdTime": "2025-02-26T00:27:45.526Z",
|
|
"id": "1dCOCdPxhTqiB3j3bWrIWM692ZbL8dyjt",
|
|
"mimeType": "application/vnd.google-apps.folder",
|
|
"modifiedTime": "2025-02-26T00:27:45.526Z",
|
|
"name": "shared-1-folder-1",
|
|
"children": [
|
|
{
|
|
"createdTime": "2025-02-26T00:28:20.571Z",
|
|
"id": "19WVyQndQsc0AxxfdrIt5CvDQd6r-BvpqnB8bWZoL7Xk",
|
|
"mimeType": "application/vnd.google-apps.document",
|
|
"modifiedTime": "2025-02-26T00:28:30.773Z",
|
|
"name": "shared-1-folder-1-doc-1",
|
|
"size": {
|
|
"unit": "bytes",
|
|
"value": 1024,
|
|
},
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"createdTime": "2025-02-26T00:27:19.287Z",
|
|
"id": "1didt_h-tDjuJ-dmYtHUSyOCPci30K_kSszvg0G3tKBM",
|
|
"mimeType": "application/vnd.google-apps.document",
|
|
"modifiedTime": "2025-02-26T00:27:26.079Z",
|
|
"name": "shared-1-doc-1",
|
|
"size": {
|
|
"unit": "bytes",
|
|
"value": 1024,
|
|
},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"name": "My Drive",
|
|
"children": [
|
|
{
|
|
"createdTime": "2025-01-24T06:34:22.305Z",
|
|
"id": "1vB6sv0MD0hYSraYvWU_fcci3GN_-Jf4g-LfyXdG8ZMo",
|
|
"mimeType": "application/vnd.google-apps.document",
|
|
"modifiedTime": "2025-02-25T21:54:30.632Z",
|
|
"name": "The Birth of MX Engineering",
|
|
"owners": [
|
|
{
|
|
"email": "one_new_tool_everyday@arcade.dev",
|
|
"name": "one_new_tool_everyday",
|
|
}
|
|
],
|
|
"size": {
|
|
"unit": "bytes",
|
|
"value": 6634,
|
|
},
|
|
},
|
|
{
|
|
"createdTime": "2025-02-25T17:57:46.036Z",
|
|
"id": "1gqioaHG53jPVeJN5gBpHoO-GWtwiJcLo",
|
|
"mimeType": "application/vnd.google-apps.folder",
|
|
"modifiedTime": "2025-02-25T17:57:46.036Z",
|
|
"name": "test folder 1",
|
|
"owners": [
|
|
{
|
|
"email": "one_new_tool_everyday@arcade.dev",
|
|
"name": "one_new_tool_everyday",
|
|
}
|
|
],
|
|
"children": [
|
|
{
|
|
"id": "1J92V9yvVWm_uNHq3CCY4wyG1H9B6iiwO",
|
|
"name": "test folder 1.1",
|
|
"mimeType": "application/vnd.google-apps.folder",
|
|
"createdTime": "2025-02-25T17:58:58.987Z",
|
|
"modifiedTime": "2025-02-25T17:58:58.987Z",
|
|
"owners": [
|
|
{
|
|
"email": "one_new_tool_everyday@arcade.dev",
|
|
"name": "one_new_tool_everyday",
|
|
}
|
|
],
|
|
"children": [
|
|
{
|
|
"id": "1wv2dmYo0skJTI59ZIcwH9vm-wt7psMwXTvihuEGeHeI",
|
|
"name": "test document 1.1.1",
|
|
"mimeType": "application/vnd.google-apps.document",
|
|
"createdTime": "2025-02-25T17:59:03.325Z",
|
|
"modifiedTime": "2025-02-25T17:59:11.445Z",
|
|
"owners": [
|
|
{
|
|
"email": "one_new_tool_everyday@arcade.dev",
|
|
"name": "one_new_tool_everyday",
|
|
}
|
|
],
|
|
"size": {
|
|
"unit": "bytes",
|
|
"value": 1024,
|
|
},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"id": "1DSmL7d07kjT6b6L-t4JIT06ElUbZ1q0K6_gEpn_UGZ8",
|
|
"name": "test document 1.2",
|
|
"mimeType": "application/vnd.google-apps.document",
|
|
"createdTime": "2025-02-25T17:58:38.628Z",
|
|
"modifiedTime": "2025-02-25T17:58:46.713Z",
|
|
"owners": [
|
|
{
|
|
"email": "one_new_tool_everyday@arcade.dev",
|
|
"name": "one_new_tool_everyday",
|
|
}
|
|
],
|
|
"size": {
|
|
"unit": "bytes",
|
|
"value": 1024,
|
|
},
|
|
},
|
|
{
|
|
"id": "1Fcxz7HsyO2Zyc-5DTD3zBQnaVrZwD29BP9KD9rPnYfE",
|
|
"name": "test document 1.1",
|
|
"mimeType": "application/vnd.google-apps.document",
|
|
"createdTime": "2025-02-25T17:57:53.850Z",
|
|
"modifiedTime": "2025-02-25T17:58:28.745Z",
|
|
"owners": [
|
|
{
|
|
"email": "one_new_tool_everyday@arcade.dev",
|
|
"name": "one_new_tool_everyday",
|
|
}
|
|
],
|
|
"size": {
|
|
"unit": "bytes",
|
|
"value": 1024,
|
|
},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"createdTime": "2025-02-18T20:48:52.786Z",
|
|
"id": "16PUe97yGQeOjQgrgd54iCoxzid4SEvu_J33P_ELd5r8",
|
|
"mimeType": "application/vnd.google-apps.presentation",
|
|
"modifiedTime": "2025-02-19T23:31:20.483Z",
|
|
"name": "Hello world presentation",
|
|
"owners": [
|
|
{
|
|
"email": "john.doe@arcade.dev",
|
|
"name": "john.doe",
|
|
}
|
|
],
|
|
"size": {
|
|
"unit": "bytes",
|
|
"value": 15774558,
|
|
},
|
|
},
|
|
{
|
|
"id": "1nG7lSvIyK05N9METPczVJa4iGgE7uoo-A6zpqjpUsDY",
|
|
"name": "Shared doc 1",
|
|
"mimeType": "application/vnd.google-apps.document",
|
|
"createdTime": "2025-02-19T18:51:44.622Z",
|
|
"modifiedTime": "2025-02-19T19:30:39.773Z",
|
|
"owners": [
|
|
{
|
|
"name": "theboss",
|
|
"email": "theboss@arcade.dev",
|
|
}
|
|
],
|
|
"size": {
|
|
"unit": "bytes",
|
|
"value": 2700,
|
|
},
|
|
},
|
|
],
|
|
},
|
|
]
|
|
}
|
|
|
|
assert result == expected_file_tree
|