ab889f9f1d
# PR Description
* Adds/updates the following files to all toolkits:
- `.pre-commit-config.yaml`
- `.ruff.toml`
- `LICENSE`
- `Makefile`
- `pyproject.toml`
* Lint all toolkits such that they pass `make check` and `make test` (a
total doozy). This includes adding some unit tests and evals.
* Github workflow for testing toolkits before merge into main (courtesy
of @sdreyer)
* Added a QOL improvement for tool developers for when they need to get
the context's auth token.
* Minor updates to `arcade new` template.
84 lines
3.2 KiB
Python
84 lines
3.2 KiB
Python
from typing import Annotated, Any, Optional
|
|
|
|
from arcade.sdk import ToolContext, tool
|
|
from arcade.sdk.auth import Google
|
|
|
|
from arcade_google.tools.utils import build_drive_service, remove_none_values
|
|
|
|
from .models import Corpora, OrderBy
|
|
|
|
|
|
# Implements: https://googleapis.github.io/google-api-python-client/docs/dyn/drive_v3.files.html#list
|
|
# Example `arcade chat` query: `list my 5 most recently modified documents`
|
|
# TODO: Support query with natural language. Currently, the tool expects a fully formed query
|
|
# string as input with the syntax defined here: https://developers.google.com/drive/api/guides/search-files
|
|
@tool(
|
|
requires_auth=Google(
|
|
scopes=["https://www.googleapis.com/auth/drive.readonly"],
|
|
)
|
|
)
|
|
async def list_documents(
|
|
context: ToolContext,
|
|
corpora: Annotated[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}
|