119 lines
4.0 KiB
Python
119 lines
4.0 KiB
Python
"""
|
|
Platform Billing Webhook
|
|
接收 LiteLLM spend 回调，写入 token_usage_logs 并扣减 user_balances
|
|
"""
|
|
import logging
|
|
from datetime import datetime
|
|
from fastapi import APIRouter, Request, HTTPException
|
|
from pydantic import BaseModel
|
|
from typing import Optional
|
|
from sqlalchemy import text
|
|
from app.database import SessionLocal
|
|
|
|
logger = logging.getLogger("aiflow.billing")
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
class SpendEvent(BaseModel):
|
|
"""LiteLLM spend log event"""
|
|
request_id: Optional[str] = None
|
|
api_key: Optional[str] = None
|
|
model: Optional[str] = None
|
|
total_tokens: int = 0
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
spend: float = 0.0
|
|
user: Optional[str] = None # platform user_id passed via metadata
|
|
metadata: Optional[dict] = None
|
|
startTime: Optional[str] = None
|
|
endTime: Optional[str] = None
|
|
|
|
|
|
@router.post("/webhook")
|
|
async def billing_webhook(event: SpendEvent, request: Request):
|
|
"""
|
|
接收 LiteLLM spend webhook 回调
|
|
记录 token 使用并扣减用户余额
|
|
"""
|
|
# 验证来源
|
|
api_key = event.api_key or ""
|
|
metadata = event.metadata or {}
|
|
|
|
# 只处理 mumuainovel 的用量
|
|
app_name = metadata.get("app", "")
|
|
if app_name != "mumuainovel" and "mumu" not in api_key:
|
|
# 也检查是否有自定义的 X-Source header
|
|
source = request.headers.get("X-Source", "")
|
|
if source != "mumuainovel":
|
|
return {"status": "skipped", "reason": "not_mumuainovel"}
|
|
|
|
user_id = event.user or metadata.get("user_id", "mumu-user")
|
|
model_name = event.model or "unknown"
|
|
spend = max(event.spend, 0)
|
|
total_tokens = event.total_tokens
|
|
prompt_tokens = event.prompt_tokens
|
|
completion_tokens = event.completion_tokens
|
|
|
|
if total_tokens == 0 and spend == 0:
|
|
return {"status": "skipped", "reason": "zero_usage"}
|
|
|
|
db = SessionLocal()
|
|
try:
|
|
# 写入 token_usage_logs
|
|
db.execute(text("""
|
|
INSERT INTO token_usage_logs (id, user_id, model_name, app_id, app_name,
|
|
prompt_tokens, completion_tokens, total_tokens, cost, created_at)
|
|
VALUES (gen_random_uuid(), :uid, :model, :app_id, :app_name,
|
|
:pt, :ct, :tt, :cost, :ts)
|
|
"""), {
|
|
"uid": user_id,
|
|
"model": model_name,
|
|
"app_id": "mumuainovel-001",
|
|
"app_name": "MuMuAINovel",
|
|
"pt": prompt_tokens,
|
|
"ct": completion_tokens,
|
|
"tt": total_tokens,
|
|
"cost": spend,
|
|
"ts": datetime.utcnow(),
|
|
})
|
|
|
|
# 扣减余额
|
|
result = db.execute(text("""
|
|
UPDATE user_balances
|
|
SET balance = balance - :cost,
|
|
total_used = total_used + :cost,
|
|
updated_at = NOW()
|
|
WHERE user_id = :uid AND balance >= :cost
|
|
RETURNING id
|
|
"""), {"cost": spend, "uid": user_id})
|
|
|
|
if result.rowcount == 0:
|
|
# 余额不足或无记录，创建新余额记录
|
|
# 检查是否存在
|
|
exists = db.execute(text(
|
|
"SELECT id FROM user_balances WHERE user_id = :uid"
|
|
), {"uid": user_id}).fetchone()
|
|
if not exists:
|
|
db.execute(text("""
|
|
INSERT INTO user_balances (id, user_id, balance, total_used, created_at, updated_at)
|
|
VALUES (gen_random_uuid(), :uid, :balance, :cost, NOW(), NOW())
|
|
"""), {"uid": user_id, "balance": 10.0 - spend, "cost": spend})
|
|
|
|
db.commit()
|
|
logger.info(f"Billing: user={user_id} model={model_name} tokens={total_tokens} spend=${spend:.6f}")
|
|
|
|
except Exception as e:
|
|
db.rollback()
|
|
logger.error(f"Billing webhook error: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
finally:
|
|
db.close()
|
|
|
|
return {"status": "ok", "user_id": user_id, "spend": spend}
|
|
|
|
|
|
@router.get("/webhook/health")
|
|
async def billing_webhook_health():
|
|
return {"status": "ok", "service": "billing-webhook"}
|