1e926da807
- Add audit status filter (Audited/Not Audited) to agent management and admin dashboard - Add token usage tracking: token_count per timeline entry, total_tokens on agents - Token badge on agent cards, Total Tokens stat in usage modal, dual-axis chart - Sort by Total Tokens option, total_tokens in CSV export - Mailgun email notifications for high token usage (optional, non-blocking) - Cooldown-based notification tracking in MongoDB token_notifications collection - Add promote_admin.py utility script for user promotion - Update CLAUDE.md documentation for all new features Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
107 lines
4.2 KiB
Python
107 lines
4.2 KiB
Python
import os
|
|
import requests
|
|
from datetime import datetime, timedelta
|
|
from database import notifications_collection, users_collection
|
|
|
|
|
|
def is_mailgun_configured() -> bool:
|
|
"""Check if Mailgun environment variables are set."""
|
|
return bool(os.getenv("MAILGUN_API_KEY")) and bool(os.getenv("MAILGUN_DOMAIN"))
|
|
|
|
|
|
def send_mailgun_email(to_emails: list[str], subject: str, html_body: str) -> bool:
|
|
"""Send email via Mailgun HTTP API. Returns True on success."""
|
|
api_key = os.getenv("MAILGUN_API_KEY")
|
|
domain = os.getenv("MAILGUN_DOMAIN")
|
|
from_email = os.getenv("MAILGUN_FROM_EMAIL", f"AgentHub <noreply@{domain}>")
|
|
|
|
response = requests.post(
|
|
f"https://api.mailgun.net/v3/{domain}/messages",
|
|
auth=("api", api_key),
|
|
data={
|
|
"from": from_email,
|
|
"to": to_emails,
|
|
"subject": subject,
|
|
"html": html_body,
|
|
},
|
|
timeout=10,
|
|
)
|
|
return response.status_code == 200
|
|
|
|
|
|
def build_threshold_email(agent_name: str, total_tokens: int, threshold: int) -> str:
|
|
"""Build HTML email body for a token threshold notification."""
|
|
return f"""
|
|
<html>
|
|
<body style="font-family: Arial, sans-serif; max-width: 600px; margin: 0 auto;">
|
|
<div style="background: linear-gradient(135deg, #f3ae3e, #e09520); padding: 20px; border-radius: 8px 8px 0 0;">
|
|
<h2 style="color: white; margin: 0;">AgentHub - High Token Usage Alert</h2>
|
|
</div>
|
|
<div style="padding: 20px; border: 1px solid #e2e8f0; border-top: none; border-radius: 0 0 8px 8px;">
|
|
<p>The following agent has exceeded the token usage threshold:</p>
|
|
<table style="width: 100%; border-collapse: collapse; margin: 16px 0;">
|
|
<tr>
|
|
<td style="padding: 8px; font-weight: bold; border-bottom: 1px solid #eee;">Agent Name</td>
|
|
<td style="padding: 8px; border-bottom: 1px solid #eee;">{agent_name}</td>
|
|
</tr>
|
|
<tr>
|
|
<td style="padding: 8px; font-weight: bold; border-bottom: 1px solid #eee;">Total Tokens</td>
|
|
<td style="padding: 8px; border-bottom: 1px solid #eee;">{total_tokens:,}</td>
|
|
</tr>
|
|
<tr>
|
|
<td style="padding: 8px; font-weight: bold; border-bottom: 1px solid #eee;">Threshold</td>
|
|
<td style="padding: 8px; border-bottom: 1px solid #eee;">{threshold:,}</td>
|
|
</tr>
|
|
</table>
|
|
<p style="color: #666; font-size: 0.9em;">
|
|
This is an automated notification from AgentHub. Please review the agent's token consumption.
|
|
</p>
|
|
</div>
|
|
</body>
|
|
</html>
|
|
"""
|
|
|
|
|
|
async def check_and_notify_threshold(agent_name: str, total_tokens: int):
|
|
"""Check token threshold and send notification if exceeded. Non-blocking, safe to call."""
|
|
if not is_mailgun_configured():
|
|
return
|
|
|
|
threshold = int(os.getenv("TOKEN_USAGE_THRESHOLD", "100000"))
|
|
if total_tokens < threshold:
|
|
return
|
|
|
|
cooldown_hours = int(os.getenv("NOTIFICATION_COOLDOWN_HOURS", "24"))
|
|
cooldown_cutoff = datetime.utcnow() - timedelta(hours=cooldown_hours)
|
|
|
|
# Check cooldown - skip if we already notified within the cooldown period
|
|
recent = await notifications_collection.find_one({
|
|
"agent_name": agent_name,
|
|
"sent_at": {"$gte": cooldown_cutoff},
|
|
})
|
|
if recent:
|
|
return
|
|
|
|
# Get admin user emails
|
|
admin_cursor = users_collection.find(
|
|
{"is_admin": True, "is_active": True},
|
|
{"email": 1},
|
|
)
|
|
admin_emails = [doc["email"] async for doc in admin_cursor]
|
|
if not admin_emails:
|
|
return
|
|
|
|
# Send email
|
|
subject = f"[AgentHub] High Token Usage: {agent_name} ({total_tokens:,} tokens)"
|
|
html_body = build_threshold_email(agent_name, total_tokens, threshold)
|
|
success = send_mailgun_email(admin_emails, subject, html_body)
|
|
|
|
# Record the notification attempt
|
|
await notifications_collection.insert_one({
|
|
"agent_name": agent_name,
|
|
"total_tokens": total_tokens,
|
|
"threshold": threshold,
|
|
"recipients": admin_emails,
|
|
"success": success,
|
|
"sent_at": datetime.utcnow(),
|
|
})
|