broker: per-key spend (chat separate) in /broker/usage

Members need to see how much each API key has spent. LiteLLM's per-key
'spend' counter is cycle-only (resets at budget_reset_at), so derive it
from the LiteLLM_SpendLogs table instead: sum spend grouped by api_key,
matching rows to member_api_keys.hash by name and labelling unmatched
rows 'chat' (the portal-held chat key is never in member_api_keys).

broker_get_usage now returns key_spend: [{name, is_chat, spend, tokens,
requests}], API keys by spend desc with the chat row last.
This commit is contained in:
InferenceBot committed 2026-11-01 18:00:25 +00:00
1 parent b065b7301f
commit 6a7eaf0ff8
1 file changed
+81
+81
View File
@@ -3522,6 +3522,74 @@ async def broker_set_matrix_mxid(request: Request):
return {"status": "saved", "matrix_mxid": mxid, "invite_sent": invite_sent} return {"status": "saved", "matrix_mxid": mxid, "invite_sent": invite_sent}
def member_per_key_spend(email: str, team_id: str) -> list[dict]:
"""Per-key spend for a member, derived from LiteLLM spend LOGS.
The per-key `spend` counter LiteLLM exposes (`LiteLLM_VerificationToken.spend`,
`/spend/keys`) is CYCLE-ONLY: it resets at the team's `budget_reset_at`, so it
cannot answer "how much has this key spent?". We instead sum `spend` over the
raw rows in `LiteLLM_SpendLogs`, grouped by the key identifier in `api_key`
(the key's LiteLLM hash — the same value we store as `member_api_keys.hash`).
A row whose `api_key` matches one of the member's stored `member_api_keys`
rows is labelled with that key's name. Any row that matches no stored key
belongs to the portal-held CHAT key (alias `member:<email>`, minted at
provisioning, never in member_api_keys) and is grouped under the label
"chat" — so a member sees "chat: $X, <keyname>: $Y".
Returns [{"name", "is_chat", "spend", "tokens", "requests"}], API keys by
spend descending with the chat row last (it is secondary in the keys view).
Empty list when the DB is not configured — never raises for the caller.
"""
if not litellm_db_ready():
return []
conn = get_db()
try:
rows = conn.execute(
"SELECT name, hash FROM member_api_keys WHERE email = ?", (email,)
).fetchall()
finally:
conn.close()
by_hash = {r[1]: r[0] for r in rows if r[1]}
pg = _pg_connect()
try:
with pg.cursor() as cur:
cur.execute(
'SELECT api_key, COALESCE(SUM(spend), 0), '
'COALESCE(SUM(completion_tokens), 0), COUNT(*) '
'FROM "LiteLLM_SpendLogs" WHERE team_id = %s GROUP BY api_key',
(team_id,),
)
raw = cur.fetchall()
finally:
pg.close()
# Aggregate by a stable internal id. Rows matching no stored key are the chat
# key; the sentinel can't collide with a real name (key names are
# [A-Za-z0-9._-], so a member key literally named "chat" stays separate).
agg: dict = {}
for api_key, spend, tokens, reqs in raw:
name = by_hash.get(api_key)
cid = name if name else "__chat__"
a = agg.setdefault(
cid,
{"name": name or "chat", "is_chat": name is None,
"spend": 0.0, "tokens": 0, "requests": 0},
)
a["spend"] += float(spend or 0.0)
a["tokens"] += int(tokens or 0)
a["requests"] += int(reqs or 0)
out = list(agg.values())
# API keys by spend desc, then the chat key last.
out.sort(key=lambda x: (x["is_chat"], -x["spend"]))
for r in out:
r["spend"] = round(r["spend"], 2)
return out
async def broker_get_usage(email: str, days: int = 0) -> 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. """Return the member's spend vs. balance, plus per-model and token detail.
@@ -3569,6 +3637,14 @@ async def broker_get_usage(email: str, days: int = 0) -> dict:
"calls": m["requests"]} "calls": m["requests"]}
for m in usage["models"] for m in usage["models"]
] ]
# Per-key spend (chat key shown separately), from the spend-log rows —
# never LiteLLM's cycle-only per-key counter. Best-effort: a failure
# must not blank the whole usage view.
try:
key_spend = member_per_key_spend(email, team_id)
except Exception as e:
logger.warning("broker_get_usage: per-key spend failed: %s", e)
key_spend = []
return { return {
"email": email, "email": email,
"balance": balance, "balance": balance,
@@ -3581,6 +3657,7 @@ async def broker_get_usage(email: str, days: int = 0) -> dict:
"total_tokens": usage["totals"]["total_tokens"], "total_tokens": usage["totals"]["total_tokens"],
"reset_at": reset_at, "reset_at": reset_at,
"models": model_list, "models": model_list,
"key_spend": key_spend,
} }
except Exception as e: except Exception as e:
logger.warning("broker_get_usage: DB path failed (%s); falling back to HTTP", e) logger.warning("broker_get_usage: DB path failed (%s); falling back to HTTP", e)
@@ -3654,6 +3731,10 @@ async def broker_get_usage(email: str, days: int = 0) -> dict:
"total_tokens": total_tokens, "total_tokens": total_tokens,
"reset_at": reset_at, "reset_at": reset_at,
"models": model_list, "models": model_list,
# Per-key spend needs the Postgres spend-log table (LiteLLM's exposed
# per-key counter is cycle-only); this HTTP fallback runs only when the
# DB is unconfigured, so there is nothing accurate to report here.
"key_spend": [],
} }