From 71d63955b370f23b4872b9a2f5809705dbbb9c93 Mon Sep 17 00:00:00 2001 From: inference-bot Date: Mon, 14 Sep 2026 10:41:26 -0600 Subject: [PATCH] Flexible balance: per-member balance (DB) synced to team budget; keys inherit team; add /admin/set-balance top-up endpoint --- app/main.py | 123 +++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 113 insertions(+), 10 deletions(-) diff --git a/app/main.py b/app/main.py index 1c29874..59f4157 100644 --- a/app/main.py +++ b/app/main.py @@ -166,9 +166,40 @@ def get_db() -> sqlite3.Connection: cols = [r[1] for r in conn.execute("PRAGMA table_info(members)").fetchall()] if "slug" not in cols: conn.execute("ALTER TABLE members ADD COLUMN slug TEXT") + # Migration: add balance column (USD) — the member's current credit + # allowance. Defaults to MEMBER_BUDGET; can be topped up by payments. + # This is the flexible counter to the fixed $15/month allowance. + if "balance" not in cols: + conn.execute("ALTER TABLE members ADD COLUMN balance REAL") return conn +def get_member_balance(email: str) -> float: + """Return the member's current balance (USD), defaulting to MEMBER_BUDGET. + + The balance is the flexible allowance — it starts at MEMBER_BUDGET but can + be adjusted (topped up by payments) without changing code. + """ + conn = get_db() + row = conn.execute( + "SELECT balance FROM members WHERE email = ?", (email,) + ).fetchone() + conn.close() + if row and row[0] is not None: + return float(row[0]) + return MEMBER_BUDGET + + +def set_member_balance(email: str, balance: float) -> None: + """Set the member's balance (USD). Used by top-up / billing adjustments.""" + conn = get_db() + conn.execute( + "UPDATE members SET balance = ? WHERE email = ?", (balance, email) + ) + conn.commit() + conn.close() + + async def fetch_members() -> list[dict]: """Return the full active-member list: email, name, slug, role. @@ -535,9 +566,10 @@ async def cloudron_get_invite_link(user_id: str) -> str: async def provision_member(email: str, name: str, slug: str = "") -> str: """Provision a member end-to-end and send our own welcome email. - Creates the Cloudron user (members group), issues a LiteLLM key, stores the - mapping, and sends our branded welcome email containing the Cloudron - account-setup link (instead of Cloudron's default invite email). + Creates the Cloudron user (members group), issues a LiteLLM key under the + member's team (budget = the member's current balance, not a hardcoded $15), + stores the mapping, and sends our branded welcome email containing the + Cloudron account-setup link (instead of Cloudron's default invite email). Idempotent: if the member already exists, reuses the existing user/key. Returns the Cloudron user id. @@ -545,8 +577,9 @@ async def provision_member(email: str, name: str, slug: str = "") -> str: user_id = await cloudron_create_user(email, name) await cloudron_set_group(user_id) await cloudron_set_active(user_id, True) - team_id = await litellm_get_or_create_team(email, MEMBER_BUDGET) - key_token = await litellm_create_key(email, MEMBER_BUDGET, team_id) + balance = get_member_balance(email) + team_id = await litellm_get_or_create_team(email, balance) + key_token = await litellm_create_key(email, balance, team_id) store_member(email, key_token, user_id, slug) # Send our own welcome email with the account-setup link (no OAuth step). @@ -595,6 +628,36 @@ async def litellm_get_or_create_team(email: str, budget: float) -> str: return r.json()["team_id"] +async def litellm_update_team_budget(team_id: str, budget: float) -> None: + """Set a team's max_budget (the enforced spend cap). + + This is the primitive behind "payment → add to balance": a top-up adjusts + the team's max_budget, which every key under the team (chat + API) draws + from. Keys inherit the team budget (their own max_budget is ignored when a + team_id is set), so the team is the single source of truth for allowance. + """ + async with httpx.AsyncClient() as client: + r = await client.post( + f"{LITELLM_BASE}/team/update", + headers={"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}, + json={"team_id": team_id, "max_budget": budget}, + ) + r.raise_for_status() + + +async def sync_team_budget_from_balance(email: str) -> float: + """Apply the member's stored balance to their team's max_budget. + + Reads the balance from the members table and pushes it to LiteLLM. Returns + the applied budget. This keeps the portal DB and LiteLLM in sync whenever + a balance changes (top-up, reset, pricing-model change). + """ + balance = get_member_balance(email) + team_id = await litellm_get_or_create_team(email, balance) + await litellm_update_team_budget(team_id, balance) + return balance + + async def litellm_create_key(email: str, budget: float, team_id: str) -> str: """Create a fresh LiteLLM virtual key under the member's team. @@ -602,6 +665,12 @@ async def litellm_create_key(email: str, budget: float, team_id: str) -> str: create a new key rather than reusing — the sk- value is unrecoverable from a stored hash, so "reuse" would store a useless hash (a bug that previously broke chat for three members). + + NOTE: the key does NOT set its own max_budget — it inherits the team's + budget (LiteLLM ignores a key's max_budget when team_id is set anyway). + The team is the single source of truth for allowance, so a balance change + (top-up) updates the team once and every key follows. The `budget` arg is + retained only for the team-creation fallback in the migration path. """ alias = f"member:{email}" async with httpx.AsyncClient() as client: @@ -611,8 +680,6 @@ async def litellm_create_key(email: str, budget: float, team_id: str) -> str: json={ "key_alias": alias, "team_id": team_id, - "max_budget": budget, - "budget_duration": "30d", "models": MEMBER_MODELS, }, ) @@ -920,6 +987,41 @@ async def admin_provision(token: str, request: Request): return {"status": "provisioned", "email": email, "user_id": user_id} +@app.post("/admin/set-balance/{token}") +async def admin_set_balance(token: str, request: Request): + """Set (or add to) a member's credit balance and sync it to LiteLLM. + + The flexible alternative to a fixed $15/month allowance. A payment (from OC + or another system) is recorded by adjusting the balance here, which pushes + the new cap to the member's team (every chat/API key under it follows). + + Body: {"email": "...", "balance": 25.0} → set absolute balance + {"email": "...", "add": 10.0} → add to current balance + """ + if not WEBHOOK_TOKEN or not secrets.compare_digest(token, WEBHOOK_TOKEN): + raise HTTPException(401, "Invalid token") + + body = await request.json() + email = (body.get("email") or "").strip().lower() + if not email: + raise HTTPException(400, "Missing email") + + current = get_member_balance(email) + if "add" in body: + new_balance = current + float(body["add"]) + elif "balance" in body: + new_balance = float(body["balance"]) + else: + raise HTTPException(400, "Provide 'balance' or 'add'") + + if new_balance < 0: + raise HTTPException(400, "Balance cannot be negative") + + set_member_balance(email, new_balance) + applied = await sync_team_budget_from_balance(email) + return {"status": "updated", "email": email, "balance": new_balance, "team_budget": applied} + + @app.post("/admin/migrate-teams/{token}") async def admin_migrate_teams(token: str): """One-off migration: put every active member under a LiteLLM team + fresh @@ -942,9 +1044,10 @@ async def admin_migrate_teams(token: str): try: # Delete any existing key(s) for this alias (by hash or sk-). await litellm_delete_keys_by_alias(email) - # (Re)create the team and a fresh sk- key. - team_id = await litellm_get_or_create_team(email, MEMBER_BUDGET) - new_key = await litellm_create_key(email, MEMBER_BUDGET, team_id) + # (Re)create the team and a fresh sk- key, at the member's balance. + balance = get_member_balance(email) + team_id = await litellm_get_or_create_team(email, balance) + new_key = await litellm_create_key(email, balance, team_id) rekey_member(email, new_key) # preserve cloudron_user_id + slug results.append({"email": email, "status": "rekeyed"}) logger.info("Migrated %s to team %s", email, team_id)