316 lines
16 KiB
Python
316 lines
16 KiB
Python
|
"""用量计量、费率快照和用户配额。
|
||
|
|
|
||
|
|
UsageEvent 是追加式账本：同一个幂等键只允许落一条记录，历史金额不会随
|
||
|
|
Provider 当前配置变化。额度只读账本聚合，不把前端传入的数字当作已用量。
|
||
|
|
"""
|
||
|
|
from datetime import datetime
|
||
|
|
from math import ceil
|
||
|
|
from typing import Iterable
|
||
|
|
|
||
|
|
from fastapi import HTTPException
|
||
|
|
from sqlmodel import Session, select
|
||
|
|
from sqlalchemy import update
|
||
|
|
|
||
|
|
from .models import (ModelCatalog, ModelProvider, PointAccount, PointTransaction,
|
||
|
|
UsageEvent, UsageQuota, now)
|
||
|
|
|
||
|
|
POINTS_PER_TOKEN = 10
|
||
|
|
MIN_AI_POINTS = POINTS_PER_TOKEN
|
||
|
|
|
||
|
|
|
||
|
|
def _points_write_lock(session: Session) -> None:
|
||
|
|
"""Serialize the short ledger transaction, never the preceding model request.
|
||
|
|
|
||
|
|
SQLite's legacy driver does not begin a transaction for SELECT. Acquire the
|
||
|
|
write reservation before idempotency/account reads, and leave commit/rollback
|
||
|
|
to the caller so the balance and ledger remain one transaction.
|
||
|
|
"""
|
||
|
|
connection = session.connection()
|
||
|
|
if connection.dialect.name == 'sqlite' and not connection.connection.driver_connection.in_transaction:
|
||
|
|
connection.exec_driver_sql('BEGIN IMMEDIATE')
|
||
|
|
|
||
|
|
|
||
|
|
def _refresh_balance(session, account):
|
||
|
|
session.expire(account, ['balance_points', 'updated_at'])
|
||
|
|
|
||
|
|
|
||
|
|
def period_key(period: str = "monthly", value: 'str | None' = None) -> str:
|
||
|
|
"""返回用于账单筛选的日/月键。"""
|
||
|
|
current = value or now()
|
||
|
|
return current[:10] if period == "daily" else current[:7]
|
||
|
|
|
||
|
|
|
||
|
|
def estimate_tokens(value: object) -> int:
|
||
|
|
"""无 tokenizer 时的保守估算；事件会标记为 estimated。"""
|
||
|
|
text = str(value or "")
|
||
|
|
return max(1, ceil(len(text) / 2)) if text else 0
|
||
|
|
|
||
|
|
|
||
|
|
def point_state(session: Session, user_id: int, organization_id: int = 1) -> dict:
|
||
|
|
"""返回用户点数余额及可折算的 Token 数；没有账户时按 0 点数处理。"""
|
||
|
|
account = session.exec(select(PointAccount).where(
|
||
|
|
PointAccount.user_id == int(user_id or 0),
|
||
|
|
PointAccount.organization_id == int(organization_id or 1),
|
||
|
|
)).first()
|
||
|
|
points = max(0, int(account.balance_points or 0)) if account else 0
|
||
|
|
return {"points": points, "token_equivalent": points // POINTS_PER_TOKEN,
|
||
|
|
"remainder_points": points % POINTS_PER_TOKEN,
|
||
|
|
"points_per_token": POINTS_PER_TOKEN}
|
||
|
|
|
||
|
|
|
||
|
|
def _point_account(session: Session, user_id: int, organization_id: int) -> PointAccount:
|
||
|
|
account = session.exec(select(PointAccount).where(
|
||
|
|
PointAccount.user_id == int(user_id),
|
||
|
|
PointAccount.organization_id == int(organization_id),
|
||
|
|
)).first()
|
||
|
|
if account:
|
||
|
|
return account
|
||
|
|
account = PointAccount(user_id=int(user_id), organization_id=int(organization_id), balance_points=0)
|
||
|
|
session.add(account)
|
||
|
|
session.flush()
|
||
|
|
return account
|
||
|
|
|
||
|
|
|
||
|
|
def grant_points(session: Session, *, user_id: int, organization_id: int,
|
||
|
|
points: int, reason: str, operator_user_id: int = 0,
|
||
|
|
idempotency_key: str = "", kind: str = "credit",
|
||
|
|
source: str = "admin", metadata: dict | None = None) -> PointTransaction:
|
||
|
|
"""给用户增加点数，幂等键确保重复提交不会重复充值。"""
|
||
|
|
amount = int(points or 0)
|
||
|
|
if amount <= 0:
|
||
|
|
raise ValueError("增加点数必须大于 0")
|
||
|
|
key = str(idempotency_key or "").strip()
|
||
|
|
if not key:
|
||
|
|
raise ValueError("点数流水缺少幂等键")
|
||
|
|
_points_write_lock(session)
|
||
|
|
existing = session.exec(select(PointTransaction).where(
|
||
|
|
PointTransaction.idempotency_key == key)).first()
|
||
|
|
if existing:
|
||
|
|
return existing
|
||
|
|
account = _point_account(session, user_id, organization_id)
|
||
|
|
session.execute(update(PointAccount).where(PointAccount.id == account.id).values(
|
||
|
|
balance_points=PointAccount.balance_points + amount, updated_at=now()),
|
||
|
|
execution_options={'synchronize_session': False})
|
||
|
|
_refresh_balance(session, account)
|
||
|
|
transaction = PointTransaction(
|
||
|
|
user_id=int(user_id), organization_id=int(organization_id),
|
||
|
|
delta_points=amount, kind=str(kind or "credit"), reason=str(reason or ""),
|
||
|
|
source=str(source or "admin"), operator_user_id=int(operator_user_id or 0),
|
||
|
|
idempotency_key=key, transaction_metadata=metadata or {},
|
||
|
|
)
|
||
|
|
session.add(transaction)
|
||
|
|
session.flush()
|
||
|
|
return transaction
|
||
|
|
|
||
|
|
|
||
|
|
def consume_points(session: Session, *, user_id: int, organization_id: int,
|
||
|
|
points: int, reason: str, idempotency_key: str,
|
||
|
|
metadata: dict | None = None) -> bool:
|
||
|
|
"""扣减点数；余额不足时保持 0 并返回 False。"""
|
||
|
|
amount = max(0, int(points or 0))
|
||
|
|
if amount == 0:
|
||
|
|
return True
|
||
|
|
key = str(idempotency_key or "").strip()
|
||
|
|
if not key:
|
||
|
|
raise ValueError("点数扣减缺少幂等键")
|
||
|
|
_points_write_lock(session)
|
||
|
|
existing = session.exec(select(PointTransaction).where(
|
||
|
|
PointTransaction.idempotency_key == key)).first()
|
||
|
|
if existing:
|
||
|
|
return existing.delta_points < 0
|
||
|
|
account = session.exec(select(PointAccount).where(
|
||
|
|
PointAccount.user_id == int(user_id),
|
||
|
|
PointAccount.organization_id == int(organization_id),
|
||
|
|
)).first()
|
||
|
|
if not account:
|
||
|
|
return False
|
||
|
|
charged = session.execute(update(PointAccount).where(
|
||
|
|
PointAccount.id == account.id, PointAccount.balance_points >= amount).values(
|
||
|
|
balance_points=PointAccount.balance_points - amount, updated_at=now()),
|
||
|
|
execution_options={'synchronize_session': False})
|
||
|
|
_refresh_balance(session, account)
|
||
|
|
if charged.rowcount != 1:
|
||
|
|
return False
|
||
|
|
session.add(PointTransaction(
|
||
|
|
user_id=int(user_id), organization_id=int(organization_id),
|
||
|
|
delta_points=-amount, kind="debit", reason=str(reason or "模型调用"),
|
||
|
|
source="llm", idempotency_key=key, transaction_metadata=metadata or {},
|
||
|
|
))
|
||
|
|
session.flush()
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def require_ai_points(session: Session, *, user_id: int, organization_id: int = 1) -> dict:
|
||
|
|
"""模型请求前的最低点数门槛：少于 10 点数直接停止调用。"""
|
||
|
|
state = point_state(session, user_id, organization_id)
|
||
|
|
if int(user_id or 0) and state["points"] < MIN_AI_POINTS:
|
||
|
|
raise ValueError(f"点数余额不足，至少需要 {MIN_AI_POINTS} 点数（1 Token）才能使用 AI")
|
||
|
|
return state
|
||
|
|
|
||
|
|
|
||
|
|
def pricing_for(session: Session, provider_id: int = 0, model: str = "") -> dict:
|
||
|
|
catalog = session.exec(select(ModelCatalog).where(
|
||
|
|
ModelCatalog.provider_id == int(provider_id or 0),
|
||
|
|
ModelCatalog.model_name == str(model or ""),
|
||
|
|
ModelCatalog.enabled == True, # noqa: E712
|
||
|
|
)).first()
|
||
|
|
if catalog:
|
||
|
|
return {"input_rate": float(catalog.input_rate or 0),
|
||
|
|
"output_rate": float(catalog.output_rate or 0),
|
||
|
|
"version": catalog.rate_version or "v1"}
|
||
|
|
provider = session.get(ModelProvider, int(provider_id or 0)) if provider_id else None
|
||
|
|
if provider and model in (provider.models or []):
|
||
|
|
return {"input_rate": 0.0, "output_rate": 0.0, "version": "v1"}
|
||
|
|
return {"input_rate": 0.0, "output_rate": 0.0, "version": "v1"}
|
||
|
|
|
||
|
|
|
||
|
|
def amount_minor(input_quantity: int, output_quantity: int,
|
||
|
|
input_rate: float, output_rate: float) -> int:
|
||
|
|
amount = (max(0, int(input_quantity or 0)) / 1_000_000 * max(0.0, float(input_rate or 0)))
|
||
|
|
amount += (max(0, int(output_quantity or 0)) / 1_000_000 * max(0.0, float(output_rate or 0)))
|
||
|
|
return int(round(amount * 100))
|
||
|
|
|
||
|
|
|
||
|
|
def record_usage_event(session: Session, *, idempotency_key: str,
|
||
|
|
user_id: int = 0, organization_id: int = 1,
|
||
|
|
campaign_id: int = 0, run_id: int = 0, segment_id: int = 0,
|
||
|
|
trace_id: str = "", llm_call_id: int | None = None,
|
||
|
|
agent_id: int = 0, provider_id: int = 0,
|
||
|
|
provider: str = "", model: str = "",
|
||
|
|
service_type: str = "llm", unit: str = "token",
|
||
|
|
quantity: float = 0.0, input_quantity: int = 0,
|
||
|
|
output_quantity: int = 0, input_rate: float = 0.0,
|
||
|
|
output_rate: float = 0.0, pricing_version: str = "v1",
|
||
|
|
status: str = "measured", metadata: dict | None = None,
|
||
|
|
amount_minor_override: int | None = None) -> UsageEvent:
|
||
|
|
"""追加一条事件；重复幂等键直接复用原事件，不更新历史账本。"""
|
||
|
|
key = str(idempotency_key or "").strip()
|
||
|
|
if not key:
|
||
|
|
raise ValueError("UsageEvent 缺少幂等键")
|
||
|
|
existing = session.exec(select(UsageEvent).where(UsageEvent.idempotency_key == key)).first()
|
||
|
|
if existing:
|
||
|
|
return existing
|
||
|
|
event = UsageEvent(
|
||
|
|
idempotency_key=key,
|
||
|
|
user_id=int(user_id or 0), organization_id=int(organization_id or 1),
|
||
|
|
campaign_id=int(campaign_id or 0), run_id=int(run_id or 0),
|
||
|
|
segment_id=int(segment_id or 0), trace_id=str(trace_id or ""),
|
||
|
|
llm_call_id=llm_call_id, agent_id=int(agent_id or 0),
|
||
|
|
provider_id=int(provider_id or 0), provider=str(provider or ""),
|
||
|
|
model=str(model or ""), service_type=str(service_type or "llm"),
|
||
|
|
unit=str(unit or "token"), quantity=float(quantity or 0),
|
||
|
|
input_quantity=max(0, int(input_quantity or 0)),
|
||
|
|
output_quantity=max(0, int(output_quantity or 0)),
|
||
|
|
input_rate=float(input_rate or 0), output_rate=float(output_rate or 0),
|
||
|
|
amount_minor=(int(amount_minor_override) if amount_minor_override is not None else
|
||
|
|
amount_minor(input_quantity, output_quantity, input_rate, output_rate)),
|
||
|
|
pricing_version=str(pricing_version or "v1"), status=str(status or "measured"),
|
||
|
|
usage_metadata=metadata or {},
|
||
|
|
)
|
||
|
|
session.add(event)
|
||
|
|
session.flush()
|
||
|
|
return event
|
||
|
|
|
||
|
|
|
||
|
|
def events_for_period(session: Session, *, user_id: int | None = None,
|
||
|
|
organization_id: int | None = None, period: str = "monthly",
|
||
|
|
service_type: str = "", status: str = "") -> list[UsageEvent]:
|
||
|
|
query = select(UsageEvent).order_by(UsageEvent.id.desc())
|
||
|
|
if user_id is not None:
|
||
|
|
query = query.where(UsageEvent.user_id == int(user_id))
|
||
|
|
if organization_id is not None:
|
||
|
|
query = query.where(UsageEvent.organization_id == int(organization_id))
|
||
|
|
if service_type:
|
||
|
|
query = query.where(UsageEvent.service_type == service_type)
|
||
|
|
if status:
|
||
|
|
query = query.where(UsageEvent.status == status)
|
||
|
|
key = period_key(period)
|
||
|
|
return [event for event in session.exec(query).all()
|
||
|
|
if (event.created_at or "").startswith(key)]
|
||
|
|
|
||
|
|
|
||
|
|
def usage_summary(events: Iterable[UsageEvent]) -> dict:
|
||
|
|
rows = list(events)
|
||
|
|
valid = [event for event in rows if event.status != "void"]
|
||
|
|
input_tokens = sum(int(event.input_quantity or 0) for event in valid)
|
||
|
|
output_tokens = sum(int(event.output_quantity or 0) for event in valid)
|
||
|
|
amount = sum(int(event.amount_minor or 0) for event in valid)
|
||
|
|
|
||
|
|
def group(field: str) -> dict:
|
||
|
|
grouped: dict[str, dict] = {}
|
||
|
|
for event in valid:
|
||
|
|
key = str(getattr(event, field) or "unknown")
|
||
|
|
item = grouped.setdefault(key, {"events": 0, "quantity": 0.0,
|
||
|
|
"input_tokens": 0, "output_tokens": 0,
|
||
|
|
"amount_minor": 0})
|
||
|
|
item["events"] += 1
|
||
|
|
item["quantity"] += float(event.quantity or 0)
|
||
|
|
item["input_tokens"] += int(event.input_quantity or 0)
|
||
|
|
item["output_tokens"] += int(event.output_quantity or 0)
|
||
|
|
item["amount_minor"] += int(event.amount_minor or 0)
|
||
|
|
return grouped
|
||
|
|
|
||
|
|
return {"events": len(valid), "input_tokens": input_tokens,
|
||
|
|
"output_tokens": output_tokens, "tokens": input_tokens + output_tokens,
|
||
|
|
"amount_minor": amount, "amount": amount / 100,
|
||
|
|
"by_model": group("model"), "by_service": group("service_type"),
|
||
|
|
"estimated_events": sum(event.status == "estimated" for event in valid),
|
||
|
|
"measured_events": sum(event.status == "measured" for event in valid),
|
||
|
|
"settled_events": sum(event.status == "settled" for event in valid)}
|
||
|
|
|
||
|
|
|
||
|
|
def event_payload(event: UsageEvent, user=None) -> dict:
|
||
|
|
return {"id": event.id, "user_id": event.user_id, "display_name": (getattr(user, "display_name", "") or getattr(user, "username", "") or ("系统" if not event.user_id else "")), "username": getattr(user, "username", "") or ("系统" if not event.user_id else ""), "organization_id": event.organization_id,
|
||
|
|
"campaign_id": event.campaign_id, "run_id": event.run_id, "segment_id": event.segment_id,
|
||
|
|
"trace_id": event.trace_id, "llm_call_id": event.llm_call_id, "agent_id": event.agent_id,
|
||
|
|
"provider_id": event.provider_id, "provider": event.provider, "model": event.model,
|
||
|
|
"service_type": event.service_type, "unit": event.unit, "quantity": event.quantity,
|
||
|
|
"input_quantity": event.input_quantity, "output_quantity": event.output_quantity,
|
||
|
|
"input_rate": event.input_rate, "output_rate": event.output_rate,
|
||
|
|
"amount_minor": event.amount_minor, "amount": event.amount_minor / 100,
|
||
|
|
"currency": event.currency, "pricing_version": event.pricing_version,
|
||
|
|
"status": event.status, "idempotency_key": event.idempotency_key,
|
||
|
|
"metadata": event.usage_metadata, "created_at": event.created_at}
|
||
|
|
|
||
|
|
|
||
|
|
def quota_for(session: Session, user_id: int, organization_id: int = 1,
|
||
|
|
period: str = "monthly") -> UsageQuota | None:
|
||
|
|
return session.exec(select(UsageQuota).where(
|
||
|
|
UsageQuota.user_id == int(user_id), UsageQuota.organization_id == int(organization_id),
|
||
|
|
UsageQuota.period == period, UsageQuota.enabled == True, # noqa: E712
|
||
|
|
)).first()
|
||
|
|
|
||
|
|
|
||
|
|
def quota_state(session: Session, user_id: int, organization_id: int = 1,
|
||
|
|
period: str = "monthly") -> dict:
|
||
|
|
quota = quota_for(session, user_id, organization_id, period)
|
||
|
|
events = events_for_period(session, user_id=user_id, organization_id=organization_id, period=period)
|
||
|
|
summary = usage_summary(events)
|
||
|
|
if not quota:
|
||
|
|
return {"configured": False, "allowed": True, "warning": False,
|
||
|
|
"quota": None, "used": summary}
|
||
|
|
token_limit = max(0, int(quota.token_limit or 0))
|
||
|
|
amount_limit = max(0, int(quota.amount_limit_minor or 0))
|
||
|
|
token_ratio = summary["tokens"] / token_limit if token_limit else 0
|
||
|
|
amount_ratio = summary["amount_minor"] / amount_limit if amount_limit else 0
|
||
|
|
ratio = max(token_ratio, amount_ratio)
|
||
|
|
exceeded = (token_limit and summary["tokens"] >= token_limit) or (amount_limit and summary["amount_minor"] >= amount_limit)
|
||
|
|
blocked = bool(exceeded and quota.overage_policy in {"soft_limit", "hard_limit"})
|
||
|
|
return {"configured": True, "allowed": not blocked, "warning": ratio >= float(quota.warning_ratio or .8),
|
||
|
|
"quota": {"id": quota.id, "period": quota.period, "token_limit": token_limit,
|
||
|
|
"amount_limit_minor": amount_limit, "storage_limit_bytes": quota.storage_limit_bytes,
|
||
|
|
"concurrent_limit": quota.concurrent_limit, "warning_ratio": quota.warning_ratio,
|
||
|
|
"overage_policy": quota.overage_policy, "enabled": quota.enabled},
|
||
|
|
"used": summary, "ratio": ratio, "exceeded": bool(exceeded)}
|
||
|
|
|
||
|
|
|
||
|
|
def check_quota(session: Session, user_id: int, organization_id: int = 1,
|
||
|
|
period: str = "monthly", estimated_tokens: int = 0) -> dict:
|
||
|
|
state = quota_state(session, user_id, organization_id, period)
|
||
|
|
quota = state.get("quota") or {}
|
||
|
|
limit = int(quota.get("token_limit") or 0)
|
||
|
|
projected = int(state["used"]["tokens"] or 0) + max(0, int(estimated_tokens or 0))
|
||
|
|
if limit and projected >= limit and quota.get("overage_policy") in {"soft_limit", "hard_limit"}:
|
||
|
|
raise HTTPException(429, "本周期用量已达到管理员设置的 Token 配额")
|
||
|
|
return state
|