arcade-google 0.1.5__py3-none-any.whl → 1.2.4__py3-none-any.whl

arcade_google/tools/drive.py

@@ -1,70 +1,194 @@
-from typing import Annotated, Optional
+from typing import Annotated, Any
 
-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 arcade_tdk import ToolContext, tool
+from arcade_tdk.auth import Google
+from googleapiclient.errors import HttpError
 
-from .models import Corpora, OrderBy
+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 import get_document_by_id
+from arcade_google.utils import (
+    build_drive_service,
+    build_file_tree,
+    build_file_tree_request_params,
+    build_files_list_params,
+)
+
+
+@tool(
+    requires_auth=Google(
+        scopes=["https://www.googleapis.com/auth/drive.file"],
+    )
+)
+async def get_file_tree_structure(
+    context: ToolContext,
+    include_shared_drives: Annotated[
+        bool, "Whether to include shared drives in the file tree structure. Defaults to False."
+    ] = False,
+    restrict_to_shared_drive_id: Annotated[
+        str | None,
+        "If provided, only include files from this shared drive in the file tree structure. "
+        "Defaults to None, which will include files and folders from all drives.",
+    ] = None,
+    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[
+        list[OrderBy] | None,
+        "Sort order. Defaults to listing the most recently modified documents first",
+    ] = None,
+    limit: Annotated[
+        int | None,
+        "The number of files and folders to list. Defaults to None, "
+        "which will list all files and folders.",
+    ] = None,
+) -> Annotated[
+    dict,
+    "A dictionary containing the file/folder tree structure in the user's Google Drive",
+]:
+    """
+    Get the file/folder tree structure of the user's Google Drive.
+    """
+    service = build_drive_service(
+        context.authorization.token if context.authorization and context.authorization.token else ""
+    )
+
+    keep_paginating = True
+    page_token = None
+    files = {}
+    file_tree: dict[str, list[dict]] = {"My Drive": []}
+
+    params = build_file_tree_request_params(
+        order_by,
+        page_token,
+        limit,
+        include_shared_drives,
+        restrict_to_shared_drive_id,
+        include_organization_domain_documents,
+    )
+
+    while keep_paginating:
+        # Get a list of files
+        results = service.files().list(**params).execute()
+
+        # Update page token
+        page_token = results.get("nextPageToken")
+        params["pageToken"] = page_token
+        keep_paginating = page_token is not None
+
+        for file in results.get("files", []):
+            files[file["id"]] = file
+
+    if not files:
+        return {"drives": []}
+
+    file_tree = build_file_tree(files)
+
+    drives = []
+
+    for drive_id, files in file_tree.items():  # type: ignore[assignment]
+        if drive_id == "My Drive":
+            drive = {"name": "My Drive", "children": files}
+        else:
+            try:
+                drive_details = service.drives().get(driveId=drive_id).execute()
+                drive_name = drive_details.get("name", "Shared Drive (name unavailable)")
+            except HttpError as e:
+                drive_name = (
+                    f"Shared Drive (name unavailable: 'HttpError {e.status_code}: {e.reason}')"
+                )
+
+            drive = {"name": drive_name, "id": drive_id, "children": files}
+
+        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
+# 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"],
+        scopes=["https://www.googleapis.com/auth/drive.file"],
     )
 )
-async def list_documents(
+async def search_documents(
     context: ToolContext,
-    corpora: Annotated[Optional[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"
+    document_contains: Annotated[
+        list[str] | None,
+        "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[
+        list[str] | None,
+        "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[
+        str | None,
+        "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[OrderBy],
+        list[OrderBy] | None,
         "Sort order. Defaults to listing the most recently modified documents first",
-    ] = OrderBy.MODIFIED_TIME_DESC,
-    supports_all_drives: Annotated[
-        Optional[bool],
-        "Whether the requesting application supports both My Drives and shared drives",
-    ] = False,
-    limit: Annotated[Optional[int], "The number of documents to list"] = 50,
+    ] = None,
+    limit: Annotated[int, "The number of documents to list"] = 50,
+    pagination_token: Annotated[
+        str | None, "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)",
+    "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.
+    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)
-    page_token = None  # The page token is used for continuing a previous request on the next page
-    files = []
-
-    service = build_drive_service(context.authorization.token)
-
-    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
+    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 page_token:
-            params["pageToken"] = page_token
+        if pagination_token:
+            params["pageToken"] = pagination_token
         else:
             params.pop("pageToken", None)
 
@@ -72,8 +196,92 @@ async def list_documents(
         batch = results.get("files", [])
         files.extend(batch[: limit - len(files)])
 
-        page_token = results.get("nextPageToken")
-        if not page_token or len(batch) < page_size:
+        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.file"],
+    )
+)
+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[
+        list[str] | None,
+        "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[
+        list[str] | None,
+        "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[
+        str | None,
+        "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[
+        list[OrderBy] | None,
+        "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[
+        str | None, "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}

arcade_google/tools/file_picker.py

@@ -0,0 +1,54 @@
+import base64
+import json
+from typing import Annotated
+
+from arcade_tdk import ToolContext, ToolMetadataKey, tool
+from arcade_tdk.auth import Google
+from arcade_tdk.errors import ToolExecutionError
+
+
+@tool(
+    requires_auth=Google(),
+    requires_metadata=[ToolMetadataKey.CLIENT_ID, ToolMetadataKey.COORDINATOR_URL],
+)
+def generate_google_file_picker_url(
+    context: ToolContext,
+) -> Annotated[dict, "Google File Picker URL for user file selection and permission granting"]:
+    """Generate a Google File Picker URL for user-driven file selection and authorization.
+
+    This tool generates a URL that directs the end-user to a Google File Picker interface where
+    where they can select or upload Google Drive files. Users can grant permission to access their
+    Drive files, providing a secure and authorized way to interact with their files.
+
+    This is particularly useful when prior tools (e.g., those accessing or modifying
+    Google Docs, Google Sheets, etc.) encountered failures due to file non-existence
+    (Requested entity was not found) or permission errors. Once the user completes the file
+    picker flow, the prior tool can be retried.
+    """
+    client_id = context.get_metadata(ToolMetadataKey.CLIENT_ID)
+    client_id_parts = client_id.split("-")
+    if not client_id_parts:
+        raise ToolExecutionError(
+            message="Invalid Google Client ID",
+            developer_message=f"Google Client ID '{client_id}' is not valid",
+        )
+    app_id = client_id_parts[0]
+    cloud_coordinator_url = context.get_metadata(ToolMetadataKey.COORDINATOR_URL).strip("/")
+
+    config = {
+        "auth": {
+            "client_id": client_id,
+            "app_id": app_id,
+        },
+    }
+    config_json = json.dumps(config)
+    config_base64 = base64.urlsafe_b64encode(config_json.encode("utf-8")).decode("utf-8")
+    url = f"{cloud_coordinator_url}/google/drive_picker?config={config_base64}"
+
+    return {
+        "url": url,
+        "llm_instructions": (
+            "Instruct the user to click the following link to open the Google Drive File Picker. "
+            "This will allow them to select files and grant access permissions: {url}"
+        ),
+    }