6c6c48cf71
All checks were successful
continuous-integration/drone/push Build is passing
- 新增告警模块 (alerts): 告警规则配置与触发 - 新增成本管理模块 (cost): 成本统计与分析 - 新增配额模块 (quota): 配额管理与限制 - 新增微信模块 (wechat): 微信相关功能接口 - 新增缓存服务 (cache): Redis 缓存封装 - 新增请求日志中间件 (request_logger) - 新增异常处理和链路追踪中间件 - 更新 dashboard 前端展示 - 更新 SDK stats_client 功能
347 lines
12 KiB
Python
347 lines
12 KiB
Python
"""配额管理服务"""
|
||
import logging
|
||
from datetime import datetime, date, timedelta
|
||
from typing import Optional, Dict, Any, Tuple
|
||
from dataclasses import dataclass
|
||
|
||
from sqlalchemy.orm import Session
|
||
from sqlalchemy import func
|
||
|
||
from ..models.tenant import Tenant, Subscription
|
||
from ..models.stats import AICallEvent
|
||
from .cache import get_cache
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class QuotaConfig:
|
||
"""配额配置"""
|
||
daily_calls: int = 0 # 每日调用限制，0表示无限制
|
||
daily_tokens: int = 0 # 每日Token限制
|
||
monthly_calls: int = 0 # 每月调用限制
|
||
monthly_tokens: int = 0 # 每月Token限制
|
||
monthly_cost: float = 0 # 每月费用限制（元）
|
||
concurrent_calls: int = 0 # 并发调用限制
|
||
|
||
|
||
@dataclass
|
||
class QuotaUsage:
|
||
"""配额使用情况"""
|
||
daily_calls: int = 0
|
||
daily_tokens: int = 0
|
||
monthly_calls: int = 0
|
||
monthly_tokens: int = 0
|
||
monthly_cost: float = 0
|
||
|
||
|
||
@dataclass
|
||
class QuotaCheckResult:
|
||
"""配额检查结果"""
|
||
allowed: bool
|
||
reason: Optional[str] = None
|
||
quota_type: Optional[str] = None
|
||
limit: int = 0
|
||
used: int = 0
|
||
remaining: int = 0
|
||
|
||
|
||
class QuotaService:
|
||
"""配额管理服务
|
||
|
||
使用示例:
|
||
quota_service = QuotaService(db)
|
||
|
||
# 检查配额
|
||
result = quota_service.check_quota("qiqi", "tools")
|
||
if not result.allowed:
|
||
raise HTTPException(status_code=429, detail=result.reason)
|
||
|
||
# 获取使用情况
|
||
usage = quota_service.get_usage("qiqi", "tools")
|
||
"""
|
||
|
||
# 默认配额（当无订阅配置时使用）
|
||
DEFAULT_QUOTA = QuotaConfig(
|
||
daily_calls=1000,
|
||
daily_tokens=100000,
|
||
monthly_calls=30000,
|
||
monthly_tokens=3000000,
|
||
monthly_cost=100
|
||
)
|
||
|
||
def __init__(self, db: Session):
|
||
self.db = db
|
||
self._cache = get_cache()
|
||
|
||
def get_subscription(self, tenant_id: str, app_code: str) -> Optional[Subscription]:
|
||
"""获取租户订阅配置"""
|
||
return self.db.query(Subscription).filter(
|
||
Subscription.tenant_id == tenant_id,
|
||
Subscription.app_code == app_code,
|
||
Subscription.status == 'active'
|
||
).first()
|
||
|
||
def get_quota_config(self, tenant_id: str, app_code: str) -> QuotaConfig:
|
||
"""获取配额配置
|
||
|
||
Args:
|
||
tenant_id: 租户ID
|
||
app_code: 应用代码
|
||
|
||
Returns:
|
||
QuotaConfig实例
|
||
"""
|
||
# 尝试从缓存获取
|
||
cache_key = f"quota:config:{tenant_id}:{app_code}"
|
||
cached = self._cache.get(cache_key)
|
||
if cached:
|
||
return QuotaConfig(**cached)
|
||
|
||
# 从订阅表获取
|
||
subscription = self.get_subscription(tenant_id, app_code)
|
||
|
||
if subscription and subscription.quota:
|
||
quota = subscription.quota
|
||
config = QuotaConfig(
|
||
daily_calls=quota.get('daily_calls', 0),
|
||
daily_tokens=quota.get('daily_tokens', 0),
|
||
monthly_calls=quota.get('monthly_calls', 0),
|
||
monthly_tokens=quota.get('monthly_tokens', 0),
|
||
monthly_cost=quota.get('monthly_cost', 0),
|
||
concurrent_calls=quota.get('concurrent_calls', 0)
|
||
)
|
||
else:
|
||
config = self.DEFAULT_QUOTA
|
||
|
||
# 缓存5分钟
|
||
self._cache.set(cache_key, config.__dict__, ttl=300)
|
||
|
||
return config
|
||
|
||
def get_usage(self, tenant_id: str, app_code: str) -> QuotaUsage:
|
||
"""获取配额使用情况
|
||
|
||
Args:
|
||
tenant_id: 租户ID
|
||
app_code: 应用代码
|
||
|
||
Returns:
|
||
QuotaUsage实例
|
||
"""
|
||
today = date.today()
|
||
month_start = today.replace(day=1)
|
||
|
||
# 今日使用量
|
||
daily_stats = self.db.query(
|
||
func.count(AICallEvent.id).label('calls'),
|
||
func.coalesce(func.sum(AICallEvent.input_tokens + AICallEvent.output_tokens), 0).label('tokens')
|
||
).filter(
|
||
AICallEvent.tenant_id == tenant_id,
|
||
AICallEvent.app_code == app_code,
|
||
func.date(AICallEvent.created_at) == today
|
||
).first()
|
||
|
||
# 本月使用量
|
||
monthly_stats = self.db.query(
|
||
func.count(AICallEvent.id).label('calls'),
|
||
func.coalesce(func.sum(AICallEvent.input_tokens + AICallEvent.output_tokens), 0).label('tokens'),
|
||
func.coalesce(func.sum(AICallEvent.cost), 0).label('cost')
|
||
).filter(
|
||
AICallEvent.tenant_id == tenant_id,
|
||
AICallEvent.app_code == app_code,
|
||
func.date(AICallEvent.created_at) >= month_start
|
||
).first()
|
||
|
||
return QuotaUsage(
|
||
daily_calls=daily_stats.calls or 0,
|
||
daily_tokens=int(daily_stats.tokens or 0),
|
||
monthly_calls=monthly_stats.calls or 0,
|
||
monthly_tokens=int(monthly_stats.tokens or 0),
|
||
monthly_cost=float(monthly_stats.cost or 0)
|
||
)
|
||
|
||
def check_quota(
|
||
self,
|
||
tenant_id: str,
|
||
app_code: str,
|
||
estimated_tokens: int = 0
|
||
) -> QuotaCheckResult:
|
||
"""检查配额是否足够
|
||
|
||
Args:
|
||
tenant_id: 租户ID
|
||
app_code: 应用代码
|
||
estimated_tokens: 预估Token消耗
|
||
|
||
Returns:
|
||
QuotaCheckResult实例
|
||
"""
|
||
config = self.get_quota_config(tenant_id, app_code)
|
||
usage = self.get_usage(tenant_id, app_code)
|
||
|
||
# 检查日调用次数
|
||
if config.daily_calls > 0:
|
||
if usage.daily_calls >= config.daily_calls:
|
||
return QuotaCheckResult(
|
||
allowed=False,
|
||
reason=f"已达到每日调用限制 ({config.daily_calls} 次)",
|
||
quota_type="daily_calls",
|
||
limit=config.daily_calls,
|
||
used=usage.daily_calls,
|
||
remaining=0
|
||
)
|
||
|
||
# 检查日Token限制
|
||
if config.daily_tokens > 0:
|
||
if usage.daily_tokens + estimated_tokens > config.daily_tokens:
|
||
return QuotaCheckResult(
|
||
allowed=False,
|
||
reason=f"已达到每日Token限制 ({config.daily_tokens:,})",
|
||
quota_type="daily_tokens",
|
||
limit=config.daily_tokens,
|
||
used=usage.daily_tokens,
|
||
remaining=max(0, config.daily_tokens - usage.daily_tokens)
|
||
)
|
||
|
||
# 检查月调用次数
|
||
if config.monthly_calls > 0:
|
||
if usage.monthly_calls >= config.monthly_calls:
|
||
return QuotaCheckResult(
|
||
allowed=False,
|
||
reason=f"已达到每月调用限制 ({config.monthly_calls} 次)",
|
||
quota_type="monthly_calls",
|
||
limit=config.monthly_calls,
|
||
used=usage.monthly_calls,
|
||
remaining=0
|
||
)
|
||
|
||
# 检查月Token限制
|
||
if config.monthly_tokens > 0:
|
||
if usage.monthly_tokens + estimated_tokens > config.monthly_tokens:
|
||
return QuotaCheckResult(
|
||
allowed=False,
|
||
reason=f"已达到每月Token限制 ({config.monthly_tokens:,})",
|
||
quota_type="monthly_tokens",
|
||
limit=config.monthly_tokens,
|
||
used=usage.monthly_tokens,
|
||
remaining=max(0, config.monthly_tokens - usage.monthly_tokens)
|
||
)
|
||
|
||
# 检查月费用限制
|
||
if config.monthly_cost > 0:
|
||
if usage.monthly_cost >= config.monthly_cost:
|
||
return QuotaCheckResult(
|
||
allowed=False,
|
||
reason=f"已达到每月费用限制 (¥{config.monthly_cost:.2f})",
|
||
quota_type="monthly_cost",
|
||
limit=int(config.monthly_cost * 100), # 转为分
|
||
used=int(usage.monthly_cost * 100),
|
||
remaining=max(0, int((config.monthly_cost - usage.monthly_cost) * 100))
|
||
)
|
||
|
||
# 所有检查通过
|
||
return QuotaCheckResult(
|
||
allowed=True,
|
||
quota_type="daily_calls",
|
||
limit=config.daily_calls,
|
||
used=usage.daily_calls,
|
||
remaining=max(0, config.daily_calls - usage.daily_calls) if config.daily_calls > 0 else -1
|
||
)
|
||
|
||
def get_quota_summary(self, tenant_id: str, app_code: str) -> Dict[str, Any]:
|
||
"""获取配额汇总信息
|
||
|
||
Returns:
|
||
包含配额配置和使用情况的字典
|
||
"""
|
||
config = self.get_quota_config(tenant_id, app_code)
|
||
usage = self.get_usage(tenant_id, app_code)
|
||
|
||
def calc_percentage(used: int, limit: int) -> float:
|
||
if limit <= 0:
|
||
return 0
|
||
return min(100, round(used / limit * 100, 1))
|
||
|
||
return {
|
||
"config": {
|
||
"daily_calls": config.daily_calls,
|
||
"daily_tokens": config.daily_tokens,
|
||
"monthly_calls": config.monthly_calls,
|
||
"monthly_tokens": config.monthly_tokens,
|
||
"monthly_cost": config.monthly_cost
|
||
},
|
||
"usage": {
|
||
"daily_calls": usage.daily_calls,
|
||
"daily_tokens": usage.daily_tokens,
|
||
"monthly_calls": usage.monthly_calls,
|
||
"monthly_tokens": usage.monthly_tokens,
|
||
"monthly_cost": round(usage.monthly_cost, 2)
|
||
},
|
||
"percentage": {
|
||
"daily_calls": calc_percentage(usage.daily_calls, config.daily_calls),
|
||
"daily_tokens": calc_percentage(usage.daily_tokens, config.daily_tokens),
|
||
"monthly_calls": calc_percentage(usage.monthly_calls, config.monthly_calls),
|
||
"monthly_tokens": calc_percentage(usage.monthly_tokens, config.monthly_tokens),
|
||
"monthly_cost": calc_percentage(int(usage.monthly_cost * 100), int(config.monthly_cost * 100))
|
||
}
|
||
}
|
||
|
||
def update_quota(
|
||
self,
|
||
tenant_id: str,
|
||
app_code: str,
|
||
quota_config: Dict[str, Any]
|
||
) -> Subscription:
|
||
"""更新配额配置
|
||
|
||
Args:
|
||
tenant_id: 租户ID
|
||
app_code: 应用代码
|
||
quota_config: 配额配置字典
|
||
|
||
Returns:
|
||
更新后的Subscription实例
|
||
"""
|
||
subscription = self.get_subscription(tenant_id, app_code)
|
||
|
||
if not subscription:
|
||
# 创建新订阅
|
||
subscription = Subscription(
|
||
tenant_id=tenant_id,
|
||
app_code=app_code,
|
||
start_date=date.today(),
|
||
quota=quota_config,
|
||
status='active'
|
||
)
|
||
self.db.add(subscription)
|
||
else:
|
||
# 更新现有订阅
|
||
subscription.quota = quota_config
|
||
|
||
self.db.commit()
|
||
self.db.refresh(subscription)
|
||
|
||
# 清除缓存
|
||
cache_key = f"quota:config:{tenant_id}:{app_code}"
|
||
self._cache.delete(cache_key)
|
||
|
||
return subscription
|
||
|
||
|
||
def check_quota_middleware(
|
||
db: Session,
|
||
tenant_id: str,
|
||
app_code: str,
|
||
estimated_tokens: int = 0
|
||
) -> QuotaCheckResult:
|
||
"""配额检查中间件函数
|
||
|
||
可在路由中使用:
|
||
result = check_quota_middleware(db, "qiqi", "tools")
|
||
if not result.allowed:
|
||
raise HTTPException(status_code=429, detail=result.reason)
|
||
"""
|
||
service = QuotaService(db)
|
||
return service.check_quota(tenant_id, app_code, estimated_tokens)
|