forked from xiaohei/taiji-AI-PAD
674 lines
19 KiB
Python
674 lines
19 KiB
Python
"""
|
||
预付费配额管理模块
|
||
配额预警与限制
|
||
|
||
注意：计费数据已迁移到新表：
|
||
- AgentBillingRecord: Agent 运行时计费
|
||
- ModelBillingRecord: 模型调用计费（LiteLLM）
|
||
旧的 BillingRecord 表已废弃，不再使用。
|
||
"""
|
||
|
||
from datetime import datetime, timedelta
|
||
from decimal import Decimal
|
||
from typing import Dict, List, Optional, Tuple
|
||
from sqlalchemy import select, func, and_, update
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from models import User, Channel, QuotaAlert, ResourceAllocation, Balance, AgentBillingRecord, ModelBillingRecord
|
||
|
||
|
||
# 配额预警阈值配置
|
||
QUOTA_THRESHOLDS = {
|
||
"balance_warning": 20, # 余额低于20%时预警
|
||
"balance_critical": 10, # 余额低于10%时严重预警
|
||
"quota_warning": 80, # 配额使用超过80%时预警
|
||
"quota_critical": 95, # 配额使用超过95%时严重预警
|
||
}
|
||
|
||
|
||
async def check_user_balance_quota(
|
||
user_id: str,
|
||
db: AsyncSession
|
||
) -> Tuple[bool, Optional[str], Dict]:
|
||
"""
|
||
检查用户余额配额
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
(是否有足够配额, 预警类型, 详情)
|
||
"""
|
||
result = await db.execute(
|
||
select(User).where(User.id == user_id)
|
||
)
|
||
user = result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
return False, "user_not_found", {"message": "用户不存在"}
|
||
|
||
# 从 Balance 表获取余额
|
||
balance_result = await db.execute(
|
||
select(Balance).where(Balance.user_id == user_id)
|
||
)
|
||
balance_obj = balance_result.scalar_one_or_none()
|
||
balance = Decimal(str(balance_obj.eu_balance)) if balance_obj else Decimal(0)
|
||
credit_limit = Decimal(str(user.credit_limit))
|
||
available = balance + credit_limit
|
||
|
||
# 获取用户平均日消费（从 AgentBillingRecord 和 ModelBillingRecord 合并统计）
|
||
thirty_days_ago = datetime.utcnow() - timedelta(days=30)
|
||
|
||
# Agent 计费
|
||
agent_cost_result = await db.execute(
|
||
select(func.sum(AgentBillingRecord.cost))
|
||
.where(
|
||
and_(
|
||
AgentBillingRecord.user_id == user_id,
|
||
AgentBillingRecord.start_time >= thirty_days_ago,
|
||
)
|
||
)
|
||
)
|
||
agent_cost = agent_cost_result.scalar() or 0
|
||
|
||
# 模型调用计费
|
||
model_cost_result = await db.execute(
|
||
select(func.sum(ModelBillingRecord.total_cost))
|
||
.where(
|
||
and_(
|
||
ModelBillingRecord.tenant_id == user_id,
|
||
ModelBillingRecord.created_at >= thirty_days_ago,
|
||
)
|
||
)
|
||
)
|
||
model_cost = model_cost_result.scalar() or 0
|
||
|
||
total_cost = Decimal(str(agent_cost)) + Decimal(str(model_cost))
|
||
daily_avg = total_cost / 30
|
||
|
||
# 预估可用天数
|
||
if daily_avg > 0:
|
||
estimated_days = float(available / daily_avg)
|
||
else:
|
||
estimated_days = 999
|
||
|
||
details = {
|
||
"balance": float(balance),
|
||
"creditLimit": float(credit_limit),
|
||
"available": float(available),
|
||
"dailyAvgCost": float(daily_avg),
|
||
"estimatedDays": round(estimated_days, 1),
|
||
}
|
||
|
||
# 判断是否需要预警
|
||
if available <= 0:
|
||
return False, "balance_exhausted", details
|
||
elif estimated_days <= 3:
|
||
return True, "balance_critical", details
|
||
elif estimated_days <= 7:
|
||
return True, "balance_warning", details
|
||
|
||
return True, None, details
|
||
|
||
|
||
async def check_channel_quota(
|
||
channel_id: str,
|
||
db: AsyncSession
|
||
) -> Tuple[bool, Optional[str], Dict]:
|
||
"""
|
||
检查渠道配额
|
||
|
||
Args:
|
||
channel_id: 渠道ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
(是否有足够配额, 预警类型, 详情)
|
||
"""
|
||
result = await db.execute(
|
||
select(Channel).where(Channel.id == channel_id)
|
||
)
|
||
channel = result.scalar_one_or_none()
|
||
|
||
if not channel:
|
||
return False, "channel_not_found", {"message": "渠道不存在"}
|
||
|
||
channel_credit = Decimal(str(channel.channel_credit))
|
||
|
||
# 获取渠道下所有租户的总消费（从 AgentBillingRecord 和 ModelBillingRecord 合并统计）
|
||
thirty_days_ago = datetime.utcnow() - timedelta(days=30)
|
||
|
||
# Agent 计费
|
||
agent_cost_result = await db.execute(
|
||
select(func.sum(AgentBillingRecord.cost))
|
||
.where(
|
||
and_(
|
||
AgentBillingRecord.channel_id == channel_id,
|
||
AgentBillingRecord.start_time >= thirty_days_ago,
|
||
)
|
||
)
|
||
)
|
||
agent_cost = agent_cost_result.scalar() or 0
|
||
|
||
# 模型调用计费
|
||
model_cost_result = await db.execute(
|
||
select(func.sum(ModelBillingRecord.total_cost))
|
||
.where(
|
||
and_(
|
||
ModelBillingRecord.channel_id == channel_id,
|
||
ModelBillingRecord.created_at >= thirty_days_ago,
|
||
)
|
||
)
|
||
)
|
||
model_cost = model_cost_result.scalar() or 0
|
||
|
||
total_cost = float(agent_cost) + float(model_cost)
|
||
|
||
details = {
|
||
"channelCredit": float(channel_credit),
|
||
"monthlyUsage": float(total_cost),
|
||
"usagePercent": round(float(total_cost) / float(channel_credit) * 100, 2) if channel_credit > 0 else 0,
|
||
}
|
||
|
||
if channel_credit > 0:
|
||
usage_percent = float(total_cost) / float(channel_credit) * 100
|
||
if usage_percent >= QUOTA_THRESHOLDS["quota_critical"]:
|
||
return True, "quota_critical", details
|
||
elif usage_percent >= QUOTA_THRESHOLDS["quota_warning"]:
|
||
return True, "quota_warning", details
|
||
|
||
return True, None, details
|
||
|
||
|
||
async def create_quota_alert(
|
||
user_id: str,
|
||
channel_id: Optional[str],
|
||
alert_type: str,
|
||
threshold_percent: int,
|
||
current_value: Decimal,
|
||
threshold_value: Decimal,
|
||
db: AsyncSession
|
||
) -> QuotaAlert:
|
||
"""
|
||
创建配额预警记录
|
||
|
||
使用行锁保护并发创建/更新预警记录。
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
channel_id: 渠道ID
|
||
alert_type: 预警类型
|
||
threshold_percent: 阈值百分比
|
||
current_value: 当前值
|
||
threshold_value: 阈值
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
QuotaAlert记录
|
||
"""
|
||
# 检查是否已有相同的活跃预警（使用行锁防止并发创建重复预警）
|
||
result = await db.execute(
|
||
select(QuotaAlert)
|
||
.where(
|
||
and_(
|
||
QuotaAlert.user_id == user_id,
|
||
QuotaAlert.alert_type == alert_type,
|
||
QuotaAlert.status == "active",
|
||
)
|
||
)
|
||
.with_for_update() # 行锁
|
||
)
|
||
existing = result.scalar_one_or_none()
|
||
|
||
if existing:
|
||
# 更新现有预警
|
||
existing.current_value = current_value
|
||
existing.threshold_value = threshold_value
|
||
existing.updated_at = datetime.utcnow()
|
||
alert = existing
|
||
else:
|
||
# 创建新预警
|
||
alert = QuotaAlert(
|
||
user_id=user_id,
|
||
channel_id=channel_id,
|
||
alert_type=alert_type,
|
||
threshold_percent=threshold_percent,
|
||
current_value=current_value,
|
||
threshold_value=threshold_value,
|
||
status="active",
|
||
)
|
||
db.add(alert)
|
||
|
||
await db.commit()
|
||
await db.refresh(alert)
|
||
|
||
return alert
|
||
|
||
|
||
async def get_active_alerts(
|
||
user_id: Optional[str],
|
||
channel_id: Optional[str],
|
||
db: AsyncSession
|
||
) -> List[Dict]:
|
||
"""
|
||
获取活跃的配额预警
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
channel_id: 渠道ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
预警列表
|
||
"""
|
||
query = select(QuotaAlert).where(QuotaAlert.status == "active")
|
||
|
||
if user_id:
|
||
query = query.where(QuotaAlert.user_id == user_id)
|
||
if channel_id:
|
||
query = query.where(QuotaAlert.channel_id == channel_id)
|
||
|
||
result = await db.execute(query.order_by(QuotaAlert.created_at.desc()))
|
||
alerts = result.scalars().all()
|
||
|
||
return [
|
||
{
|
||
"id": str(alert.id),
|
||
"alertType": alert.alert_type,
|
||
"thresholdPercent": alert.threshold_percent,
|
||
"currentValue": float(alert.current_value) if alert.current_value else None,
|
||
"thresholdValue": float(alert.threshold_value) if alert.threshold_value else None,
|
||
"status": alert.status,
|
||
"createdAt": alert.created_at.isoformat(),
|
||
}
|
||
for alert in alerts
|
||
]
|
||
|
||
|
||
async def acknowledge_alert(
|
||
alert_id: str,
|
||
db: AsyncSession
|
||
) -> bool:
|
||
"""
|
||
确认预警
|
||
|
||
Args:
|
||
alert_id: 预警ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
result = await db.execute(
|
||
select(QuotaAlert).where(QuotaAlert.id == alert_id)
|
||
)
|
||
alert = result.scalar_one_or_none()
|
||
|
||
if not alert:
|
||
return False
|
||
|
||
alert.status = "acknowledged"
|
||
alert.acknowledged_at = datetime.utcnow()
|
||
|
||
await db.commit()
|
||
return True
|
||
|
||
|
||
async def resolve_alert(
|
||
alert_id: str,
|
||
db: AsyncSession
|
||
) -> bool:
|
||
"""
|
||
解决预警
|
||
|
||
Args:
|
||
alert_id: 预警ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
result = await db.execute(
|
||
select(QuotaAlert).where(QuotaAlert.id == alert_id)
|
||
)
|
||
alert = result.scalar_one_or_none()
|
||
|
||
if not alert:
|
||
return False
|
||
|
||
alert.status = "resolved"
|
||
alert.resolved_at = datetime.utcnow()
|
||
|
||
await db.commit()
|
||
return True
|
||
|
||
|
||
async def check_rate_limit(
|
||
user_id: str,
|
||
resource_type: str,
|
||
db: AsyncSession
|
||
) -> Tuple[bool, int, int]:
|
||
"""
|
||
检查速率限制
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
resource_type: 资源类型 (api, model)
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
(是否允许, 当前使用量, 限制量)
|
||
"""
|
||
# 获取最近1分钟的调用次数（从 AgentBillingRecord 和 ModelBillingRecord 合并统计）
|
||
one_minute_ago = datetime.utcnow() - timedelta(minutes=1)
|
||
|
||
# Agent 调用次数
|
||
agent_result = await db.execute(
|
||
select(func.count(AgentBillingRecord.id))
|
||
.where(
|
||
and_(
|
||
AgentBillingRecord.user_id == user_id,
|
||
AgentBillingRecord.start_time >= one_minute_ago,
|
||
)
|
||
)
|
||
)
|
||
agent_rpm = agent_result.scalar() or 0
|
||
|
||
# 模型调用次数
|
||
model_result = await db.execute(
|
||
select(func.count(ModelBillingRecord.id))
|
||
.where(
|
||
and_(
|
||
ModelBillingRecord.tenant_id == user_id,
|
||
ModelBillingRecord.created_at >= one_minute_ago,
|
||
)
|
||
)
|
||
)
|
||
model_rpm = model_result.scalar() or 0
|
||
|
||
current_rpm = agent_rpm + model_rpm
|
||
|
||
# 获取用户的RPM限制（从资源分配表）
|
||
limit_result = await db.execute(
|
||
select(ResourceAllocation.rpm)
|
||
.where(
|
||
and_(
|
||
ResourceAllocation.target_id == user_id,
|
||
ResourceAllocation.target_type == "tenant",
|
||
ResourceAllocation.resource_type == resource_type,
|
||
)
|
||
)
|
||
)
|
||
limit_row = limit_result.first()
|
||
rpm_limit = limit_row[0] if limit_row and limit_row[0] else 60 # 默认60 RPM
|
||
|
||
return current_rpm < rpm_limit, current_rpm, rpm_limit
|
||
|
||
|
||
async def get_quota_summary(
|
||
user_id: str,
|
||
db: AsyncSession
|
||
) -> Dict:
|
||
"""
|
||
获取配额汇总信息
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
配额汇总
|
||
"""
|
||
# 获取余额信息
|
||
has_quota, alert_type, balance_info = await check_user_balance_quota(user_id, db)
|
||
|
||
# 获取速率限制
|
||
rate_allowed, current_rpm, rpm_limit = await check_rate_limit(user_id, "api", db)
|
||
|
||
# 获取活跃预警数
|
||
alerts = await get_active_alerts(user_id, None, db)
|
||
|
||
return {
|
||
"hasQuota": has_quota,
|
||
"alertType": alert_type,
|
||
"balance": balance_info,
|
||
"rateLimit": {
|
||
"currentRpm": current_rpm,
|
||
"rpmLimit": rpm_limit,
|
||
"allowed": rate_allowed,
|
||
},
|
||
"activeAlerts": len(alerts),
|
||
"alerts": alerts,
|
||
}
|
||
|
||
|
||
async def get_user_quota_summary(
|
||
user_id: str,
|
||
db: AsyncSession
|
||
) -> Dict:
|
||
"""
|
||
获取用户配额综合信息（用于API返回）
|
||
|
||
Returns:
|
||
{
|
||
"hasQuota": bool,
|
||
"alertType": str | None,
|
||
"balance": {...},
|
||
"rateLimit": {...},
|
||
"activeAlerts": int,
|
||
"alerts": [...]
|
||
}
|
||
"""
|
||
from models import User, QuotaAlert, ResourceUsage
|
||
|
||
# 1. 获取用户信息
|
||
result = await db.execute(
|
||
select(User).where(User.id == user_id)
|
||
)
|
||
user = result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
raise ValueError(f"用户不存在: {user_id}")
|
||
|
||
# 2. 余额信息（从 Balance 表获取）
|
||
balance_result = await db.execute(
|
||
select(Balance).where(Balance.user_id == user_id)
|
||
)
|
||
balance_obj = balance_result.scalar_one_or_none()
|
||
balance = Decimal(str(balance_obj.eu_balance)) if balance_obj else Decimal(0)
|
||
credit_limit = Decimal(str(user.credit_limit or 0))
|
||
available = balance + credit_limit
|
||
|
||
# 计算近30天平均日消费（从 AgentBillingRecord 和 ModelBillingRecord 合并统计）
|
||
thirty_days_ago = datetime.utcnow() - timedelta(days=30)
|
||
|
||
# Agent 计费
|
||
agent_cost_result = await db.execute(
|
||
select(func.sum(AgentBillingRecord.cost))
|
||
.where(
|
||
and_(
|
||
AgentBillingRecord.user_id == user_id,
|
||
AgentBillingRecord.start_time >= thirty_days_ago,
|
||
)
|
||
)
|
||
)
|
||
agent_cost = agent_cost_result.scalar() or 0
|
||
|
||
# 模型调用计费
|
||
model_cost_result = await db.execute(
|
||
select(func.sum(ModelBillingRecord.total_cost))
|
||
.where(
|
||
and_(
|
||
ModelBillingRecord.tenant_id == user_id,
|
||
ModelBillingRecord.created_at >= thirty_days_ago,
|
||
)
|
||
)
|
||
)
|
||
model_cost = model_cost_result.scalar() or 0
|
||
|
||
total_cost = Decimal(str(agent_cost)) + Decimal(str(model_cost))
|
||
daily_avg = total_cost / 30 if total_cost > 0 else Decimal("0")
|
||
|
||
# 预估可用天数
|
||
if daily_avg > 0:
|
||
estimated_days = float(available / daily_avg)
|
||
else:
|
||
estimated_days = 999
|
||
|
||
balance_info = {
|
||
"balance": float(balance),
|
||
"creditLimit": float(credit_limit),
|
||
"available": float(available),
|
||
"dailyAvgCost": round(float(daily_avg), 2),
|
||
"estimatedDays": round(estimated_days, 1)
|
||
}
|
||
|
||
# 3. 速率限制信息（查询近1分钟的请求数）
|
||
one_minute_ago = datetime.utcnow() - timedelta(minutes=1)
|
||
rpm_result = await db.execute(
|
||
select(func.count(ResourceUsage.id))
|
||
.where(
|
||
and_(
|
||
ResourceUsage.user_id == user_id,
|
||
ResourceUsage.created_at >= one_minute_ago
|
||
)
|
||
)
|
||
)
|
||
current_rpm = rpm_result.scalar() or 0
|
||
|
||
rate_limit_info = {
|
||
"currentRpm": current_rpm,
|
||
"rpmLimit": user.rpm_limit,
|
||
"allowed": current_rpm < user.rpm_limit
|
||
}
|
||
|
||
# 4. 活跃预警
|
||
alerts_result = await db.execute(
|
||
select(QuotaAlert)
|
||
.where(
|
||
and_(
|
||
QuotaAlert.user_id == user_id,
|
||
QuotaAlert.status == "active"
|
||
)
|
||
)
|
||
.order_by(QuotaAlert.created_at.desc())
|
||
)
|
||
alerts = alerts_result.scalars().all()
|
||
|
||
alert_type = None
|
||
if alerts:
|
||
# 优先级：balance_exhausted > balance_critical > balance_warning
|
||
priority_map = {
|
||
"balance_exhausted": 3,
|
||
"balance_critical": 2,
|
||
"balance_warning": 1
|
||
}
|
||
alerts_sorted = sorted(alerts, key=lambda a: priority_map.get(a.alert_type, 0), reverse=True)
|
||
alert_type = alerts_sorted[0].alert_type
|
||
|
||
return {
|
||
"hasQuota": available > 0,
|
||
"alertType": alert_type,
|
||
"balance": balance_info,
|
||
"rateLimit": rate_limit_info,
|
||
"activeAlerts": len(alerts),
|
||
"alerts": [
|
||
{
|
||
"id": str(alert.id),
|
||
"alertType": alert.alert_type,
|
||
"thresholdPercent": alert.threshold_percent,
|
||
"currentValue": float(alert.current_value) if alert.current_value else None,
|
||
"thresholdValue": float(alert.threshold_value) if alert.threshold_value else None,
|
||
"status": alert.status,
|
||
"createdAt": alert.created_at.isoformat()
|
||
}
|
||
for alert in alerts
|
||
]
|
||
}
|
||
|
||
|
||
async def get_channel_quota_summary(
|
||
channel_id: str,
|
||
db: AsyncSession
|
||
) -> Dict:
|
||
"""
|
||
获取渠道配额综合信息
|
||
|
||
Returns:
|
||
{
|
||
"hasQuota": bool,
|
||
"alertType": str | None,
|
||
"channelCredit": Decimal,
|
||
"monthlyUsage": Decimal,
|
||
"usagePercent": float
|
||
}
|
||
"""
|
||
from models import Channel, User, QuotaAlert
|
||
|
||
# 1. 获取渠道信息
|
||
result = await db.execute(
|
||
select(Channel).where(Channel.id == channel_id)
|
||
)
|
||
channel = result.scalar_one_or_none()
|
||
|
||
if not channel:
|
||
raise ValueError(f"渠道不存在: {channel_id}")
|
||
|
||
# 2. 计算本月使用量（从 AgentBillingRecord 和 ModelBillingRecord 合并统计）
|
||
now = datetime.utcnow()
|
||
month_start = datetime(now.year, now.month, 1)
|
||
|
||
# Agent 计费
|
||
agent_usage_result = await db.execute(
|
||
select(func.sum(AgentBillingRecord.cost))
|
||
.join(User, AgentBillingRecord.user_id == User.id)
|
||
.where(
|
||
and_(
|
||
User.channel_id == channel_id,
|
||
AgentBillingRecord.start_time >= month_start
|
||
)
|
||
)
|
||
)
|
||
agent_usage = agent_usage_result.scalar() or 0
|
||
|
||
# 模型调用计费
|
||
model_usage_result = await db.execute(
|
||
select(func.sum(ModelBillingRecord.total_cost))
|
||
.join(User, ModelBillingRecord.tenant_id == User.id)
|
||
.where(
|
||
and_(
|
||
User.channel_id == channel_id,
|
||
ModelBillingRecord.created_at >= month_start
|
||
)
|
||
)
|
||
)
|
||
model_usage = model_usage_result.scalar() or 0
|
||
|
||
monthly_usage = Decimal(str(agent_usage)) + Decimal(str(model_usage))
|
||
|
||
channel_credit = Decimal(str(channel.channel_credit or 0)) # 修复：使用正确的字段名 channel_credit
|
||
usage_percent = float(monthly_usage / channel_credit * 100) if channel_credit > 0 else 0
|
||
|
||
# 3. 检查活跃预警
|
||
alerts_result = await db.execute(
|
||
select(QuotaAlert)
|
||
.where(
|
||
and_(
|
||
QuotaAlert.channel_id == channel_id,
|
||
QuotaAlert.status == "active"
|
||
)
|
||
)
|
||
)
|
||
alerts = alerts_result.scalars().all()
|
||
alert_type = alerts[0].alert_type if alerts else None
|
||
|
||
return {
|
||
"hasQuota": channel_credit > monthly_usage,
|
||
"alertType": alert_type,
|
||
"channelCredit": float(channel_credit),
|
||
"monthlyUsage": float(monthly_usage),
|
||
"usagePercent": round(usage_percent, 2)
|
||
}
|
||
|