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:
1 parent
b065b7301f
commit
6a7eaf0ff8
1 file changed
+81
+81
@@ -3522,6 +3522,74 @@ async def broker_set_matrix_mxid(request: Request):
|
||||
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:
|
||||
"""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"]}
|
||||
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 {
|
||||
"email": email,
|
||||
"balance": balance,
|
||||
@@ -3581,6 +3657,7 @@ async def broker_get_usage(email: str, days: int = 0) -> dict:
|
||||
"total_tokens": usage["totals"]["total_tokens"],
|
||||
"reset_at": reset_at,
|
||||
"models": model_list,
|
||||
"key_spend": key_spend,
|
||||
}
|
||||
except Exception as 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,
|
||||
"reset_at": reset_at,
|
||||
"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": [],
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user