refactored all tools to use fastmcp native responses

This commit is contained in:
Taylor Wilsdon
2025-06-06 17:32:09 -04:00
parent 047c44b472
commit 100fb4399d
7 changed files with 210 additions and 452 deletions

View File

@@ -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}")