feat: Search with file type

gdrive/drive_helpers.py

@@ -229,6 +229,82 @@ def build_drive_list_params(
 
 SHORTCUT_MIME_TYPE = "application/vnd.google-apps.shortcut"
 FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
+
+# RFC 6838 token-style MIME type validation (safe for Drive query interpolation).
+MIME_TYPE_PATTERN = re.compile(r"^[A-Za-z0-9!#$&^_.+-]+/[A-Za-z0-9!#$&^_.+-]+$")
+
+# Mapping from friendly type names to Google Drive MIME types.
+# Raw MIME type strings (containing '/') are always accepted as-is.
+FILE_TYPE_MIME_MAP: Dict[str, str] = {
+    "folder": "application/vnd.google-apps.folder",
+    "folders": "application/vnd.google-apps.folder",
+    "document": "application/vnd.google-apps.document",
+    "doc": "application/vnd.google-apps.document",
+    "documents": "application/vnd.google-apps.document",
+    "docs": "application/vnd.google-apps.document",
+    "spreadsheet": "application/vnd.google-apps.spreadsheet",
+    "sheet": "application/vnd.google-apps.spreadsheet",
+    "spreadsheets": "application/vnd.google-apps.spreadsheet",
+    "sheets": "application/vnd.google-apps.spreadsheet",
+    "presentation": "application/vnd.google-apps.presentation",
+    "presentations": "application/vnd.google-apps.presentation",
+    "slide": "application/vnd.google-apps.presentation",
+    "slides": "application/vnd.google-apps.presentation",
+    "form": "application/vnd.google-apps.form",
+    "forms": "application/vnd.google-apps.form",
+    "drawing": "application/vnd.google-apps.drawing",
+    "drawings": "application/vnd.google-apps.drawing",
+    "pdf": "application/pdf",
+    "pdfs": "application/pdf",
+    "shortcut": "application/vnd.google-apps.shortcut",
+    "shortcuts": "application/vnd.google-apps.shortcut",
+    "script": "application/vnd.google-apps.script",
+    "scripts": "application/vnd.google-apps.script",
+    "site": "application/vnd.google-apps.site",
+    "sites": "application/vnd.google-apps.site",
+    "jam": "application/vnd.google-apps.jam",
+    "jamboard": "application/vnd.google-apps.jam",
+    "jamboards": "application/vnd.google-apps.jam",
+}
+
+
+def resolve_file_type_mime(file_type: str) -> str:
+    """
+    Resolve a friendly file type name or raw MIME type string to a Drive MIME type.
+
+    If `file_type` contains '/' it is returned as-is (treated as a raw MIME type).
+    Otherwise it is looked up in FILE_TYPE_MIME_MAP.
+
+    Args:
+        file_type: A friendly name ('folder', 'document', 'pdf', …) or a raw MIME
+                   type string ('application/vnd.google-apps.document', …).
+
+    Returns:
+        str: The resolved MIME type string.
+
+    Raises:
+        ValueError: If the value is not a recognised friendly name and contains no '/'.
+    """
+    normalized = file_type.strip()
+    if not normalized:
+        raise ValueError("file_type cannot be empty.")
+
+    if "/" in normalized:
+        normalized_mime = normalized.lower()
+        if not MIME_TYPE_PATTERN.fullmatch(normalized_mime):
+            raise ValueError(
+                f"Invalid MIME type '{file_type}'. Expected format like 'application/pdf'."
+            )
+        return normalized_mime
+    lower = normalized.lower()
+    if lower not in FILE_TYPE_MIME_MAP:
+        valid = ", ".join(sorted(FILE_TYPE_MIME_MAP.keys()))
+        raise ValueError(
+            f"Unknown file_type '{file_type}'. Pass a MIME type directly (e.g. "
+            f"'application/pdf') or use one of the friendly names: {valid}"
+        )
+    return FILE_TYPE_MIME_MAP[lower]
+
 BASE_SHORTCUT_FIELDS = (
     "id, mimeType, parents, shortcutDetails(targetId, targetMimeType)"
 )

gdrive/drive_tools.py

@@ -36,6 +36,7 @@
     format_permission_info,
     get_drive_image_url,
     resolve_drive_item,
+    resolve_file_type_mime,
     resolve_folder_id,
     validate_expiration_time,
     validate_share_role,
@@ -61,6 +62,7 @@ async def search_drive_files(
     drive_id: Optional[str] = None,
     include_items_from_all_drives: bool = True,
     corpora: Optional[str] = None,
+    file_type: Optional[str] = None,
     detailed: bool = True,
 ) -> str:
     """
@@ -76,14 +78,19 @@ async def search_drive_files(
         corpora (Optional[str]): Bodies of items to query (e.g., 'user', 'domain', 'drive', 'allDrives').
                                  If 'drive_id' is specified and 'corpora' is None, it defaults to 'drive'.
                                  Otherwise, Drive API default behavior applies. Prefer 'user' or 'drive' over 'allDrives' for efficiency.
+        file_type (Optional[str]): Restrict results to a specific file type. Accepts a friendly
+                                   name ('folder', 'document'/'doc', 'spreadsheet'/'sheet',
+                                   'presentation'/'slides', 'form', 'drawing', 'pdf', 'shortcut',
+                                   'script', 'site', 'jam'/'jamboard') or any raw MIME type
+                                   string (e.g. 'application/pdf'). Defaults to None (all types).
         detailed (bool): Whether to include size, modified time, and link in results. Defaults to True.
 
     Returns:
         str: A formatted list of found files/folders with their details (ID, name, type, and optionally size, modified time, link).
              Includes a nextPageToken line when more results are available.
     """
     logger.info(
-        f"[search_drive_files] Invoked. Email: '{user_google_email}', Query: '{query}'"
+        f"[search_drive_files] Invoked. Email: '{user_google_email}', Query: '{query}', file_type: '{file_type}'"
     )
 
     # Check if the query looks like a structured Drive query or free text
@@ -103,6 +110,11 @@ async def search_drive_files(
             f"[search_drive_files] Reformatting free text query '{query}' to '{final_query}'"
         )
 
+    if file_type is not None:
+        mime = resolve_file_type_mime(file_type)
+        final_query = f"({final_query}) and mimeType = '{mime}'"
+        logger.info(f"[search_drive_files] Added mimeType filter: '{mime}'")
+
     list_params = build_drive_list_params(
         query=final_query,
         page_size=page_size,
@@ -429,6 +441,7 @@ async def list_drive_items(
     drive_id: Optional[str] = None,
     include_items_from_all_drives: bool = True,
     corpora: Optional[str] = None,
+    file_type: Optional[str] = None,
     detailed: bool = True,
 ) -> str:
     """
@@ -444,19 +457,29 @@ async def list_drive_items(
         drive_id (Optional[str]): ID of the shared drive. If provided, the listing is scoped to this drive.
         include_items_from_all_drives (bool): Whether items from all accessible shared drives should be included if `drive_id` is not set. Defaults to True.
         corpora (Optional[str]): Corpus to query ('user', 'drive', 'allDrives'). If `drive_id` is set and `corpora` is None, 'drive' is used. If None and no `drive_id`, API defaults apply.
+        file_type (Optional[str]): Restrict results to a specific file type. Accepts a friendly
+                                   name ('folder', 'document'/'doc', 'spreadsheet'/'sheet',
+                                   'presentation'/'slides', 'form', 'drawing', 'pdf', 'shortcut',
+                                   'script', 'site', 'jam'/'jamboard') or any raw MIME type
+                                   string (e.g. 'application/pdf'). Defaults to None (all types).
         detailed (bool): Whether to include size, modified time, and link in results. Defaults to True.
 
     Returns:
         str: A formatted list of files/folders in the specified folder.
              Includes a nextPageToken line when more results are available.
     """
     logger.info(
-        f"[list_drive_items] Invoked. Email: '{user_google_email}', Folder ID: '{folder_id}'"
+        f"[list_drive_items] Invoked. Email: '{user_google_email}', Folder ID: '{folder_id}', File Type: '{file_type}'"
     )
 
     resolved_folder_id = await resolve_folder_id(service, folder_id)
     final_query = f"'{resolved_folder_id}' in parents and trashed=false"
 
+    if file_type is not None:
+        mime = resolve_file_type_mime(file_type)
+        final_query = f"({final_query}) and mimeType = '{mime}'"
+        logger.info(f"[list_drive_items] Added mimeType filter: '{mime}'")
+
     list_params = build_drive_list_params(
         query=final_query,
         page_size=page_size,