refactored all tools to use fastmcp native responses

2025-06-06 17:32:09 -04:00
parent 047c44b472
commit 100fb4399d
7 changed files with 210 additions and 452 deletions
--- a/gdrive/drive_tools.py
+++ b/gdrive/drive_tools.py
@@ -86,7 +86,7 @@ async def search_drive_files(
    drive_id: Optional[str] = None,
    include_items_from_all_drives: bool = True,
    corpora: Optional[str] = None,
-) -> types.CallToolResult:
+) -> str:
    """
    Searches for files and folders within a user's Google Drive, including shared drives.

@@ -101,9 +101,7 @@ async def search_drive_files(
                                 Otherwise, Drive API default behavior applies. Prefer 'user' or 'drive' over 'allDrives' for efficiency.

    Returns:
-        types.CallToolResult: Contains a list of found files/folders with their details (ID, name, type, size, modified time, link),
-                               an error message if the API call fails,
-                               or an authentication guidance message if credentials are required.
+        str: A formatted list of found files/folders with their details (ID, name, type, size, modified time, link).
    """
    tool_name = "search_drive_files"
    logger.info(f"[{tool_name}] Invoked. Email: '{user_google_email}', Query: '{query}'")
@@ -146,7 +144,7 @@ async def search_drive_files(
        )
        files = results.get('files', [])
        if not files:
-            return types.CallToolResult(content=[types.TextContent(type="text", text=f"No files found for '{query}'.")])
+            return f"No files found for '{query}'."

        formatted_files_text_parts = [f"Found {len(files)} files for {user_google_email} matching '{query}':"]
        for item in files:
@@ -155,19 +153,19 @@ async def search_drive_files(
                f"- Name: \"{item['name']}\" (ID: {item['id']}, Type: {item['mimeType']}{size_str}, Modified: {item.get('modifiedTime', 'N/A')}) Link: {item.get('webViewLink', '#')}"
            )
        text_output = "\n".join(formatted_files_text_parts)
-        return types.CallToolResult(content=[types.TextContent(type="text", text=text_output)])
+        return text_output
    except HttpError as error:
        logger.error(f"API error searching Drive files: {error}", exc_info=True)
-        return types.CallToolResult(isError=True, content=[types.TextContent(type="text", text=f"API error: {error}")])
+        raise Exception(f"API error: {error}")
    except Exception as e:
        logger.exception(f"Unexpected error searching Drive files: {e}")
-        return types.CallToolResult(isError=True, content=[types.TextContent(type="text", text=f"Unexpected error: {e}")])
+        raise Exception(f"Unexpected error: {e}")

@server.tool()
 async def get_drive_file_content(
    user_google_email: str,
    file_id: str,
-) -> types.CallToolResult:
+) -> str:
    """
    Retrieves the content of a specific Google Drive file by ID, supporting files in shared drives.

@@ -181,7 +179,7 @@ async def get_drive_file_content(
        file_id: Drive file ID.

    Returns:
-        types.CallToolResult with plain-text content (or error info).
+        str: The file content as plain text with metadata header.
    """
    tool_name = "get_drive_file_content"
    logger.info(f"[{tool_name}] Invoked. File ID: '{file_id}'")
@@ -244,25 +242,17 @@ async def get_drive_file_content(
            f'File: "{file_name}" (ID: {file_id}, Type: {mime_type})\n'
            f'Link: {file_metadata.get("webViewLink", "#")}\n\n--- CONTENT ---\n'
        )
-        return types.CallToolResult(
-            content=[types.TextContent(type="text", text=header + body_text)]
-        )
+        return header + body_text

    except HttpError as error:
        logger.error(
            f"API error getting Drive file content for {file_id}: {error}",
            exc_info=True,
        )
-        return types.CallToolResult(
-            isError=True,
-            content=[types.TextContent(type="text", text=f"API error: {error}")],
-        )
+        raise Exception(f"API error: {error}")
    except Exception as e:
        logger.exception(f"Unexpected error getting Drive file content for {file_id}: {e}")
-        return types.CallToolResult(
-            isError=True,
-            content=[types.TextContent(type="text", text=f"Unexpected error: {e}")],
-        )
+        raise Exception(f"Unexpected error: {e}")


@server.tool()
@@ -273,7 +263,7 @@ async def list_drive_items(
    drive_id: Optional[str] = None,
    include_items_from_all_drives: bool = True,
    corpora: Optional[str] = None,
-) -> types.CallToolResult:
+) -> str:
    """
    Lists files and folders, supporting shared drives.
    If `drive_id` is specified, lists items within that shared drive. `folder_id` is then relative to that drive (or use drive_id as folder_id for root).
@@ -288,7 +278,7 @@ async def list_drive_items(
        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.

    Returns:
-        types.CallToolResult: Contains a list of files/folders or an error.
+        str: A formatted list of files/folders in the specified folder.
    """
    tool_name = "list_drive_items"
    logger.info(f"[{tool_name}] Invoked. Email: '{user_google_email}', Folder ID: '{folder_id}'")
@@ -320,7 +310,7 @@ async def list_drive_items(
        )
        files = results.get('files', [])
        if not files:
-            return types.CallToolResult(content=[types.TextContent(type="text", text=f"No items found in folder '{folder_id}'.")])
+            return f"No items found in folder '{folder_id}'."

        formatted_items_text_parts = [f"Found {len(files)} items in folder '{folder_id}' for {user_google_email}:"]
        for item in files:
@@ -329,13 +319,13 @@ async def list_drive_items(
                f"- Name: \"{item['name']}\" (ID: {item['id']}, Type: {item['mimeType']}{size_str}, Modified: {item.get('modifiedTime', 'N/A')}) Link: {item.get('webViewLink', '#')}"
            )
        text_output = "\n".join(formatted_items_text_parts)
-        return types.CallToolResult(content=[types.TextContent(type="text", text=text_output)])
+        return text_output
    except HttpError as error:
        logger.error(f"API error listing Drive items in folder {folder_id}: {error}", exc_info=True)
-        return types.CallToolResult(isError=True, content=[types.TextContent(type="text", text=f"API error: {error}")])
+        raise Exception(f"API error: {error}")
    except Exception as e:
        logger.exception(f"Unexpected error listing Drive items in folder {folder_id}: {e}")
-        return types.CallToolResult(isError=True, content=[types.TextContent(type="text", text=f"Unexpected error: {e}")])
+        raise Exception(f"Unexpected error: {e}")

@server.tool()
 async def create_drive_file(
@@ -344,7 +334,7 @@ async def create_drive_file(
    content: str,
    folder_id: str = 'root',
    mime_type: str = 'text/plain',
-) -> types.CallToolResult:
+) -> str:
    """
    Creates a new file in Google Drive, supporting creation within shared drives.

@@ -356,7 +346,7 @@ async def create_drive_file(
        mime_type (str): The MIME type of the file. Defaults to 'text/plain'.

    Returns:
-        A CallToolResult confirming creation or an error/auth guidance message.
+        str: Confirmation message of the successful file creation with file link.
    """
    tool_name = "create_drive_file"
    logger.info(f"[{tool_name}] Invoked. Email: '{user_google_email}', File Name: {file_name}, Folder ID: {folder_id}")
@@ -392,11 +382,11 @@ async def create_drive_file(
        link = created_file.get('webViewLink', 'No link available')
        confirmation_message = f"Successfully created file '{created_file.get('name', file_name)}' (ID: {created_file.get('id', 'N/A')}) in folder '{folder_id}' for {user_email}. Link: {link}"
        logger.info(f"Successfully created file. Link: {link}")
-        return types.CallToolResult(content=[types.TextContent(type="text", text=confirmation_message)])
+        return confirmation_message

    except HttpError as error:
        logger.error(f"API error creating Drive file '{file_name}': {error}", exc_info=True)
-        return types.CallToolResult(isError=True, content=[types.TextContent(type="text", text=f"API error: {error}")])
+        raise Exception(f"API error: {error}")
    except Exception as e:
        logger.exception(f"Unexpected error creating Drive file '{file_name}': {e}")
-        return types.CallToolResult(isError=True, content=[types.TextContent(type="text", text=f"Unexpected error: {e}")])
+        raise Exception(f"Unexpected error: {e}")