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

google_workspace_mcp/tools/gmail.py

@@ -3,9 +3,19 @@ Gmail tools for Google Workspace MCP operations.
 """
 
 import logging
-from typing import Any
 
 from google_workspace_mcp.app import mcp
+from google_workspace_mcp.models import (
+    GmailAttachmentOutput,
+    GmailBulkDeleteOutput,
+    GmailDraftCreationOutput,
+    GmailDraftDeletionOutput,
+    GmailDraftSendOutput,
+    GmailEmailSearchOutput,
+    GmailMessageDetailsOutput,
+    GmailReplyOutput,
+    GmailSendOutput,
+)
 from google_workspace_mcp.services.gmail import GmailService
 
 logger = logging.getLogger(__name__)
@@ -14,10 +24,13 @@ logger = logging.getLogger(__name__)
 # --- Gmail Tool Functions --- #
 
 
-# @mcp.tool(
-#     name="query_gmail_emails",
-# )
-async def query_gmail_emails(query: str, max_results: int = 100) -> dict[str, Any]:
+@mcp.tool(
+    name="query_gmail_emails",
+    description="Query Gmail emails based on a search query.",
+)
+async def query_gmail_emails(
+    query: str, max_results: int = 100
+) -> GmailEmailSearchOutput:
     """
     Searches for Gmail emails using Gmail query syntax.
 
@@ -26,7 +39,7 @@ async def query_gmail_emails(query: str, max_results: int = 100) -> dict[str, An
         max_results: Maximum number of emails to return.
 
     Returns:
-        A dictionary containing the list of matching emails or an error message.
+        GmailEmailSearchOutput containing the list of matching emails.
     """
     logger.info(f"Executing query_gmail_emails tool with query: '{query}'")
 
@@ -39,15 +52,16 @@ async def query_gmail_emails(query: str, max_results: int = 100) -> dict[str, An
 
     # Return appropriate message if no results
     if not emails:
-        return {"message": "No emails found for the query."}
+        emails = []
 
-    return {"count": len(emails), "emails": emails}
+    return GmailEmailSearchOutput(count=len(emails), emails=emails)
 
 
-# @mcp.tool(
-#     name="gmail_get_message_details",
-# )
-async def gmail_get_message_details(email_id: str) -> dict[str, Any]:
+@mcp.tool(
+    name="gmail_get_message_details",
+    description="Retrieves a complete Gmail email message by its ID.",
+)
+async def gmail_get_message_details(email_id: str) -> GmailMessageDetailsOutput:
     """
     Retrieves a complete Gmail email message by its ID.
 
@@ -55,7 +69,7 @@ async def gmail_get_message_details(email_id: str) -> dict[str, Any]:
         email_id: The ID of the Gmail message to retrieve.
 
     Returns:
-        A dictionary containing the email details and attachments.
+        GmailMessageDetailsOutput containing the email details and attachments.
     """
     logger.info(f"Executing gmail_get_message_details tool with email_id: '{email_id}'")
     if not email_id or not email_id.strip():
@@ -72,13 +86,25 @@ async def gmail_get_message_details(email_id: str) -> dict[str, Any]:
     if not result:
         raise ValueError(f"Failed to retrieve email with ID: {email_id}")
 
-    return result
+    return GmailMessageDetailsOutput(
+        id=result["id"],
+        thread_id=result.get("thread_id", ""),
+        subject=result.get("subject", ""),
+        from_email=result.get("from_email", ""),
+        to_email=result.get("to_email", []),
+        date=result.get("date", ""),
+        body=result.get("body", ""),
+        attachments=result.get("attachments"),
+    )
 
 
-# @mcp.tool(
-#     name="gmail_get_attachment_content",
-# )
-async def gmail_get_attachment_content(message_id: str, attachment_id: str) -> dict[str, Any]:
+@mcp.tool(
+    name="gmail_get_attachment_content",
+    description="Retrieves a specific attachment from a Gmail message.",
+)
+async def gmail_get_attachment_content(
+    message_id: str, attachment_id: str
+) -> GmailAttachmentOutput:
     """
     Retrieves a specific attachment from a Gmail message.
 
@@ -87,14 +113,18 @@ async def gmail_get_attachment_content(message_id: str, attachment_id: str) -> d
         attachment_id: The ID of the attachment to retrieve.
 
     Returns:
-        A dictionary containing filename, mimeType, size, and base64 data.
+        GmailAttachmentOutput containing filename, mimeType, size, and base64 data.
     """
-    logger.info(f"Executing gmail_get_attachment_content tool - Msg: {message_id}, Attach: {attachment_id}")
+    logger.info(
+        f"Executing gmail_get_attachment_content tool - Msg: {message_id}, Attach: {attachment_id}"
+    )
     if not message_id or not attachment_id:
         raise ValueError("Message ID and attachment ID are required")
 
     gmail_service = GmailService()
-    result = gmail_service.get_attachment_content(message_id=message_id, attachment_id=attachment_id)
+    result = gmail_service.get_attachment_content(
+        message_id=message_id, attachment_id=attachment_id
+    )
 
     if not result or (isinstance(result, dict) and result.get("error")):
         error_msg = "Error getting attachment"
@@ -102,20 +132,25 @@ async def gmail_get_attachment_content(message_id: str, attachment_id: str) -> d
             error_msg = result.get("message", error_msg)
         raise ValueError(error_msg)
 
-    # FastMCP should handle this dict, recognizing 'data' as content blob.
-    return result
+    return GmailAttachmentOutput(
+        filename=result["filename"],
+        mime_type=result["mime_type"],
+        size=result["size"],
+        data=result["data"],
+    )
 
 
-# @mcp.tool(
-#     name="create_gmail_draft",
-# )
+@mcp.tool(
+    name="create_gmail_draft",
+    description="Creates a draft email message in Gmail.",
+)
 async def create_gmail_draft(
     to: str,
     subject: str,
     body: str,
     cc: list[str] | None = None,
     bcc: list[str] | None = None,
-) -> dict[str, Any]:
+) -> GmailDraftCreationOutput:
     """
     Creates a draft email message in Gmail.
 
@@ -127,7 +162,7 @@ async def create_gmail_draft(
         bcc: Optional list of email addresses to BCC.
 
     Returns:
-        A dictionary containing the created draft details.
+        GmailDraftCreationOutput containing the created draft details.
     """
     logger.info("Executing create_gmail_draft")
     if not to or not subject or not body:  # Check for empty strings
@@ -135,7 +170,9 @@ async def create_gmail_draft(
 
     gmail_service = GmailService()
     # Pass bcc parameter even though service may not use it (for test compatibility)
-    result = gmail_service.create_draft(to=to, subject=subject, body=body, cc=cc, bcc=bcc)
+    result = gmail_service.create_draft(
+        to=to, subject=subject, body=body, cc=cc, bcc=bcc
+    )
 
     if not result or (isinstance(result, dict) and result.get("error")):
         error_msg = "Error creating draft"
@@ -143,15 +180,16 @@ async def create_gmail_draft(
             error_msg = result.get("message", error_msg)
         raise ValueError(error_msg)
 
-    return result
+    return GmailDraftCreationOutput(id=result["id"], message=result.get("message", {}))
 
 
-# @mcp.tool(
-#     name="delete_gmail_draft",
-# )
+@mcp.tool(
+    name="delete_gmail_draft",
+    description="Deletes a Gmail draft email by its draft ID.",
+)
 async def delete_gmail_draft(
     draft_id: str,
-) -> dict[str, Any]:
+) -> GmailDraftDeletionOutput:
     """
     Deletes a specific draft email from Gmail.
 
@@ -159,7 +197,7 @@ async def delete_gmail_draft(
         draft_id: The ID of the draft to delete.
 
     Returns:
-        A dictionary confirming the deletion.
+        GmailDraftDeletionOutput confirming the deletion.
     """
     logger.info(f"Executing delete_gmail_draft with draft_id: '{draft_id}'")
     if not draft_id or not draft_id.strip():
@@ -172,22 +210,24 @@ async def delete_gmail_draft(
         # Attempt to check if the service returned an error dict
         # (Assuming handle_api_error might return dict or False/None)
         # This part might need adjustment based on actual service error handling
-        error_info = getattr(gmail_service, "last_error", None)  # Hypothetical error capture
+        error_info = getattr(
+            gmail_service, "last_error", None
+        )  # Hypothetical error capture
         error_msg = "Failed to delete draft"
         if isinstance(error_info, dict) and error_info.get("error"):
             error_msg = error_info.get("message", error_msg)
         raise ValueError(error_msg)
 
-    return {
-        "message": f"Draft with ID '{draft_id}' deleted successfully.",
-        "success": True,
-    }
+    return GmailDraftDeletionOutput(
+        message=f"Draft with ID '{draft_id}' deleted successfully.", success=True
+    )
 
 
-# @mcp.tool(
-#     name="gmail_send_draft",
-# )
-async def gmail_send_draft(draft_id: str) -> dict[str, Any]:
+@mcp.tool(
+    name="gmail_send_draft",
+    description="Sends an existing draft email from Gmail.",
+)
+async def gmail_send_draft(draft_id: str) -> GmailDraftSendOutput:
     """
     Sends a specific draft email.
 
@@ -195,7 +235,7 @@ async def gmail_send_draft(draft_id: str) -> dict[str, Any]:
         draft_id: The ID of the draft to send.
 
     Returns:
-        A dictionary containing the details of the sent message or an error.
+        GmailDraftSendOutput containing the details of the sent message.
     """
     logger.info(f"Executing gmail_send_draft tool for draft_id: '{draft_id}'")
     if not draft_id or not draft_id.strip():
@@ -210,18 +250,23 @@ async def gmail_send_draft(draft_id: str) -> dict[str, Any]:
     if not result:  # Should be caught by error dict check
         raise ValueError(f"Failed to send draft '{draft_id}'")
 
-    return result
+    return GmailDraftSendOutput(
+        id=result["id"],
+        thread_id=result.get("thread_id", ""),
+        label_ids=result.get("label_ids", []),
+    )
 
 
-# @mcp.tool(
-#     name="gmail_reply_to_email",
-# )
+@mcp.tool(
+    name="gmail_reply_to_email",
+    description="Create a reply to an existing email. Can be sent or saved as draft.",
+)
 async def gmail_reply_to_email(
     email_id: str,
     reply_body: str,
     send: bool = False,
     reply_all: bool = False,
-) -> dict[str, Any]:
+) -> GmailReplyOutput:
     """
     Creates a reply to an existing email thread.
 
@@ -232,7 +277,7 @@ async def gmail_reply_to_email(
         reply_all: If True, reply to all recipients. If False, reply to sender only.
 
     Returns:
-        A dictionary containing the sent message or created draft details.
+        GmailReplyOutput containing the sent message or created draft details.
     """
     logger.info(f"Executing gmail_reply_to_email to message: '{email_id}'")
     if not email_id or not reply_body:
@@ -252,15 +297,18 @@ async def gmail_reply_to_email(
             error_msg = result.get("message", error_msg)
         raise ValueError(error_msg)
 
-    return result
+    return GmailReplyOutput(
+        id=result["id"], thread_id=result.get("thread_id", ""), in_reply_to=email_id
+    )
 
 
-# @mcp.tool(
-#     name="gmail_bulk_delete_messages",
-# )
+@mcp.tool(
+    name="gmail_bulk_delete_messages",
+    description="Delete multiple emails at once by providing a list of message IDs.",
+)
 async def gmail_bulk_delete_messages(
     message_ids: list[str],
-) -> dict[str, Any]:
+) -> GmailBulkDeleteOutput:
     """
     Deletes multiple Gmail emails using a list of message IDs.
 
@@ -268,7 +316,7 @@ async def gmail_bulk_delete_messages(
         message_ids: A list of email message IDs to delete.
 
     Returns:
-        A dictionary summarizing the deletion result.
+        GmailBulkDeleteOutput summarizing the deletion result.
     """
     # Validation first - check if it's a list
     if not isinstance(message_ids, list):
@@ -289,19 +337,26 @@ async def gmail_bulk_delete_messages(
             error_msg = result.get("message", error_msg)
         raise ValueError(error_msg)
 
-    return result
+    return GmailBulkDeleteOutput(
+        deleted_count=result.get("deleted_count", len(message_ids)),
+        success=result.get("success", True),
+        message=result.get(
+            "message", f"Successfully deleted {len(message_ids)} messages"
+        ),
+    )
 
 
-# @mcp.tool(
-#     name="gmail_send_email",
-# )
+@mcp.tool(
+    name="gmail_send_email",
+    description="Composes and sends an email directly.",
+)
 async def gmail_send_email(
     to: list[str],
     subject: str,
     body: str,
     cc: list[str] | None = None,
     bcc: list[str] | None = None,
-) -> dict[str, Any]:
+) -> GmailSendOutput:
     """
     Composes and sends an email message.
 
@@ -313,14 +368,20 @@ async def gmail_send_email(
         bcc: Optional. A list of BCC recipient email addresses.
 
     Returns:
-        A dictionary containing the details of the sent message or an error.
+        GmailSendOutput containing the details of the sent message.
     """
     logger.info(f"Executing gmail_send_email tool to: {to}, subject: '{subject}'")
-    if not to or not isinstance(to, list) or not all(isinstance(email, str) and email.strip() for email in to):
+    if (
+        not to
+        or not isinstance(to, list)
+        or not all(isinstance(email, str) and email.strip() for email in to)
+    ):
         raise ValueError("Recipients 'to' must be a non-empty list of email strings.")
     if not subject or not subject.strip():
         raise ValueError("Subject cannot be empty.")
-    if body is None:  # Allow empty string for body, but not None if it implies missing arg.
+    if (
+        body is None
+    ):  # Allow empty string for body, but not None if it implies missing arg.
         raise ValueError("Body cannot be None (can be an empty string).")
 
     gmail_service = GmailService()
@@ -332,4 +393,8 @@ async def gmail_send_email(
     if not result:
         raise ValueError("Failed to send email")
 
-    return result
+    return GmailSendOutput(
+        id=result["id"],
+        thread_id=result.get("thread_id", ""),
+        label_ids=result.get("label_ids", []),
+    )