google-workspace-mcp 1.1.5__py3-none-any.whl → 1.2.0__py3-none-any.whl

google_workspace_mcp/tools/drive.py

@@ -6,6 +6,16 @@ import logging
 from typing import Any
 
 from google_workspace_mcp.app import mcp  # Import from central app module
+from google_workspace_mcp.models import (
+    DriveFileContentOutput,
+    DriveFileDeletionOutput,
+    DriveFileUploadOutput,
+    DriveFolderCreationOutput,
+    DriveFolderFindOutput,
+    DriveFolderSearchOutput,
+    DriveSearchOutput,
+    DriveSharedDrivesOutput,
+)
 from google_workspace_mcp.services.drive import DriveService
 
 logger = logging.getLogger(__name__)
@@ -16,82 +26,74 @@ logger = logging.getLogger(__name__)
 
 @mcp.tool(
     name="drive_search_files",
+    description="""Search for files in Google Drive with intelligent query handling.
+
+QUERY FORMATS SUPPORTED:
+1. Simple text searches (recommended): "sprint planning meeting notes" - searches content and filenames
+2. Drive API syntax: name contains 'project' AND modifiedTime > '2024-01-01'
+
+CRITICAL LIMITATIONS:
+- Parentheses ( ) are NOT supported in file searches
+- Mixed syntax (text + operators) is not allowed
+- For OR logic with multiple terms, use separate tool calls
+
+VALID EXAMPLES:
+✅ "project documents 2024" (simple text)
+✅ name contains 'sprint' (API syntax)
+✅ fullText contains 'meeting' AND modifiedTime > '2024-01-01' (combined API syntax)
+✅ mimeType = 'application/pdf' (file type filter)
+
+INVALID EXAMPLES (will cause errors):
+❌ "project (sprint OR planning)" - parentheses not supported
+❌ "ArcLio sprint OR planning" - mixed text and operators
+❌ "meeting modifiedTime > '2024-01-01'" - mixed syntax
+
+COMMON FILE TYPE FILTERS:
+- Google Docs: mimeType = 'application/vnd.google-apps.document'
+- PDFs: mimeType = 'application/pdf'
+- Images: mimeType contains 'image/'
+
+DATE FORMAT: Use RFC 3339 format like '2024-01-01' for date searches.""",
 )
 async def drive_search_files(
     query: str,
     page_size: int = 10,
     shared_drive_id: str | None = None,
-    include_shared_drives: bool = True,
-    include_trashed: bool = False,
-) -> dict[str, Any]:
+) -> DriveSearchOutput:
     """
-    Search for files in Google Drive with optional shared drive support. Trashed files are excluded by default.
-
-
-    Examples:
-    - "budget report" → works as-is
-    - "John's Documents" → automatically handled
+    Search for files in Google Drive.
 
     Args:
-        query: Search query string. Can be a simple text search or complex query with operators.
-               Apostrophes are automatically escaped for you.
-        page_size: Maximum number of files to return (1 to 1000, default 10).
-        shared_drive_id: Optional shared drive ID to search within a specific shared drive.
-        include_shared_drives: Whether to include shared drives and folders in search (default True).
-                              Set to False to search only personal files.
-        include_trashed: Whether to include trashed files in search results (default False).
+        query: Search query - either simple text or valid Drive API syntax (no parentheses)
+        page_size: Maximum number of files to return (1 to 1000, default 10)
+        shared_drive_id: Optional shared drive ID to search within specific shared drive
 
     Returns:
-        A dictionary containing a list of files or an error message.
+        DriveSearchOutput containing a list of matching files with metadata
     """
     logger.info(
-        f"Executing drive_search_files with query: '{query}', page_size: {page_size}, "
-        f"shared_drive_id: {shared_drive_id}, include_shared_drives: {include_shared_drives}, "
-        f"include_trashed: {include_trashed}"
+        f"Executing drive_search_files with query: '{query}', page_size: {page_size}, shared_drive_id: {shared_drive_id}"
     )
 
     if not query or not query.strip():
         raise ValueError("Query cannot be empty")
 
-    # Logic to build a robust query
-    # If the query looks like a simple term (no spaces, no operators), wrap it.
-    # Otherwise, assume the user has provided a full query expression.
-    clean_query = query.strip()
-    if (
-        " " not in clean_query
-        and ":" not in clean_query
-        and "=" not in clean_query
-        and ">" not in clean_query
-        and "<" not in clean_query
-    ):
-        # This is likely a simple term, wrap it for a full-text search.
-        final_query = f"fullText contains '{clean_query.replace("'", "\\'")}'"
-    else:
-        # Assume it's a complex query and use it as-is.
-        final_query = clean_query.replace("'", "\\'")
-
-    # Append the trashed filter
-    if not include_trashed:
-        final_query = f"{final_query} and trashed=false"
-
     drive_service = DriveService()
     files = drive_service.search_files(
-        query=final_query,
-        page_size=page_size,
-        shared_drive_id=shared_drive_id,
-        include_shared_drives=include_shared_drives,
+        query=query, page_size=page_size, shared_drive_id=shared_drive_id
     )
 
     if isinstance(files, dict) and files.get("error"):
         raise ValueError(f"Search failed: {files.get('message', 'Unknown error')}")
 
-    return {"files": files}
+    return DriveSearchOutput(files=files or [])
 
 
-# @mcp.tool(
-#     name="drive_read_file_content",
-# )
-async def drive_read_file_content(file_id: str) -> dict[str, Any]:
+@mcp.tool(
+    name="drive_read_file_content",
+    description="Read the content of a file from Google Drive.",
+)
+async def drive_read_file_content(file_id: str) -> DriveFileContentOutput:
     """
     Read the content of a file from Google Drive.
 
@@ -99,7 +101,7 @@ async def drive_read_file_content(file_id: str) -> dict[str, Any]:
         file_id: The ID of the file to read.
 
     Returns:
-        A dictionary containing the file content and metadata or an error.
+        DriveFileContentOutput containing the file content and metadata.
     """
     logger.info(f"Executing drive_read_file_content tool with file_id: '{file_id}'")
     if not file_id or not file_id.strip():
@@ -114,18 +116,24 @@ async def drive_read_file_content(file_id: str) -> dict[str, Any]:
     if isinstance(result, dict) and result.get("error"):
         raise ValueError(result.get("message", "Error reading file"))
 
-    return result
+    return DriveFileContentOutput(
+        file_id=result["file_id"],
+        name=result["name"],
+        content=result["content"],
+        mime_type=result["mime_type"],
+    )
 
 
-# @mcp.tool(
-#     name="drive_upload_file",
-# )
+@mcp.tool(
+    name="drive_upload_file",
+    description="Uploads a file to Google Drive by providing its content directly.",
+)
 async def drive_upload_file(
     filename: str,
     content_base64: str,
     parent_folder_id: str | None = None,
     shared_drive_id: str | None = None,
-) -> dict[str, Any]:
+) -> DriveFileUploadOutput:
     """
     Uploads a file to Google Drive using its base64 encoded content.
 
@@ -136,7 +144,7 @@ async def drive_upload_file(
         shared_drive_id: Optional shared drive ID to upload the file to a shared drive.
 
     Returns:
-        A dictionary containing the uploaded file metadata or an error.
+        DriveFileUploadOutput containing the uploaded file metadata.
     """
     logger.info(
         f"Executing drive_upload_file with filename: '{filename}', parent_folder_id: {parent_folder_id}, shared_drive_id: {shared_drive_id}"
@@ -157,17 +165,23 @@ async def drive_upload_file(
     if isinstance(result, dict) and result.get("error"):
         raise ValueError(result.get("message", "Error uploading file"))
 
-    return result
+    return DriveFileUploadOutput(
+        id=result["id"],
+        name=result["name"],
+        web_view_link=result["web_view_link"],
+        size=result["size"],
+    )
 
 
-# @mcp.tool(
-#     name="drive_create_folder",
-# )
+@mcp.tool(
+    name="drive_create_folder",
+    description="Create a new folder in Google Drive.",
+)
 async def drive_create_folder(
     folder_name: str,
     parent_folder_id: str | None = None,
     shared_drive_id: str | None = None,
-) -> dict[str, Any]:
+) -> DriveFolderCreationOutput:
     """
     Create a new folder in Google Drive.
 
@@ -177,7 +191,7 @@ async def drive_create_folder(
         shared_drive_id: Optional shared drive ID to create the folder in a shared drive.
 
     Returns:
-        A dictionary containing the created folder information.
+        DriveFolderCreationOutput containing the created folder information.
     """
     logger.info(
         f"Executing drive_create_folder with folder_name: '{folder_name}', parent_folder_id: {parent_folder_id}, shared_drive_id: {shared_drive_id}"
@@ -198,15 +212,18 @@ async def drive_create_folder(
             f"Folder creation failed: {result.get('message', 'Unknown error')}"
         )
 
-    return result
+    return DriveFolderCreationOutput(
+        id=result["id"], name=result["name"], web_view_link=result["web_view_link"]
+    )
 
 
-# @mcp.tool(
-#     name="drive_delete_file",
-# )
+@mcp.tool(
+    name="drive_delete_file",
+    description="Delete a file from Google Drive using its file ID.",
+)
 async def drive_delete_file(
     file_id: str,
-) -> dict[str, Any]:
+) -> DriveFileDeletionOutput:
     """
     Delete a file from Google Drive.
 
@@ -214,7 +231,7 @@ async def drive_delete_file(
         file_id: The ID of the file to delete.
 
     Returns:
-        A dictionary confirming the deletion or an error.
+        DriveFileDeletionOutput confirming the deletion.
     """
     logger.info(f"Executing drive_delete_file with file_id: '{file_id}'")
     if not file_id or not file_id.strip():
@@ -226,13 +243,18 @@ async def drive_delete_file(
     if isinstance(result, dict) and result.get("error"):
         raise ValueError(result.get("message", "Error deleting file"))
 
-    return result
+    return DriveFileDeletionOutput(
+        success=result.get("success", True),
+        message=result.get("message", f"File '{file_id}' deleted successfully"),
+        file_id=file_id,
+    )
 
 
 @mcp.tool(
     name="drive_list_shared_drives",
+    description="Lists shared drives accessible by the user.",
 )
-async def drive_list_shared_drives(page_size: int = 100) -> dict[str, Any]:
+async def drive_list_shared_drives(page_size: int = 100) -> DriveSharedDrivesOutput:
     """
     Lists shared drives (formerly Team Drives) that the user has access to.
 
@@ -240,8 +262,7 @@ async def drive_list_shared_drives(page_size: int = 100) -> dict[str, Any]:
         page_size: Maximum number of shared drives to return (1 to 100, default 100).
 
     Returns:
-        A dictionary containing a list of shared drives with their 'id' and 'name',
-        or an error message.
+        DriveSharedDrivesOutput containing a list of shared drives.
     """
     logger.info(f"Executing drive_list_shared_drives tool with page_size: {page_size}")
 
@@ -252,9 +273,9 @@ async def drive_list_shared_drives(page_size: int = 100) -> dict[str, Any]:
         raise ValueError(drives.get("message", "Error listing shared drives"))
 
     if not drives:
-        return {"message": "No shared drives found or accessible."}
+        drives = []
 
-    return {"count": len(drives), "shared_drives": drives}
+    return DriveSharedDrivesOutput(count=len(drives), shared_drives=drives)
 
 
 @mcp.tool(
@@ -264,7 +285,7 @@ async def drive_search_files_in_folder(
     folder_id: str,
     query: str = "",
     page_size: int = 10,
-) -> dict[str, Any]:
+) -> DriveFolderSearchOutput:
     """
     Search for files or folders within a specific folder ID. Trashed files are excluded.
     This works for both regular folders and Shared Drives (when using the Shared Drive's ID as the folder_id).
@@ -277,7 +298,7 @@ async def drive_search_files_in_folder(
         page_size: Maximum number of files to return (1 to 1000, default 10).
 
     Returns:
-        A dictionary containing a list of files and folders.
+        DriveFolderSearchOutput containing a list of files and folders.
     """
     logger.info(
         f"Executing drive_search_files_in_folder with folder_id: '{folder_id}', "
@@ -309,7 +330,7 @@ async def drive_search_files_in_folder(
             f"Folder search failed: {files.get('message', 'Unknown error')}"
         )
 
-    return {"folder_id": folder_id, "files": files}
+    return DriveFolderSearchOutput(folder_id=folder_id, files=files or [])
 
 
 # @mcp.tool(
@@ -358,7 +379,7 @@ async def drive_find_folder_by_name(
     file_query: str = "",
     page_size: int = 10,
     shared_drive_id: str | None = None,
-) -> dict[str, Any]:
+) -> DriveFolderFindOutput:
     """
     Finds folders by name using a two-step search: first an exact match, then a partial match.
     Automatically handles apostrophes in folder names and search queries. Trashed items are excluded.
@@ -375,7 +396,7 @@ async def drive_find_folder_by_name(
         shared_drive_id: Optional shared drive ID to search within a specific shared drive.
 
     Returns:
-        A dictionary containing folders_found and, if requested, file search results.
+        DriveFolderFindOutput containing folders found and optionally file search results.
     """
     logger.info(
         f"Executing drive_find_folder_by_name with folder_name: '{folder_name}', "
@@ -417,17 +438,17 @@ async def drive_find_folder_by_name(
             f"Folder search failed: {folders.get('message', 'Unknown error')}"
         )
 
-    result = {
-        "folder_name": folder_name,
-        "folders_found": folders,
-        "folder_count": len(folders) if folders else 0,
-    }
+    result = DriveFolderFindOutput(
+        folder_name=folder_name,
+        folders_found=folders or [],
+        folder_count=len(folders) if folders else 0,
+    )
 
     if not include_files:
         return result
 
     if not folders:
-        result["message"] = f"No folders found with name matching '{folder_name}'"
+        result.message = f"No folders found with name matching '{folder_name}'"
         return result
 
     target_folder = folders[0]
@@ -444,9 +465,8 @@ async def drive_find_folder_by_name(
             and ":" not in clean_file_query
             and "=" not in clean_file_query
         ):
-            wrapped_file_query = (
-                f"fullText contains '{clean_file_query.replace("'", "\\'")}'"
-            )
+            escaped_file_query = clean_file_query.replace("'", "\\'")
+            wrapped_file_query = f"fullText contains '{escaped_file_query}'"
         else:
             wrapped_file_query = clean_file_query.replace("'", "\\'")
         combined_query = f"{wrapped_file_query} and {folder_constraint}"
@@ -462,8 +482,8 @@ async def drive_find_folder_by_name(
             f"File search in folder failed: {files.get('message', 'Unknown error')}"
         )
 
-    result["target_folder"] = target_folder
-    result["files"] = files
-    result["file_count"] = len(files) if files else 0
+    result.target_folder = target_folder
+    result.files = files or []
+    result.file_count = len(files) if files else 0
 
     return result