144 lines
4.6 KiB
Python
144 lines
4.6 KiB
Python
"""AI 功能用量与配额（批次 B2，为批次 C billing 预留接口）。
|
|
|
|
设计：每次 AI 调用记一条 `AiUsage` 流水；配额从 `quota_for()` 取——
|
|
批次 C 的 billing 落地后，把 `apps/billing/quota.py` 放进来即自动接管，
|
|
本模块无需改动（导入失败则回落到 settings.AI_FREE_MONTHLY_QUOTA）。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import date
|
|
|
|
from django.conf import settings
|
|
from django.db import models
|
|
from django.db.models import Count
|
|
|
|
from apps.core.base_models import TenantScopedModel
|
|
|
|
# 用量类型
|
|
KIND_PARSE_ORDER = "parse_order"
|
|
KIND_ASK = "ask"
|
|
KIND_ADVICE = "advice"
|
|
|
|
|
|
try: # billing（批次 C1）已安装时复用其异常类，视图层只需 catch 一处
|
|
from apps.billing.quota import QuotaExceeded as _QuotaBase
|
|
except Exception: # pragma: no cover - billing 未安装时的兜底基类
|
|
_QuotaBase = Exception
|
|
|
|
|
|
class QuotaExceeded(_QuotaBase):
|
|
"""AI 用量超出套餐配额（调用方翻译为 403 + 升级引导）。"""
|
|
|
|
def __init__(self, kind: str, used: int, quota: int, **extra):
|
|
self.kind = kind
|
|
self.used = used
|
|
self.quota = quota
|
|
self.extra = extra
|
|
# 直接调 Exception.__init__：父类（billing）的签名不同，不适合链式调用
|
|
Exception.__init__(self, f"ai quota exceeded for {kind}: {used}/{quota}")
|
|
|
|
def as_dict(self) -> dict:
|
|
return {
|
|
"code": "quota_exceeded",
|
|
"kind": self.kind,
|
|
"detail": f"AI 功能本月已用 {self.used}/{self.quota} 次，升级套餐可解除限制",
|
|
"used": self.used,
|
|
"limit": self.quota,
|
|
"upgrade_url": "/#/pricing",
|
|
}
|
|
|
|
|
|
class AiUsage(TenantScopedModel):
|
|
"""AI 调用流水（一次调用一条）。"""
|
|
|
|
kind = models.CharField(max_length=32, help_text="parse_order/ask/advice")
|
|
period = models.CharField(max_length=7, help_text="账期 YYYY-MM（便于按月聚合）")
|
|
detail = models.JSONField(default=dict, blank=True, help_text="调用元数据（不含敏感原文）")
|
|
|
|
class Meta:
|
|
db_table = "ai_usage"
|
|
verbose_name_plural = "AI 调用流水"
|
|
indexes = [
|
|
models.Index(fields=["tenant", "kind", "period"]),
|
|
]
|
|
ordering = ["-created_at"]
|
|
|
|
def __str__(self):
|
|
return f"{self.kind}@{self.period}"
|
|
|
|
|
|
def current_period(today: date | None = None) -> str:
|
|
return (today or date.today()).strftime("%Y-%m")
|
|
|
|
|
|
def quota_for(tenant, kind: str) -> int:
|
|
"""该租户该 AI 功能的月度配额（0 = 不限）。
|
|
|
|
优先走批次 C 的 billing（若已安装），否则用全局免费额度。
|
|
"""
|
|
try:
|
|
from apps.billing.quota import quota_for as billing_quota # type: ignore
|
|
|
|
return int(billing_quota(tenant, f"ai_{kind}"))
|
|
except Exception:
|
|
return int(getattr(settings, "AI_FREE_MONTHLY_QUOTA", 10))
|
|
|
|
|
|
def monthly_count(tenant, kind: str, *, period: str | None = None) -> int:
|
|
"""本月该功能的已用次数。"""
|
|
return AiUsage.objects.filter(
|
|
tenant=tenant, kind=kind, period=period or current_period()
|
|
).count()
|
|
|
|
|
|
def check_quota(tenant, kind: str, *, period: str | None = None) -> dict:
|
|
"""检查并返回配额状态 {used, quota, remaining, allowed}。
|
|
|
|
超限抛 QuotaExceeded（不在此处吞掉，让视图层翻译 HTTP 语义）。
|
|
"""
|
|
period = period or current_period()
|
|
quota = quota_for(tenant, kind)
|
|
used = monthly_count(tenant, kind, period=period)
|
|
if quota > 0 and used >= quota:
|
|
raise QuotaExceeded(kind, used, quota)
|
|
return {
|
|
"used": used,
|
|
"quota": quota,
|
|
"remaining": (quota - used) if quota > 0 else None,
|
|
"allowed": True,
|
|
}
|
|
|
|
|
|
def record_usage(tenant, kind: str, *, detail: dict | None = None) -> AiUsage:
|
|
"""登记一次调用。"""
|
|
return AiUsage.objects.create(
|
|
tenant=tenant,
|
|
kind=kind,
|
|
period=current_period(),
|
|
detail=detail or {},
|
|
)
|
|
|
|
|
|
def usage_summary(tenant) -> dict:
|
|
"""该租户各 AI 功能本月用量总览（前端展示"免费版 10 次/月"进度）。"""
|
|
period = current_period()
|
|
rows = (
|
|
AiUsage.objects.filter(tenant=tenant, period=period)
|
|
.values("kind")
|
|
.annotate(n=Count("id"))
|
|
)
|
|
used = {r["kind"]: r["n"] for r in rows}
|
|
kinds = [KIND_PARSE_ORDER, KIND_ASK, KIND_ADVICE]
|
|
return {
|
|
"period": period,
|
|
"items": [
|
|
{
|
|
"kind": k,
|
|
"used": used.get(k, 0),
|
|
"quota": quota_for(tenant, k),
|
|
}
|
|
for k in kinds
|
|
],
|
|
}
|