Add Google Drive utility functions

plugins/_gmail_plugin/drive_tools.py

@@ -15,6 +15,8 @@
 from googleapiclient.discovery import build
 from googleapiclient.http import (
     MediaIoBaseUpload,
+    MediaFileUpload,
+
 )
 
 # Load the spaCy model
@@ -118,7 +120,7 @@ async def upload_file(drive_service, user_input):
         # If the folder_name is not 'root', find the folder ID
         if folder_name != 'root':
             # Search for the folder to get its ID
-            folder_id = 'KitchenSinkGPT'
+            folder_id = 'GPT_ALL'
             file_metadata['parents'] = [folder_id]
 
         file = drive_service.files().create(
@@ -208,6 +210,63 @@ async def list_files(drive_service, folder_name='root', max_results=10):
         return f"An error occurred while listing the files: {e}"
 
 
+async def create_folder(drive_service, folder_name, parent_id='root'):
+    """
+    Create a new folder in Google Drive.
+    """
+    file_metadata = {
+        'name': folder_name,
+        'mimeType': 'application/vnd.google-apps.folder',
+        'parents': [parent_id]
+    }
+    folder = drive_service.files().create(body=file_metadata, fields='id').execute()
+    logging.info("Created folder: %s", folder_name)
+    return folder.get('id')
+
+
+async def search_files(drive_service, query):
+    """
+    Search for files and folders in Google Drive.
+    """
+    results = drive_service.files().list(q=query, fields="files(id, name)").execute()
+    items = results.get('files', [])
+    return items
+
+
+async def share_file(drive_service, file_id, user_email, role='reader'):
+    """
+    Share a file or folder with a user.
+    """
+    user_permission = {
+        'type': 'user',
+        'role': role,
+        'emailAddress': user_email
+    }
+    drive_service.permissions().create(
+        fileId=file_id,
+        body=user_permission,
+        fields='id'
+    ).execute()
+    logging.info("Shared file %s with %s", file_id, user_email)
+
+
+async def move_file(drive_service, file_id, folder_id):
+    """
+    Move a file or folder to a different folder.
+    """
+    # Retrieve the existing parents to remove
+    file = drive_service.files().get(fileId=file_id, fields='parents').execute()
+    previous_parents = ",".join(file.get('parents'))
+    # Move the file to the new folder
+    file = drive_service.files().update(
+        fileId=file_id,
+        addParents=folder_id,
+        removeParents=previous_parents,
+        fields='id, parents'
+    ).execute()
+    logging.info("Moved file %s to folder %s", file_id, folder_id)
+
+
 drive_tools_list = [
     {
         "type": "function",
@@ -266,10 +325,100 @@ async def list_files(drive_service, folder_name='root', max_results=10):
             }
         },
     },
+    {
+        "type": "function",
+        "function": {
+            "name": "create_folder",
+            "description": "Create a new folder in Google Drive.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "folder_name": {
+                        "type": "string",
+                        "description": "The name of the folder to create.",
+                    },
+                    "parent_id": {
+                        "type": "string",
+                        "description": "The ID of the parent folder. Defaults to 'root' if not specified.",
+                        "default": "root"
+                    }
+                },
+                "required": ["folder_name"]
+            }
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "search_files",
+            "description": "Search for files and folders in Google Drive.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "The search query to use.",
+                    }
+                },
+                "required": ["query"]
+            }
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "share_file",
+            "description": "Share a file or folder with a user.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "file_id": {
+                        "type": "string",
+                        "description": "The ID of the file or folder to share.",
+                    },
+                    "user_email": {
+                        "type": "string",
+                        "description": "The email address of the user to share with.",
+                    },
+                    "role": {
+                        "type": "string",
+                        "description": "The role to grant to the user. Defaults to 'reader' if not specified.",
+                        "default": "reader"
+                    }
+                },
+                "required": ["file_id", "user_email"]
+            }
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "move_file",
+            "description": "Move a file or folder to a different folder.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "file_id": {
+                        "type": "string",
+                        "description": "The ID of the file or folder to move.",
+                    },
+                    "folder_id": {
+                        "type": "string",
+                        "description": "The ID of the folder to move the file or folder to.",
+                    }
+                },
+                "required": ["file_id", "folder_id"]
+            }
+        },
+    }
 ]
 
 available_functions = {
     "upload_file": upload_file,
     "download_file": download_file,
     "list_files": list_files,
+    "create_folder": create_folder,
+    "search_files": search_files,
+    "share_file": share_file,
+    "move_file": move_file,
 }