Spaces:
Running
Running
File size: 2,942 Bytes
9c05550 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """
Knowledge Universe — Admin Usage Report (Upstash REST Version)
Run: python scripts/admin_usage_report.py
"""
import os
import json
from datetime import datetime, timezone
from upstash_redis import Redis
from dotenv import load_dotenv
# Load variables from .env securely
load_dotenv()
# Pull the variables from the environment safely
UPSTASH_URL = os.getenv("UPSTASH_URL")
UPSTASH_TOKEN = os.getenv("UPSTASH_TOKEN")
# Fail-safe: Alert you if the .env variables are missing
if not UPSTASH_URL or not UPSTASH_TOKEN:
raise ValueError("🚨 ERROR: UPSTASH_URL or UPSTASH_TOKEN is missing from your .env file!")
# Initialize the client securely
client = Redis(url=UPSTASH_URL, token=UPSTASH_TOKEN)
def get_current_month():
return datetime.now(timezone.utc).strftime("%Y-%m")
def run_report():
print(f"Connecting to Upstash REST API at {UPSTASH_URL}...")
# 1. Get all customer keys
customer_keys = client.keys("ku:key:*")
if not customer_keys:
print("\nNo customers found in the database yet.")
return
users = []
current_month = get_current_month()
# 2. Fetch data for each customer
for key in customer_keys:
raw_data = client.get(key)
if not raw_data:
continue
# Upstash REST client might auto-parse JSON, so we handle both dict and string
customer = raw_data if isinstance(raw_data, dict) else json.loads(raw_data)
c_id = customer.get("customer_id")
# Look up their usage for the current month
usage_key = f"ku:usage:{c_id}:{current_month}"
usage = client.get(usage_key)
calls_used = int(usage) if usage else 0
# Format the date
created_dt = datetime.fromisoformat(customer.get("created_at").replace('Z', '+00:00'))
date_str = created_dt.strftime("%b %d, %H:%M")
users.append({
"email": customer.get("email"),
"tier": customer.get("tier").upper(),
"used": calls_used,
"limit": customer.get("calls_limit"),
"joined": date_str
})
# 3. Sort by usage (highest first) to find the "Whales"
users.sort(key=lambda x: x["used"], reverse=True)
# 4. Print the beautiful table
print(f"\n{'='*80}")
print(f" KNOWLEDGE UNIVERSE USAGE REPORT ({current_month})")
print(f"{'='*80}")
print(f"{'Email':<30} | {'Tier':<8} | {'Used':<6} | {'Limit':<7} | {'Joined'}")
print("-" * 80)
total_calls = 0
for u in users:
total_calls += u['used']
used_str = f"{u['used']}"
if u['used'] > 50:
used_str = f"🔥 {u['used']}"
print(f"{u['email']:<30} | {u['tier']:<8} | {used_str:<6} | {u['limit']:<7} | {u['joined']}")
print("-" * 80)
print(f"Total Users: {len(users)} | Total API Calls This Month: {total_calls}\n")
if __name__ == "__main__":
run_report()
|