diff --git a/app/__pycache__/main.cpython-313.pyc b/app/__pycache__/main.cpython-313.pyc new file mode 100644 index 0000000..e558e47 Binary files /dev/null and b/app/__pycache__/main.cpython-313.pyc differ diff --git a/app/main.py b/app/main.py index dbf1784..3de3dd5 100644 --- a/app/main.py +++ b/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 - spend logs by model for token counts. This is the read-only primitive the - member dashboard uses (no master key on the client). + This is the read-only primitive the member dashboard uses (no master key on + the client). Reads go straight to LiteLLM's Postgres when configured (fast, + 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: diff --git a/requirements.txt b/requirements.txt index 4f514b5..5ed9ad2 100644 --- a/requirements.txt +++ b/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