"""Extended admin helpers for per-user management. What lives here --------------- Pure collectors that the ``/admin/users`` table and the ``/admin/users/{user_id}`` drilldown page both consume, so they stay in lockstep. No business logic — just roll-ups that read from ``BotInstance``, ``AuditLog``, ``DeploymentLog`` etc. Returned shapes --------------- :: collect_user_usage(db, user) -> { "user": User, "bots_total": int, "bots_running": int, "bots_by_status": {status: count, ...}, "cpu_alloc_cores": float, "cpu_quota_cores": float, "cpu_percent": float, # 0..100 of quota used "ram_alloc_mb": int, "ram_quota_mb": int, "ram_percent": float, "storage_used_mb": int, "storage_quota_mb": int, "storage_percent": float, "bots_at_max": bool, # True if bots_total >= plan_max_bots } collect_user_traffic(db, user, hours=24) -> { "hours": int, "total_requests": int, "error_count": int, "avg_latency_ms": float, } The quota percentages are clamped to [0, 100] for the bar UI; values above 100 indicate the user is over their plan (still useful information — admin should see "100%+!" red bar). """ from __future__ import annotations import json from typing import Any from sqlalchemy import select, func, and_ from sqlalchemy.ext.asyncio import AsyncSession from models import ( AuditLog, BotInstance, BotStatus, DeploymentMode, User, ) def _percent(used: float, quota: float) -> float: """Used / quota * 100, clamped to [0, 200]. ``>100`` for over-plan.""" if quota <= 0: return 100.0 if used > 0 else 0.0 pct = round(used / quota * 100.0, 1) return min(200.0, max(0.0, pct)) async def collect_user_usage(db: AsyncSession, user: User) -> dict[str, Any]: """Compute current usage vs plan for one user.""" bots = list((await db.execute( select(BotInstance).where(BotInstance.owner_id == user.id) )).scalars().all()) bots_by_status: dict[str, int] = {} for b in bots: key = b.status.value if hasattr(b.status, "value") else str(b.status) bots_by_status[key] = bots_by_status.get(key, 0) + 1 # Live process metrics for multitentant bots; alloc for legacy ones # uses the bot's configured cpu/ram (allocated at create-time). cpu_alloc = 0.0 ram_alloc = 0 storage_used = 0 for b in bots: if b.deployment_mode == DeploymentMode.MULTITENANT: try: import bot_runner st = bot_runner.bot_status(b) if st.get("running"): cpu_alloc += float(st.get("cpu_percent") or 0.0) / 100.0 # cpu_percent -> cores (approx) ram_alloc += int(float(st.get("rss_mb") or 0)) except Exception: pass else: # Legacy / HF-Space bot: alloc from the configured slice. cpu_alloc += float(b.cpu_cores or 0.0) ram_alloc += int(b.ram_mb or 0) storage_used += int(b.storage_used_mb or 0) cpu_quota = float(user.cpu_quota_cores or 0.0) ram_quota = int(user.ram_quota_mb or 0) storage_quota = int(user.storage_quota_mb or 0) # CPU percent: number of allocated cores / quota. Capped at 200%. cpu_pct = _percent(cpu_alloc, cpu_quota) if cpu_quota > 0 else (100.0 if cpu_alloc > 0 else 0.0) return { "user": user, "bots_total": len(bots), "bots_running": sum(1 for b in bots if b.status == BotStatus.RUNNING), "bots_by_status": bots_by_status, "cpu_alloc_cores": round(cpu_alloc, 2), "cpu_quota_cores": cpu_quota, "cpu_percent": cpu_pct, "ram_alloc_mb": int(ram_alloc), "ram_quota_mb": ram_quota, "ram_percent": _percent(ram_alloc, ram_quota), "storage_used_mb": int(storage_used), "storage_quota_mb": storage_quota, "storage_percent": _percent(storage_used, storage_quota), "bots_at_max": (user.plan_max_bots > 0 and len(bots) >= user.plan_max_bots), } async def collect_user_traffic( db: AsyncSession, user: User, hours: int = 24, ) -> dict[str, Any]: """Recent proxy traffic for a user's bots.""" bot_ids = (await db.execute( select(BotInstance.id).where(BotInstance.owner_id == user.id) )).scalars().all() if not bot_ids: return { "hours": hours, "total_requests": 0, "error_count": 0, "avg_latency_ms": 0.0, } from datetime import datetime, timedelta end = datetime.utcnow() start = end - timedelta(hours=hours) rows = (await db.execute( select(AuditLog).where( and_( AuditLog.action == "proxy_request", AuditLog.target_type == "bot", AuditLog.target_id.in_(bot_ids), AuditLog.created_at >= start, AuditLog.created_at < end, ) ) )).scalars().all() total = errors = 0 latencies: list[float] = [] for r in rows: total += 1 try: details = json.loads(r.details_json) if r.details_json else {} except (ValueError, TypeError): details = {} if int(details.get("status") or 0) >= 500: errors += 1 lat = details.get("latency_ms") if isinstance(lat, (int, float)): latencies.append(float(lat)) return { "hours": hours, "total_requests": total, "error_count": errors, "avg_latency_ms": round(sum(latencies) / len(latencies), 1) if latencies else 0.0, } # --------------------------------------------------------------------------- # Plan presets — quick-apply buttons on the users page # --------------------------------------------------------------------------- PLAN_PRESETS: dict[str, dict[str, Any]] = { "free": {"cpu_cores": 0.4, "ram_mb": 512, "storage_mb": 2048, "max_bots": 1, "label": "Free"}, "trial": {"cpu_cores": 0.5, "ram_mb": 1024, "storage_mb": 4096, "max_bots": 2, "label": "Trial"}, "std": {"cpu_cores": 1.0, "ram_mb": 2 * 1024, "storage_mb": 10 * 1024, "max_bots": 5, "label": "Standard"}, "pro": {"cpu_cores": 4.0, "ram_mb": 8 * 1024, "storage_mb": 50 * 1024, "max_bots": 20, "label": "Pro"}, "biz": {"cpu_cores": 8.0, "ram_mb": 16 * 1024, "storage_mb": 200 * 1024, "max_bots": 50, "label": "Business"}, "unlim": {"cpu_cores": 16.0,"ram_mb": 64 * 1024, "storage_mb": 1024 * 1024,"max_bots": 1000,"label": "Unlimited"}, } __all__ = [ "collect_user_usage", "collect_user_traffic", "PLAN_PRESETS", ]