apply ruff formatting

This commit is contained in:
Taylor Wilsdon
2025-12-13 13:49:28 -08:00
parent 1d80a24ca4
commit 6b8352a354
50 changed files with 4010 additions and 2842 deletions

View File

@@ -69,7 +69,9 @@ def _extract_message_bodies(payload):
if body_data:
try:
decoded_data = base64.urlsafe_b64decode(body_data).decode("utf-8", errors="ignore")
decoded_data = base64.urlsafe_b64decode(body_data).decode(
"utf-8", errors="ignore"
)
if mime_type == "text/plain" and not text_body:
text_body = decoded_data
elif mime_type == "text/html" and not html_body:
@@ -84,7 +86,9 @@ def _extract_message_bodies(payload):
# Check the main payload if it has body data directly
if payload.get("body", {}).get("data"):
try:
decoded_data = base64.urlsafe_b64decode(payload["body"]["data"]).decode("utf-8", errors="ignore")
decoded_data = base64.urlsafe_b64decode(payload["body"]["data"]).decode(
"utf-8", errors="ignore"
)
mime_type = payload.get("mimeType", "")
if mime_type == "text/plain" and not text_body:
text_body = decoded_data
@@ -93,10 +97,7 @@ def _extract_message_bodies(payload):
except Exception as e:
logger.warning(f"Failed to decode main payload body: {e}")
return {
"text": text_body,
"html": html_body
}
return {"text": text_body, "html": html_body}
def _format_body_content(text_body: str, html_body: str) -> str:
@@ -115,7 +116,9 @@ def _format_body_content(text_body: str, html_body: str) -> str:
elif html_body.strip():
# Truncate very large HTML to keep responses manageable
if len(html_body) > HTML_BODY_TRUNCATE_LIMIT:
html_body = html_body[:HTML_BODY_TRUNCATE_LIMIT] + "\n\n[HTML content truncated...]"
html_body = (
html_body[:HTML_BODY_TRUNCATE_LIMIT] + "\n\n[HTML content truncated...]"
)
return f"[HTML Content Converted]\n{html_body}"
else:
return "[No readable content found]"
@@ -137,12 +140,14 @@ def _extract_attachments(payload: dict) -> List[Dict[str, Any]]:
"""Recursively search for attachments in message parts"""
# Check if this part is an attachment
if part.get("filename") and part.get("body", {}).get("attachmentId"):
attachments.append({
"filename": part["filename"],
"mimeType": part.get("mimeType", "application/octet-stream"),
"size": part.get("body", {}).get("size", 0),
"attachmentId": part["body"]["attachmentId"]
})
attachments.append(
{
"filename": part["filename"],
"mimeType": part.get("mimeType", "application/octet-stream"),
"size": part.get("body", {}).get("size", 0),
"attachmentId": part["body"]["attachmentId"],
}
)
# Recursively search sub-parts
if "parts" in part:
@@ -204,7 +209,7 @@ def _prepare_gmail_message(
"""
# Handle reply subject formatting
reply_subject = subject
if in_reply_to and not subject.lower().startswith('re:'):
if in_reply_to and not subject.lower().startswith("re:"):
reply_subject = f"Re: {subject}"
# Prepare the email
@@ -255,7 +260,9 @@ def _generate_gmail_web_url(item_id: str, account_index: int = 0) -> str:
return f"https://mail.google.com/mail/u/{account_index}/#all/{item_id}"
def _format_gmail_results_plain(messages: list, query: str, next_page_token: Optional[str] = None) -> str:
def _format_gmail_results_plain(
messages: list, query: str, next_page_token: Optional[str] = None
) -> str:
"""Format Gmail search results in clean, LLM-friendly plain text."""
if not messages:
return f"No messages found for query: '{query}'"
@@ -269,11 +276,13 @@ def _format_gmail_results_plain(messages: list, query: str, next_page_token: Opt
for i, msg in enumerate(messages, 1):
# Handle potential null/undefined message objects
if not msg or not isinstance(msg, dict):
lines.extend([
f" {i}. Message: Invalid message data",
" Error: Message object is null or malformed",
"",
])
lines.extend(
[
f" {i}. Message: Invalid message data",
" Error: Message object is null or malformed",
"",
]
)
continue
# Handle potential null/undefined values from Gmail API
@@ -318,7 +327,9 @@ def _format_gmail_results_plain(messages: list, query: str, next_page_token: Opt
# Add pagination info if there's a next page
if next_page_token:
lines.append("")
lines.append(f"📄 PAGINATION: To get the next page, call search_gmail_messages again with page_token='{next_page_token}'")
lines.append(
f"📄 PAGINATION: To get the next page, call search_gmail_messages again with page_token='{next_page_token}'"
)
return "\n".join(lines)
@@ -327,7 +338,11 @@ def _format_gmail_results_plain(messages: list, query: str, next_page_token: Opt
@handle_http_errors("search_gmail_messages", is_read_only=True, service_type="gmail")
@require_google_service("gmail", "gmail_read")
async def search_gmail_messages(
service, query: str, user_google_email: str, page_size: int = 10, page_token: Optional[str] = None
service,
query: str,
user_google_email: str,
page_size: int = 10,
page_token: Optional[str] = None,
) -> str:
"""
Searches messages in a user's Gmail account based on a query.
@@ -349,22 +364,15 @@ async def search_gmail_messages(
)
# Build the API request parameters
request_params = {
"userId": "me",
"q": query,
"maxResults": page_size
}
request_params = {"userId": "me", "q": query, "maxResults": page_size}
# Add page token if provided
if page_token:
request_params["pageToken"] = page_token
logger.info("[search_gmail_messages] Using page_token for pagination")
response = await asyncio.to_thread(
service.users()
.messages()
.list(**request_params)
.execute
service.users().messages().list(**request_params).execute
)
# Handle potential null response (but empty dict {} is valid)
@@ -384,12 +392,16 @@ async def search_gmail_messages(
logger.info(f"[search_gmail_messages] Found {len(messages)} messages")
if next_page_token:
logger.info("[search_gmail_messages] More results available (next_page_token present)")
logger.info(
"[search_gmail_messages] More results available (next_page_token present)"
)
return formatted_output
@server.tool()
@handle_http_errors("get_gmail_message_content", is_read_only=True, service_type="gmail")
@handle_http_errors(
"get_gmail_message_content", is_read_only=True, service_type="gmail"
)
@require_google_service("gmail", "gmail_read")
async def get_gmail_message_content(
service, message_id: str, user_google_email: str
@@ -472,7 +484,7 @@ async def get_gmail_message_content(
if attachments:
content_lines.append("\n--- ATTACHMENTS ---")
for i, att in enumerate(attachments, 1):
size_kb = att['size'] / 1024
size_kb = att["size"] / 1024
content_lines.append(
f"{i}. {att['filename']} ({att['mimeType']}, {size_kb:.1f} KB)\n"
f" Attachment ID: {att['attachmentId']}\n"
@@ -483,7 +495,9 @@ async def get_gmail_message_content(
@server.tool()
@handle_http_errors("get_gmail_messages_content_batch", is_read_only=True, service_type="gmail")
@handle_http_errors(
"get_gmail_messages_content_batch", is_read_only=True, service_type="gmail"
)
@require_google_service("gmail", "gmail_read")
async def get_gmail_messages_content_batch(
service,
@@ -581,7 +595,7 @@ async def get_gmail_messages_content_batch(
except ssl.SSLError as ssl_error:
if attempt < max_retries - 1:
# Exponential backoff: 1s, 2s, 4s
delay = 2 ** attempt
delay = 2**attempt
logger.warning(
f"[get_gmail_messages_content_batch] SSL error for message {mid} on attempt {attempt + 1}: {ssl_error}. Retrying in {delay}s..."
)
@@ -623,7 +637,9 @@ async def get_gmail_messages_content_batch(
to = headers.get("To", "")
cc = headers.get("Cc", "")
msg_output = f"Message ID: {mid}\nSubject: {subject}\nFrom: {sender}\n"
msg_output = (
f"Message ID: {mid}\nSubject: {subject}\nFrom: {sender}\n"
)
if to:
msg_output += f"To: {to}\n"
if cc:
@@ -647,12 +663,16 @@ async def get_gmail_messages_content_batch(
# Format body content with HTML fallback
body_data = _format_body_content(text_body, html_body)
msg_output = f"Message ID: {mid}\nSubject: {subject}\nFrom: {sender}\n"
msg_output = (
f"Message ID: {mid}\nSubject: {subject}\nFrom: {sender}\n"
)
if to:
msg_output += f"To: {to}\n"
if cc:
msg_output += f"Cc: {cc}\n"
msg_output += f"Web Link: {_generate_gmail_web_url(mid)}\n\n{body_data}\n"
msg_output += (
f"Web Link: {_generate_gmail_web_url(mid)}\n\n{body_data}\n"
)
output_messages.append(msg_output)
@@ -664,7 +684,9 @@ async def get_gmail_messages_content_batch(
@server.tool()
@handle_http_errors("get_gmail_attachment_content", is_read_only=True, service_type="gmail")
@handle_http_errors(
"get_gmail_attachment_content", is_read_only=True, service_type="gmail"
)
@require_google_service("gmail", "gmail_read")
async def get_gmail_attachment_content(
service,
@@ -703,7 +725,9 @@ async def get_gmail_attachment_content(
.execute
)
except Exception as e:
logger.error(f"[get_gmail_attachment_content] Failed to download attachment: {e}")
logger.error(
f"[get_gmail_attachment_content] Failed to download attachment: {e}"
)
return (
f"Error: Failed to download attachment. The attachment ID may have changed.\n"
f"Please fetch the message content again to get an updated attachment ID.\n\n"
@@ -711,12 +735,13 @@ async def get_gmail_attachment_content(
)
# Format response with attachment data
size_bytes = attachment.get('size', 0)
size_bytes = attachment.get("size", 0)
size_kb = size_bytes / 1024 if size_bytes else 0
base64_data = attachment.get('data', '')
base64_data = attachment.get("data", "")
# Check if we're in stateless mode (can't save files)
from auth.oauth_config import is_stateless_mode
if is_stateless_mode():
result_lines = [
"Attachment downloaded successfully!",
@@ -725,17 +750,19 @@ async def get_gmail_attachment_content(
"\n⚠️ Stateless mode: File storage disabled.",
"\nBase64-encoded content (first 100 characters shown):",
f"{base64_data[:100]}...",
"\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch."
"\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch.",
]
logger.info(f"[get_gmail_attachment_content] Successfully downloaded {size_kb:.1f} KB attachment (stateless mode)")
logger.info(
f"[get_gmail_attachment_content] Successfully downloaded {size_kb:.1f} KB attachment (stateless mode)"
)
return "\n".join(result_lines)
# Save attachment and generate URL
try:
from core.attachment_storage import get_attachment_storage, get_attachment_url
storage = get_attachment_storage()
# Try to get filename and mime type from message (optional - attachment IDs are ephemeral)
filename = None
mime_type = None
@@ -757,18 +784,18 @@ async def get_gmail_attachment_content(
break
except Exception:
# If we can't get metadata, use defaults
logger.debug(f"Could not fetch attachment metadata for {attachment_id}, using defaults")
logger.debug(
f"Could not fetch attachment metadata for {attachment_id}, using defaults"
)
# Save attachment
file_id = storage.save_attachment(
base64_data=base64_data,
filename=filename,
mime_type=mime_type
base64_data=base64_data, filename=filename, mime_type=mime_type
)
# Generate URL
attachment_url = get_attachment_url(file_id)
result_lines = [
"Attachment downloaded successfully!",
f"Message ID: {message_id}",
@@ -776,14 +803,19 @@ async def get_gmail_attachment_content(
f"\n📎 Download URL: {attachment_url}",
"\nThe attachment has been saved and is available at the URL above.",
"The file will expire after 1 hour.",
"\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch."
"\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch.",
]
logger.info(f"[get_gmail_attachment_content] Successfully saved {size_kb:.1f} KB attachment as {file_id}")
logger.info(
f"[get_gmail_attachment_content] Successfully saved {size_kb:.1f} KB attachment as {file_id}"
)
return "\n".join(result_lines)
except Exception as e:
logger.error(f"[get_gmail_attachment_content] Failed to save attachment: {e}", exc_info=True)
logger.error(
f"[get_gmail_attachment_content] Failed to save attachment: {e}",
exc_info=True,
)
# Fallback to showing base64 preview
result_lines = [
"Attachment downloaded successfully!",
@@ -793,7 +825,7 @@ async def get_gmail_attachment_content(
"\nBase64-encoded content (first 100 characters shown):",
f"{base64_data[:100]}...",
f"\nError: {str(e)}",
"\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch."
"\nNote: Attachment IDs are ephemeral. Always use IDs from the most recent message fetch.",
]
return "\n".join(result_lines)
@@ -808,13 +840,20 @@ async def send_gmail_message(
subject: str = Body(..., description="Email subject."),
body: str = Body(..., description="Email body content (plain text or HTML)."),
body_format: Literal["plain", "html"] = Body(
"plain", description="Email body format. Use 'plain' for plaintext or 'html' for HTML content."
"plain",
description="Email body format. Use 'plain' for plaintext or 'html' for HTML content.",
),
cc: Optional[str] = Body(None, description="Optional CC email address."),
bcc: Optional[str] = Body(None, description="Optional BCC email address."),
thread_id: Optional[str] = Body(None, description="Optional Gmail thread ID to reply within."),
in_reply_to: Optional[str] = Body(None, description="Optional Message-ID of the message being replied to."),
references: Optional[str] = Body(None, description="Optional chain of Message-IDs for proper threading."),
thread_id: Optional[str] = Body(
None, description="Optional Gmail thread ID to reply within."
),
in_reply_to: Optional[str] = Body(
None, description="Optional Message-ID of the message being replied to."
),
references: Optional[str] = Body(
None, description="Optional chain of Message-IDs for proper threading."
),
) -> str:
"""
Sends an email using the user's Gmail account. Supports both new emails and replies.
@@ -906,14 +945,21 @@ async def draft_gmail_message(
subject: str = Body(..., description="Email subject."),
body: str = Body(..., description="Email body (plain text)."),
body_format: Literal["plain", "html"] = Body(
"plain", description="Email body format. Use 'plain' for plaintext or 'html' for HTML content."
"plain",
description="Email body format. Use 'plain' for plaintext or 'html' for HTML content.",
),
to: Optional[str] = Body(None, description="Optional recipient email address."),
cc: Optional[str] = Body(None, description="Optional CC email address."),
bcc: Optional[str] = Body(None, description="Optional BCC email address."),
thread_id: Optional[str] = Body(None, description="Optional Gmail thread ID to reply within."),
in_reply_to: Optional[str] = Body(None, description="Optional Message-ID of the message being replied to."),
references: Optional[str] = Body(None, description="Optional chain of Message-IDs for proper threading."),
thread_id: Optional[str] = Body(
None, description="Optional Gmail thread ID to reply within."
),
in_reply_to: Optional[str] = Body(
None, description="Optional Message-ID of the message being replied to."
),
references: Optional[str] = Body(
None, description="Optional chain of Message-IDs for proper threading."
),
) -> str:
"""
Creates a draft email in the user's Gmail account. Supports both new drafts and reply drafts.
@@ -1115,7 +1161,9 @@ async def get_gmail_thread_content(
@server.tool()
@require_google_service("gmail", "gmail_read")
@handle_http_errors("get_gmail_threads_content_batch", is_read_only=True, service_type="gmail")
@handle_http_errors(
"get_gmail_threads_content_batch", is_read_only=True, service_type="gmail"
)
async def get_gmail_threads_content_batch(
service,
thread_ids: List[str],
@@ -1181,7 +1229,7 @@ async def get_gmail_threads_content_batch(
except ssl.SSLError as ssl_error:
if attempt < max_retries - 1:
# Exponential backoff: 1s, 2s, 4s
delay = 2 ** attempt
delay = 2**attempt
logger.warning(
f"[get_gmail_threads_content_batch] SSL error for thread {tid} on attempt {attempt + 1}: {ssl_error}. Retrying in {delay}s..."
)
@@ -1412,9 +1460,7 @@ async def list_gmail_filters(service, user_google_email: str) -> str:
if action.get("forward"):
action_lines.append(f"Forward to: {action['forward']}")
if action.get("removeLabelIds"):
action_lines.append(
f"Remove labels: {', '.join(action['removeLabelIds'])}"
)
action_lines.append(f"Remove labels: {', '.join(action['removeLabelIds'])}")
if action.get("addLabelIds"):
action_lines.append(f"Add labels: {', '.join(action['addLabelIds'])}")
@@ -1456,9 +1502,11 @@ async def create_gmail_filter(
filter_body = {"criteria": criteria, "action": action}
created_filter = await asyncio.to_thread(
service.users().settings().filters().create(
userId="me", body=filter_body
).execute
service.users()
.settings()
.filters()
.create(userId="me", body=filter_body)
.execute
)
filter_id = created_filter.get("id", "(unknown)")
@@ -1469,7 +1517,9 @@ async def create_gmail_filter(
@handle_http_errors("delete_gmail_filter", service_type="gmail")
@require_google_service("gmail", "gmail_settings_basic")
async def delete_gmail_filter(
service, user_google_email: str, filter_id: str = Field(..., description="ID of the filter to delete.")
service,
user_google_email: str,
filter_id: str = Field(..., description="ID of the filter to delete."),
) -> str:
"""
Deletes a Gmail filter by ID.
@@ -1509,8 +1559,12 @@ async def modify_gmail_message_labels(
service,
user_google_email: str,
message_id: str,
add_label_ids: List[str] = Field(default=[], description="Label IDs to add to the message."),
remove_label_ids: List[str] = Field(default=[], description="Label IDs to remove from the message."),
add_label_ids: List[str] = Field(
default=[], description="Label IDs to add to the message."
),
remove_label_ids: List[str] = Field(
default=[], description="Label IDs to remove from the message."
),
) -> str:
"""
Adds or removes labels from a Gmail message.
@@ -1561,8 +1615,12 @@ async def batch_modify_gmail_message_labels(
service,
user_google_email: str,
message_ids: List[str],
add_label_ids: List[str] = Field(default=[], description="Label IDs to add to messages."),
remove_label_ids: List[str] = Field(default=[], description="Label IDs to remove from messages."),
add_label_ids: List[str] = Field(
default=[], description="Label IDs to add to messages."
),
remove_label_ids: List[str] = Field(
default=[], description="Label IDs to remove from messages."
),
) -> str:
"""
Adds or removes labels from multiple Gmail messages in a single batch request.