Read usage aggregations from LiteLLM Postgres directly (HTTP fallback kept)

app/main.py

@@ -167,6 +167,123 @@ LISTMONK_API_USER = os.environ.get("LISTMONK_API_USER", "portal-sync")
 LISTMONK_API_TOKEN = os.environ.get("LISTMONK_API_TOKEN", "")
 LISTMONK_LIST_ID = int(os.environ.get("LISTMONK_LIST_ID", "3"))  # "Members" list
 
+# Direct read access to LiteLLM's Postgres (for usage aggregations).
+# READ-ONLY purpose: never write to LiteLLM's tables — its HTTP API remains the
+# only interface for writes/management. Falls back to the HTTP path when unset.
+LITELLM_DB_HOST = os.environ.get("LITELLM_DB_HOST", "")
+LITELLM_DB_PORT = int(os.environ.get("LITELLM_DB_PORT", "5432"))
+LITELLM_DB_USER = os.environ.get("LITELLM_DB_USER", "")
+LITELLM_DB_PASSWORD = os.environ.get("LITELLM_DB_PASSWORD", "")
+LITELLM_DB_NAME = os.environ.get("LITELLM_DB_NAME", "")
+
+
+def litellm_db_ready() -> bool:
+    """True when LiteLLM DB read credentials are configured."""
+    return bool(LITELLM_DB_HOST and LITELLM_DB_USER and LITELLM_DB_NAME)
+
+
+def _pg_connect():
+    """Open a short-lived connection to LiteLLM's Postgres (read-only use)."""
+    import psycopg2
+    return psycopg2.connect(
+        host=LITELLM_DB_HOST, port=LITELLM_DB_PORT, user=LITELLM_DB_USER,
+        password=LITELLM_DB_PASSWORD, dbname=LITELLM_DB_NAME,
+        connect_timeout=10, sslmode="prefer",
+    )
+
+
+def aggregate_model_usage(days: int = 0, team_id: str | None = None) -> dict:
+    """Aggregate LiteLLM spend logs per model, straight from Postgres.
+
+    Returns {"models": [...], "totals": {...}} in the same shape the HTTP path
+    produced, so callers don't change. `days` bounds the window (0 = all time);
+    `team_id` restricts to one member's team (used by the dashboard path).
+
+    Prefers the `model_group` column, which carries our provider/model alias
+    (e.g. greenpt/green-r); falls back to `model` for pre-alias rows — matching
+    the logic the HTTP path used, expressed in SQL.
+    """
+    if not litellm_db_ready():
+        raise RuntimeError("LiteLLM DB credentials not configured")
+
+    where = []
+    params: list = []
+    if days and days > 0:
+        from datetime import datetime, timedelta, timezone
+        cutoff = (datetime.now(timezone.utc) - timedelta(days=days)).strftime(
+            "%Y-%m-%dT%H:%M:%S"
+        )
+        where.append("COALESCE(\"startTime\", '') >= %s")
+        params.append(cutoff)
+    if team_id:
+        where.append("team_id = %s")
+        params.append(team_id)
+    where_sql = ("WHERE " + " AND ".join(where)) if where else ""
+
+    sql = f"""
+        SELECT
+            COALESCE(NULLIF(model_group, ''), NULLIF(model, ''), 'unknown') AS model,
+            COALESCE(SUM(spend), 0) AS spend,
+            COALESCE(SUM(total_tokens), 0) AS tokens,
+            COUNT(*) AS requests
+        FROM "LiteLLM_SpendLogs"
+        {where_sql}
+        GROUP BY 1
+        ORDER BY 2 DESC
+    """
+
+    conn = _pg_connect()
+    try:
+        with conn.cursor() as cur:
+            cur.execute(sql, params)
+            rows = cur.fetchall()
+    finally:
+        conn.close()
+
+    models = [
+        {"model": r[0], "spend": float(r[1] or 0.0), "tokens": int(r[2] or 0),
+         "requests": int(r[3] or 0)}
+        for r in rows
+    ]
+    return {
+        "models": models,
+        "totals": {
+            "total_spend": sum(m["spend"] for m in models),
+            "total_tokens": sum(m["tokens"] for m in models),
+            "total_requests": sum(m["requests"] for m in models),
+        },
+    }
+
+
+def get_team_budget_from_db(team_id: str) -> dict | None:
+    """Fetch a team's spend/budget/reset directly from LiteLLM's Postgres.
+
+    Same data as the /team/list HTTP path (LiteLLM computes `spend`), without
+    downloading the whole team list. Returns None when the team is unknown.
+    """
+    if not litellm_db_ready():
+        return None
+    sql = """
+        SELECT max_budget, spend, budget_reset_at
+        FROM "LiteLLM_TeamTable"
+        WHERE team_id = %s
+        LIMIT 1
+    """
+    conn = _pg_connect()
+    try:
+        with conn.cursor() as cur:
+            cur.execute(sql, (team_id,))
+            row = cur.fetchone()
+    finally:
+        conn.close()
+    if not row:
+        return None
+    return {
+        "max_budget": float(row[0]) if row[0] is not None else None,
+        "spend": float(row[1] or 0.0),
+        "budget_reset_at": row[2],
+    }
+
 
 async def listmonk_sync(email: str, subscribe: bool) -> None:
     """Subscribe or unsubscribe a member to the newsletter list.
@@ -1330,9 +1447,20 @@ async def admin_model_usage(request: Request, days: int = 0):
     sorted by spend descending. Optional ?days=N limits to the last N days
     (0 or omitted = all-time). No per-member breakdown — this is co-op-level
     so admins can see which models the membership actually uses.
+
+    Reads go straight to LiteLLM's Postgres when configured (fast, bounded);
+    otherwise falls back to the HTTP API path.
     """
     _verify_admin_token(request)
 
+    # Fast path: direct Postgres aggregation.
+    if litellm_db_ready():
+        try:
+            return aggregate_model_usage(days=days)
+        except Exception as e:
+            logger.warning("admin model-usage: DB path failed (%s); falling back to HTTP", e)
+
+    # Fallback path: LiteLLM HTTP API.
     cutoff = None
     if days and days > 0:
         from datetime import datetime, timedelta, timezone
@@ -1479,12 +1607,14 @@ async def broker_revoke(name: str, request: Request):
     return {"status": "revoked", "name": name}
 
 
-async def broker_get_usage(email: str) -> dict:
+async def broker_get_usage(email: str, days: int = 0) -> dict:
     """Return the member's spend vs. balance, plus per-model and token detail.
 
-    Reads the member's team (spend + max_budget) and aggregates their recent
+    This is the read-only primitive the member dashboard uses (no master key on
-    spend logs by model for token counts. This is the read-only primitive the
+    the client). Reads go straight to LiteLLM's Postgres when configured (fast,
-    member dashboard uses (no master key on the client).
+    bounded); otherwise falls back to the HTTP API path. `days` bounds the
+    per-model breakdown window (0 = all time); balance/spend always come from
+    the team's live totals.
     """
     team_id = await litellm_get_or_create_team(email, get_member_balance(email))
 
@@ -1492,6 +1622,34 @@ async def broker_get_usage(email: str) -> dict:
     spend = 0.0
     reset_at = None
 
+    # Fast path: direct Postgres reads.
+    if litellm_db_ready():
+        try:
+            team = get_team_budget_from_db(team_id)
+            if team:
+                if team["max_budget"] is not None:
+                    balance = float(team["max_budget"])
+                spend = float(team["spend"] or 0.0)
+                reset_at = team["budget_reset_at"]
+            usage = aggregate_model_usage(days=days, team_id=team_id)
+            model_list = [
+                {"model": m["model"], "spend": m["spend"], "tokens": m["tokens"],
+                 "calls": m["requests"]}
+                for m in usage["models"]
+            ]
+            return {
+                "email": email,
+                "balance": balance,
+                "spend": spend,
+                "remaining": max(balance - spend, 0.0),
+                "total_tokens": usage["totals"]["total_tokens"],
+                "reset_at": reset_at,
+                "models": model_list,
+            }
+        except Exception as e:
+            logger.warning("broker_get_usage: DB path failed (%s); falling back to HTTP", e)
+
+    # Fallback path: LiteLLM HTTP API.
     # /team/list and /spend/logs have grown slow as usage accumulates; the
     # httpx default 5s read timeout was too short and made /broker/usage 500.
     client_timeout = httpx.Timeout(connect=15.0, read=60.0, write=15.0, pool=15.0)
@@ -1558,9 +1716,9 @@ async def broker_get_usage(email: str) -> dict:
 
 @app.get("/broker/usage")
 @limiter.limit("60/minute")
-async def broker_usage(request: Request):
+async def broker_usage(request: Request, days: int = 0):
     email = _verify_broker_secret(request)
-    return await broker_get_usage(email)
+    return await broker_get_usage(email, days=days)
 
 
 async def reconcile_memberships() -> dict:

requirements.txt

@@ -4,3 +4,4 @@ httpx==0.27.2
 pydantic==2.9.2
 PyJWT==2.9.0
 slowapi==0.1.9
+psycopg2-binary==2.9.9