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}
|
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": [],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user