From 1e926da8075f172a2319f2029cbd8b15c3db9f8b Mon Sep 17 00:00:00 2001 From: nickviljoen Date: Sun, 15 Feb 2026 21:48:04 +0200 Subject: [PATCH] Add token tracking, audit filter, and high usage email notifications - 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 --- CLAUDE.md | 40 +++++++++++- database.py | 1 + main.py | 71 ++++++++++++++++----- models.py | 5 +- notifications.py | 107 ++++++++++++++++++++++++++++++++ promote_admin.py | 35 +++++++++++ templates/admin/dashboard.html | 16 ++++- templates/agent_management.html | 97 +++++++++++++++++++++-------- 8 files changed, 324 insertions(+), 48 deletions(-) create mode 100644 notifications.py create mode 100644 promote_admin.py diff --git a/CLAUDE.md b/CLAUDE.md index 5f56b6d..81f2048 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -27,6 +27,13 @@ Create `.env` file with: - `ALGORITHM`: JWT algorithm (default: HS256) - `ACCESS_TOKEN_EXPIRE_MINUTES`: Token expiration (default: 60) +#### Optional: Email Notifications (Mailgun) +- `MAILGUN_API_KEY`: Mailgun API key (if not set, notifications are silently disabled) +- `MAILGUN_DOMAIN`: Mailgun sending domain (e.g. your-domain.mailgun.org) +- `MAILGUN_FROM_EMAIL`: Sender address (default: `AgentHub `) +- `TOKEN_USAGE_THRESHOLD`: Token count that triggers an alert (default: 100000) +- `NOTIFICATION_COOLDOWN_HOURS`: Hours between repeat alerts for the same agent (default: 24) + ### Default Login Credentials - Admin: `admin@agenthub.com` / `admin123` - Test User: `test@example.com` / `testpass123` @@ -51,8 +58,11 @@ Create `.env` file with: **models.py**: Pydantic models for: - `AiAgent`: Core agent model with comprehensive fields +- `UsageTimelineEntry`: Daily usage data including `message_count` and `token_count` - `UserCreate/UserResponse`: User management models -- `AiAgentCreate/AiAgentResponse`: API request/response models +- `AiAgentCreate/AiAgentResponse`: API request/response models (includes `total_tokens`) +- `AgentCollectorCreate`: Collector API input model (includes `total_tokens`) +- `AgentUsageStatsResponse`: Usage statistics response (includes `total_tokens`) **crud.py**: Database operations using Motor (async MongoDB driver): - User CRUD: authentication, creation, management @@ -61,6 +71,13 @@ Create `.env` file with: - All operations use ObjectId for MongoDB document IDs **database.py**: MongoDB connection setup with Motor async client +- Collections: `users`, `agents`, `agent_usage`, `token_notifications` + +**notifications.py**: Optional Mailgun email notification system: +- `is_mailgun_configured()`: Returns False if env vars not set (gracefully disabled) +- `send_mailgun_email()`: POST to Mailgun HTTP API with 10s timeout +- `build_threshold_email()`: HTML email template for threshold alerts +- `check_and_notify_threshold()`: Checks token usage against threshold, enforces cooldown via `token_notifications` collection, sends to admin users **auth.py**: JWT authentication with: - bcrypt password hashing @@ -101,8 +118,25 @@ Located in `templates/` directory: - Status tracking (Active, Inactive, Development, Deprecated) - Rich metadata: tags, userbase, department, contact person - Search functionality across multiple fields +- Filtering by status and audit status (Audited / Not Audited) - Admin can view/manage all agents +### Token Usage Tracking +- `total_tokens` field on agents tracks cumulative LLM token consumption +- `token_count` per day in usage timeline entries alongside `message_count` +- Token badge displayed on agent cards (gold/coins icon) +- Usage modal shows Total Tokens stat alongside messages/conversations/users +- Dual-axis chart (messages left axis, tokens right axis) when token data exists +- Sort agents by Total Tokens +- CSV export includes `total_tokens` column + +### High Usage Email Notifications +- Entirely optional — silently disabled when Mailgun env vars are not set +- Triggered from the Agent Collector endpoint (POST `/agents`) when `total_tokens` exceeds threshold +- Non-blocking — notification failure never breaks the collector API +- Cooldown tracking in MongoDB `token_notifications` collection (default 24h, configurable) +- Sends to all active admin users' email addresses + ### User Management - User registration with validation - Admin user creation capabilities @@ -114,6 +148,7 @@ Located in `templates/` directory: - Async operations using Motor driver - Indexed queries for performance - Data aggregation for statistics +- Collections: `users`, `agents`, `agent_usage`, `token_notifications` ## Development Guidelines @@ -156,4 +191,5 @@ Key dependencies from requirements.txt: - **bcrypt**: Password encryption - **pydantic**: Data validation - **jinja2**: Template engine -- **python-multipart**: Form handling \ No newline at end of file +- **python-multipart**: Form handling +- **requests**: HTTP client (used for Mailgun API calls) \ No newline at end of file diff --git a/database.py b/database.py index ba09b45..fc34d2b 100644 --- a/database.py +++ b/database.py @@ -12,6 +12,7 @@ db = client[MONGODB_DBNAME] users_collection = db.get_collection("users") agents_collection = db.get_collection("agents") agent_usage_collection = db.get_collection("agent_usage") +notifications_collection = db.get_collection("token_notifications") async def check_database_health(): """Check MongoDB connection health""" diff --git a/main.py b/main.py index cba5b10..d01aea6 100644 --- a/main.py +++ b/main.py @@ -10,6 +10,7 @@ import models import auth import config import msal_auth +import notifications from datetime import datetime import os import re @@ -225,13 +226,14 @@ def create_agent_response(agent: dict) -> models.AiAgentResponse: risk_factor=agent.get("risk_factor"), last_edited_by=agent.get("last_edited_by"), created_by=agent["created_by"], - # Usage tracking fields (new) + # Usage tracking fields usage_timeline=agent.get("usage_timeline"), conversation_count=agent.get("conversation_count"), unique_users=agent.get("unique_users"), total_messages=agent.get("total_messages"), first_used=agent.get("first_used"), - last_used=agent.get("last_used") + last_used=agent.get("last_used"), + total_tokens=agent.get("total_tokens") ) async def verify_agent_collector_api_key(x_api_key: str = Header(alias="X-API-Key")): @@ -259,7 +261,7 @@ def map_agent_collector_to_internal(collector_data: models.AgentCollectorCreate) # Convert usage_timeline from Pydantic models to dicts usage_timeline = None if collector_data.usage_timeline: - usage_timeline = [{"date": entry.date, "message_count": entry.message_count} + usage_timeline = [{"date": entry.date, "message_count": entry.message_count, "token_count": entry.token_count} for entry in collector_data.usage_timeline] return { @@ -277,13 +279,14 @@ def map_agent_collector_to_internal(collector_data: models.AgentCollectorCreate) "agent_tags": collector_data.tags, "agent_metadata": collector_data.metadata, "url": collector_data.url, - # Usage tracking fields (new) + # Usage tracking fields "usage_timeline": usage_timeline, "conversation_count": collector_data.conversation_count, "unique_users": collector_data.unique_users, "total_messages": collector_data.total_messages, "first_used": collector_data.first_used, "last_used": collector_data.last_used, + "total_tokens": collector_data.total_tokens, } @@ -1082,7 +1085,8 @@ async def export_agents_csv(current_user: dict = Depends(require_admin)): "quality_audit_updated_by_name", "risk_factor", "last_edited_by", - "created_by" + "created_by", + "total_tokens" ] writer = csv.DictWriter(output, fieldnames=fieldnames) @@ -1116,7 +1120,8 @@ async def export_agents_csv(current_user: dict = Depends(require_admin)): "quality_audit_updated_by_name": agent.get("quality_audit_updated_by_name", ""), "risk_factor": str(agent.get("risk_factor", "")) if agent.get("risk_factor") is not None else "", "last_edited_by": agent.get("last_edited_by", ""), - "created_by": agent.get("created_by", "") + "created_by": agent.get("created_by", ""), + "total_tokens": str(agent.get("total_tokens", "")) if agent.get("total_tokens") is not None else "" } writer.writerow(row) @@ -1324,6 +1329,7 @@ async def create_agent_collector( "total_messages": internal_data.get("total_messages"), "first_used": internal_data.get("first_used"), "last_used": internal_data.get("last_used"), + "total_tokens": internal_data.get("total_tokens"), "updated_at": datetime.utcnow() } @@ -1340,6 +1346,13 @@ async def create_agent_collector( ) print(f"🔗 URL DEBUG: Agent '{agent.name}' - Update executed: matched={result.matched_count}, modified={result.modified_count}") + # Check token threshold and send notification if needed (non-blocking) + if internal_data.get("total_tokens"): + try: + await notifications.check_and_notify_threshold(agent.name, internal_data["total_tokens"]) + except Exception as notify_err: + print(f"Notification check failed for '{agent.name}': {notify_err}") + return models.AgentUsageTrackingResponse( status="usage_logged", message="Agent already exists, usage tracked", @@ -1357,7 +1370,14 @@ async def create_agent_collector( # Create agent using collector-specific function created_agent = await crud.create_agent_from_collector(internal_data) - + + # Check token threshold and send notification if needed (non-blocking) + if internal_data.get("total_tokens"): + try: + await notifications.check_and_notify_threshold(agent.name, internal_data["total_tokens"]) + except Exception as notify_err: + print(f"Notification check failed for '{agent.name}': {notify_err}") + return models.AgentCollectorResponse( status="success", message="Agent data collected successfully", @@ -1424,7 +1444,8 @@ async def get_agent_usage( last_usage=agent.get("last_used"), usage_by_period=usage_by_period, conversation_count=agent.get("conversation_count"), - unique_users=agent.get("unique_users") + unique_users=agent.get("unique_users"), + total_tokens=agent.get("total_tokens") ) else: # FALLBACK TO CALCULATED DATA (old system) @@ -1449,7 +1470,8 @@ async def get_agent_usage( last_usage=stats["last_usage"].isoformat() if stats["last_usage"] else None, usage_by_period=usage_by_period, conversation_count=None, # Old system doesn't track this - unique_users=None # Old system doesn't track this + unique_users=None, # Old system doesn't track this + total_tokens=None # Old system doesn't track this ) except HTTPException: raise @@ -1479,16 +1501,31 @@ async def get_agent_usage_chart( # USE PRE-COMPUTED TIMELINE (new system) usage_timeline = agent.get("usage_timeline", []) - # Format for Chart.js + # Format for Chart.js with dual datasets + datasets = [{ + "label": "Message Count", + "data": [entry["message_count"] for entry in usage_timeline], + "backgroundColor": "rgba(54, 162, 235, 0.2)", + "borderColor": "rgba(54, 162, 235, 1)", + "borderWidth": 1, + "yAxisID": "y" + }] + + # Add token dataset if any entries have token data + token_data = [entry.get("token_count", 0) for entry in usage_timeline] + if any(t > 0 for t in token_data): + datasets.append({ + "label": "Token Count", + "data": token_data, + "backgroundColor": "rgba(255, 193, 7, 0.2)", + "borderColor": "rgba(255, 193, 7, 1)", + "borderWidth": 1, + "yAxisID": "y1" + }) + chart_data = { "labels": [entry["date"] for entry in usage_timeline], - "datasets": [{ - "label": "Message Count", - "data": [entry["message_count"] for entry in usage_timeline], - "backgroundColor": "rgba(54, 162, 235, 0.2)", - "borderColor": "rgba(54, 162, 235, 1)", - "borderWidth": 1 - }] + "datasets": datasets } return chart_data diff --git a/models.py b/models.py index bf675eb..0ccb554 100644 --- a/models.py +++ b/models.py @@ -5,6 +5,7 @@ from typing import Optional, List class UsageTimelineEntry(BaseModel): date: str = Field(..., description="Date in YYYY-MM-DD format") message_count: int = Field(..., ge=0, description="Number of messages on this date") + token_count: int = Field(0, ge=0, description="Number of tokens consumed on this date") class AiAgent(BaseModel): @@ -131,6 +132,7 @@ class AiAgentResponse(BaseModel): total_messages: Optional[int] = None first_used: Optional[str] = None last_used: Optional[str] = None + total_tokens: Optional[int] = None # Agent Collector API Models (for compatibility with agent_collector app) class AgentCollectorCreate(BaseModel): @@ -158,6 +160,7 @@ class AgentCollectorCreate(BaseModel): total_messages: Optional[int] = Field(default=None, ge=0) first_used: Optional[str] = None # ISO 8601 datetime string last_used: Optional[str] = None # ISO 8601 datetime string + total_tokens: Optional[int] = Field(default=None, ge=0) class AgentCollectorResponse(BaseModel): status: str = "success" @@ -189,4 +192,4 @@ class AgentUsageStatsResponse(BaseModel): usage_by_period: dict conversation_count: Optional[int] = None unique_users: Optional[int] = None - + total_tokens: Optional[int] = None diff --git a/notifications.py b/notifications.py new file mode 100644 index 0000000..3a75c49 --- /dev/null +++ b/notifications.py @@ -0,0 +1,107 @@ +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 ") + + 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""" + + +
+

AgentHub - High Token Usage Alert

+
+
+

The following agent has exceeded the token usage threshold:

+ + + + + + + + + + + + + +
Agent Name{agent_name}
Total Tokens{total_tokens:,}
Threshold{threshold:,}
+

+ This is an automated notification from AgentHub. Please review the agent's token consumption. +

+
+ + + """ + + +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(), + }) diff --git a/promote_admin.py b/promote_admin.py new file mode 100644 index 0000000..0da2e0a --- /dev/null +++ b/promote_admin.py @@ -0,0 +1,35 @@ +"""One-time script to promote a user to admin. Run on the server where .env has the correct MONGODB_URI. + +Usage: python promote_admin.py nick.viljoen@brandtech.plus +""" +import asyncio +import sys +from dotenv import load_dotenv + +load_dotenv() + +from database import users_collection + + +async def promote(email: str): + user = await users_collection.find_one({"email": email}) + if not user: + print(f"User '{email}' not found in database.") + return + + if user.get("is_admin"): + print(f"User '{email}' is already an admin.") + return + + await users_collection.update_one( + {"email": email}, + {"$set": {"is_admin": True}}, + ) + print(f"User '{email}' promoted to admin successfully.") + + +if __name__ == "__main__": + if len(sys.argv) != 2: + print("Usage: python promote_admin.py ") + sys.exit(1) + asyncio.run(promote(sys.argv[1])) diff --git a/templates/admin/dashboard.html b/templates/admin/dashboard.html index 3ea73d2..db2d8ed 100644 --- a/templates/admin/dashboard.html +++ b/templates/admin/dashboard.html @@ -145,6 +145,11 @@ +
@@ -663,6 +668,7 @@ function setupEventListeners() { document.getElementById('userSearch').addEventListener('input', filterUsers); document.getElementById('agentSearch').addEventListener('input', filterAgents); document.getElementById('agentStatusFilter').addEventListener('change', filterAgents); + document.getElementById('agentAuditFilter').addEventListener('change', filterAgents); document.getElementById('editUserForm').addEventListener('submit', handleEditUserSubmit); document.getElementById('editAgentForm').addEventListener('submit', handleEditAgentSubmit); document.getElementById('createUserForm').addEventListener('submit', handleCreateUserSubmit); @@ -841,14 +847,18 @@ function filterUsers() { function filterAgents() { const searchTerm = document.getElementById('agentSearch').value.toLowerCase(); const statusFilter = document.getElementById('agentStatusFilter').value; - + const auditFilter = document.getElementById('agentAuditFilter').value; + let filtered = allAgents.filter(agent => { const matchesSearch = agent.agent_name.toLowerCase().includes(searchTerm) || (agent.agent_description || '').toLowerCase().includes(searchTerm); const matchesStatus = !statusFilter || agent.agent_status === statusFilter; - return matchesSearch && matchesStatus; + const matchesAudit = !auditFilter || + (auditFilter === 'audited' && agent.quality_audit_status) || + (auditFilter === 'not_audited' && !agent.quality_audit_status); + return matchesSearch && matchesStatus && matchesAudit; }); - + displayAgents(filtered); } diff --git a/templates/agent_management.html b/templates/agent_management.html index 5e1e6be..a71df4d 100644 --- a/templates/agent_management.html +++ b/templates/agent_management.html @@ -77,13 +77,13 @@
-
+
-
+
+
+ +
@@ -459,6 +467,7 @@ document.addEventListener('DOMContentLoaded', function() { function setupEventListeners() { document.getElementById('searchInput').addEventListener('input', filterAgents); document.getElementById('statusFilter').addEventListener('change', filterAgents); + document.getElementById('auditFilter').addEventListener('change', filterAgents); document.getElementById('sortBy').addEventListener('change', sortAndDisplayAgents); document.getElementById('refreshBtn').addEventListener('click', function() { // Reload both datasets to ensure counts are accurate @@ -616,6 +625,7 @@ function displayAgents(agentsToShow) { ${agent.quality_audit_status ? '' : ''} ${agent.quality_audit_status && agent.risk_factor ? getRiskFactorBadge(agent.risk_factor) : ''} ${agent.total_messages ? `${agent.total_messages.toLocaleString()}` : ''} + ${agent.total_tokens ? `${agent.total_tokens.toLocaleString()}` : ''} ${agent.unique_users ? `${agent.unique_users}` : ''} ${agent.url ? `Open Agent` : ''}
@@ -1053,13 +1063,17 @@ async function deleteAgent() { function filterAgents() { const searchTerm = document.getElementById('searchInput').value.toLowerCase(); const statusFilter = document.getElementById('statusFilter').value; + const auditFilter = document.getElementById('auditFilter').value; let filtered = agents.filter(agent => { const matchesSearch = agent.agent_name.toLowerCase().includes(searchTerm) || (agent.agent_description || '').toLowerCase().includes(searchTerm) || (agent.agent_tags || []).some(tag => tag.toLowerCase().includes(searchTerm)); const matchesStatus = !statusFilter || agent.agent_status === statusFilter; - return matchesSearch && matchesStatus; + const matchesAudit = !auditFilter || + (auditFilter === 'audited' && agent.quality_audit_status) || + (auditFilter === 'not_audited' && !agent.quality_audit_status); + return matchesSearch && matchesStatus && matchesAudit; }); displayAgents(filtered); @@ -1077,6 +1091,10 @@ function sortAndDisplayAgents() { const aVal = a.total_messages || 0; const bVal = b.total_messages || 0; return bVal - aVal; // Descending order + } else if (sortBy === 'total_tokens') { + const aVal = a.total_tokens || 0; + const bVal = b.total_tokens || 0; + return bVal - aVal; // Descending order } else if (sortBy === 'unique_users') { // High to low (most users first), treat null/undefined as 0 const aVal = a.unique_users || 0; @@ -1219,33 +1237,38 @@ async function loadUsageChart(agentName, period = 'daily', startDate = null, end // Display statistics const statsHtml = `
-
+
${stats.total_usage_count || 0}
Total Messages
-
+
-
${stats.conversation_count !== undefined ? stats.conversation_count : 'N/A'}
+
${stats.total_tokens != null ? stats.total_tokens.toLocaleString() : 'N/A'}
+ Total Tokens +
+
+
+
+
${stats.conversation_count != null ? stats.conversation_count : 'N/A'}
Conversations
-
+
-
${stats.unique_users !== undefined ? stats.unique_users : 'N/A'}
+
${stats.unique_users != null ? stats.unique_users : 'N/A'}
Unique Users
-
-
-
${stats.first_usage ? formatDate(stats.first_usage) : 'Never'}
- First Used -
-
-
+
+
+ First Used: ${stats.first_usage ? formatDate(stats.first_usage) : 'Never'} +
+
+
Last Used: ${stats.last_usage ? formatDate(stats.last_usage) : 'Never'}
@@ -1277,12 +1300,43 @@ async function loadUsageChart(agentName, period = 'daily', startDate = null, end function renderUsageChart(chartData) { const ctx = document.getElementById('usageChart').getContext('2d'); - + // Destroy existing chart if it exists if (usageChart) { usageChart.destroy(); } - + + // Check if we have a second dataset (tokens) for dual-axis + const hasDualAxis = chartData.datasets.length > 1 && chartData.datasets.some(ds => ds.yAxisID === 'y1'); + + const scales = { + y: { + beginAtZero: true, + position: 'left', + title: { + display: hasDualAxis, + text: 'Messages' + }, + ticks: { + stepSize: 1 + } + } + }; + + if (hasDualAxis) { + scales.y1 = { + beginAtZero: true, + position: 'right', + title: { + display: true, + text: 'Tokens' + }, + grid: { + drawOnChartArea: false + } + }; + } + usageChart = new Chart(ctx, { type: 'line', data: chartData, @@ -1295,14 +1349,7 @@ function renderUsageChart(chartData) { position: 'top' } }, - scales: { - y: { - beginAtZero: true, - ticks: { - stepSize: 1 - } - } - } + scales: scales } }); }