577 lines
15 KiB
Python
577 lines
15 KiB
Python
"""
|
|
预付费配额管理模块
|
|
配额预警与限制
|
|
"""
|
|
|
|
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, BillingRecord, ResourceAllocation
|
|
|
|
|
|
# 配额预警阈值配置
|
|
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 = Decimal(str(user.balance))
|
|
credit_limit = Decimal(str(user.credit_limit))
|
|
available = balance + credit_limit
|
|
|
|
# 获取用户平均日消费
|
|
thirty_days_ago = datetime.utcnow() - timedelta(days=30)
|
|
cost_result = await db.execute(
|
|
select(func.sum(BillingRecord.cost))
|
|
.where(
|
|
and_(
|
|
BillingRecord.tenant_id == user_id,
|
|
BillingRecord.timestamp >= thirty_days_ago,
|
|
)
|
|
)
|
|
)
|
|
total_cost = cost_result.scalar() or 0
|
|
daily_avg = Decimal(str(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))
|
|
|
|
# 获取渠道下所有租户的总消费
|
|
thirty_days_ago = datetime.utcnow() - timedelta(days=30)
|
|
cost_result = await db.execute(
|
|
select(func.sum(BillingRecord.cost))
|
|
.where(
|
|
and_(
|
|
BillingRecord.channel_id == channel_id,
|
|
BillingRecord.timestamp >= thirty_days_ago,
|
|
)
|
|
)
|
|
)
|
|
total_cost = cost_result.scalar() or 0
|
|
|
|
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",
|
|
)
|
|
)
|
|
)
|
|
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分钟的调用次数
|
|
one_minute_ago = datetime.utcnow() - timedelta(minutes=1)
|
|
|
|
result = await db.execute(
|
|
select(func.count(BillingRecord.id))
|
|
.where(
|
|
and_(
|
|
BillingRecord.tenant_id == user_id,
|
|
BillingRecord.timestamp >= one_minute_ago,
|
|
)
|
|
)
|
|
)
|
|
current_rpm = result.scalar() or 0
|
|
|
|
# 获取用户的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, BillingRecord, 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 = Decimal(str(user.balance))
|
|
credit_limit = Decimal(str(user.credit_limit or 0))
|
|
available = balance + credit_limit
|
|
|
|
# 计算近30天平均日消费
|
|
thirty_days_ago = datetime.utcnow() - timedelta(days=30)
|
|
cost_result = await db.execute(
|
|
select(func.sum(BillingRecord.cost))
|
|
.where(
|
|
and_(
|
|
BillingRecord.tenant_id == user_id,
|
|
BillingRecord.timestamp >= thirty_days_ago,
|
|
)
|
|
)
|
|
)
|
|
total_cost = cost_result.scalar() or 0
|
|
daily_avg = Decimal(str(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, BillingRecord
|
|
|
|
# 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. 计算本月使用量
|
|
now = datetime.utcnow()
|
|
month_start = datetime(now.year, now.month, 1)
|
|
|
|
usage_result = await db.execute(
|
|
select(func.sum(BillingRecord.cost))
|
|
.join(User, BillingRecord.tenant_id == User.id)
|
|
.where(
|
|
and_(
|
|
User.channel_id == channel_id,
|
|
BillingRecord.timestamp >= month_start
|
|
)
|
|
)
|
|
)
|
|
monthly_usage = Decimal(str(usage_result.scalar() or 0))
|
|
|
|
channel_credit = Decimal(str(channel.credit_limit or 0))
|
|
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)
|
|
}
|
|
|