forked from xiaohei/taiji-AI-PAD
396 lines
14 KiB
Python
396 lines
14 KiB
Python
"""
|
||
资源管控模块
|
||
提供请求级别的资源管控、速率限制和配额检查
|
||
|
||
注意：计费数据已迁移到新表：
|
||
- AgentBillingRecord: Agent 运行时计费
|
||
- ModelBillingRecord: 模型调用计费（LiteLLM）
|
||
旧的 BillingRecord 表已废弃，不再使用。
|
||
"""
|
||
|
||
from datetime import datetime, timedelta
|
||
from decimal import Decimal
|
||
from typing import Dict, Optional, Tuple
|
||
from sqlalchemy import select, func, and_, update
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from fastapi import HTTPException, status
|
||
import structlog
|
||
|
||
from models import User, Channel, ResourceUsage, AgentBillingRecord
|
||
from app.quota_manager import check_user_balance_quota, check_channel_quota
|
||
from app.db_utils import with_retry
|
||
|
||
logger = structlog.get_logger()
|
||
|
||
|
||
# 速率限制窗口（秒）
|
||
RATE_LIMIT_WINDOW = 60 # 1分钟窗口
|
||
|
||
|
||
class ResourceController:
|
||
"""资源控制器 - 统一管理资源配额和速率限制"""
|
||
|
||
def __init__(self):
|
||
self._rate_limit_cache: Dict[str, Dict] = {}
|
||
|
||
async def check_and_enforce(
|
||
self,
|
||
user_id: str,
|
||
resource_type: str, # 'tool', 'agent', 'model'
|
||
resource_id: Optional[str] = None,
|
||
estimated_cost: Optional[Decimal] = None,
|
||
db: AsyncSession = None
|
||
) -> Tuple[bool, Optional[str], Dict]:
|
||
"""
|
||
检查并执行资源管控
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
resource_type: 资源类型
|
||
resource_id: 资源ID
|
||
estimated_cost: 预估成本
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
(是否允许, 拒绝原因, 详细信息)
|
||
"""
|
||
# 1. 检查用户余额配额
|
||
balance_ok, alert_type, balance_info = await check_user_balance_quota(user_id, db)
|
||
if not balance_ok:
|
||
logger.warning(
|
||
"resource_control_balance_failed",
|
||
user_id=user_id,
|
||
alert_type=alert_type,
|
||
balance_info=balance_info
|
||
)
|
||
return False, "insufficient_balance", {
|
||
"message": "账户余额不足",
|
||
"alert_type": alert_type,
|
||
**balance_info
|
||
}
|
||
|
||
# 2. 检查速率限制（RPM）
|
||
rpm_ok, rpm_reason, rpm_info = await self._check_rate_limit(
|
||
user_id, resource_type, db
|
||
)
|
||
if not rpm_ok:
|
||
logger.warning(
|
||
"resource_control_rate_limit_exceeded",
|
||
user_id=user_id,
|
||
resource_type=resource_type,
|
||
rpm_info=rpm_info
|
||
)
|
||
return False, rpm_reason, {
|
||
"message": "超过速率限制",
|
||
**rpm_info
|
||
}
|
||
|
||
# 3. 检查渠道配额（如果用户属于某个渠道）
|
||
result = await db.execute(
|
||
select(User).where(User.id == user_id)
|
||
)
|
||
user = result.scalar_one_or_none()
|
||
if user and user.channel_id:
|
||
channel_ok, channel_alert, channel_info = await check_channel_quota(
|
||
str(user.channel_id), db
|
||
)
|
||
if not channel_ok or channel_alert == "quota_critical":
|
||
logger.warning(
|
||
"resource_control_channel_quota_exceeded",
|
||
user_id=user_id,
|
||
channel_id=str(user.channel_id),
|
||
channel_info=channel_info
|
||
)
|
||
return False, "channel_quota_exceeded", {
|
||
"message": "渠道配额不足",
|
||
**channel_info
|
||
}
|
||
|
||
# 4. 检查预估成本是否会导致余额不足
|
||
if estimated_cost and estimated_cost > 0:
|
||
available = Decimal(str(balance_info.get("available", 0)))
|
||
if estimated_cost > available:
|
||
logger.warning(
|
||
"resource_control_insufficient_for_cost",
|
||
user_id=user_id,
|
||
estimated_cost=float(estimated_cost),
|
||
available=float(available)
|
||
)
|
||
return False, "insufficient_funds_for_operation", {
|
||
"message": "余额不足以支付本次操作",
|
||
"estimated_cost": float(estimated_cost),
|
||
"available": float(available)
|
||
}
|
||
|
||
# 所有检查通过
|
||
logger.info(
|
||
"resource_control_passed",
|
||
user_id=user_id,
|
||
resource_type=resource_type,
|
||
resource_id=resource_id
|
||
)
|
||
|
||
return True, None, {
|
||
"allowed": True,
|
||
"balance_info": balance_info,
|
||
"rpm_info": rpm_info
|
||
}
|
||
|
||
async def _check_rate_limit(
|
||
self,
|
||
user_id: str,
|
||
resource_type: str,
|
||
db: AsyncSession
|
||
) -> Tuple[bool, Optional[str], Dict]:
|
||
"""
|
||
检查速率限制（每分钟请求数 RPM）
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
resource_type: 资源类型
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
(是否允许, 拒绝原因, 速率信息)
|
||
"""
|
||
# 获取用户配置的RPM限制
|
||
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": "用户不存在"}
|
||
|
||
# 默认RPM限制（可从用户配置或系统配置获取）
|
||
rpm_limit = getattr(user, 'rpm_limit', 60) # 默认60 RPM
|
||
|
||
# 计算当前窗口的起始时间
|
||
now = datetime.utcnow()
|
||
window_start = now - timedelta(seconds=RATE_LIMIT_WINDOW)
|
||
|
||
# 查询当前窗口内的API调用次数
|
||
result = await db.execute(
|
||
select(func.count(ResourceUsage.id))
|
||
.where(
|
||
and_(
|
||
ResourceUsage.user_id == user_id,
|
||
ResourceUsage.period_start >= window_start,
|
||
ResourceUsage.period_start <= now
|
||
)
|
||
)
|
||
)
|
||
|
||
current_rpm = result.scalar() or 0
|
||
|
||
# 判断是否超过限制
|
||
if current_rpm >= rpm_limit:
|
||
return False, "rate_limit_exceeded", {
|
||
"current_rpm": current_rpm,
|
||
"rpm_limit": rpm_limit,
|
||
"window_seconds": RATE_LIMIT_WINDOW,
|
||
"retry_after": RATE_LIMIT_WINDOW
|
||
}
|
||
|
||
return True, None, {
|
||
"current_rpm": current_rpm,
|
||
"rpm_limit": rpm_limit,
|
||
"remaining": rpm_limit - current_rpm,
|
||
"window_seconds": RATE_LIMIT_WINDOW
|
||
}
|
||
|
||
async def record_resource_consumption(
|
||
self,
|
||
user_id: str,
|
||
resource_type: str,
|
||
resource_id: Optional[str],
|
||
cost: Decimal,
|
||
execution_time_ms: float,
|
||
cpu_usage: float = 0.0,
|
||
memory_usage: float = 0.0,
|
||
network_io: float = 0.0,
|
||
db: AsyncSession = None
|
||
) -> None:
|
||
"""
|
||
记录资源消耗
|
||
|
||
使用行锁和重试机制防止并发更新丢失。
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
resource_type: 资源类型
|
||
resource_id: 资源ID
|
||
cost: 成本
|
||
execution_time_ms: 执行时间（毫秒）
|
||
cpu_usage: CPU使用率
|
||
memory_usage: 内存使用（MB）
|
||
network_io: 网络IO（KB）
|
||
db: 数据库会话
|
||
"""
|
||
async def _do_record():
|
||
now = datetime.utcnow()
|
||
|
||
# 记录到ResourceUsage表（每小时聚合）
|
||
period_start = now.replace(minute=0, second=0, microsecond=0)
|
||
period_end = period_start + timedelta(hours=1)
|
||
|
||
# 计算增量值
|
||
cpu_increment = cpu_usage * (execution_time_ms / 1000.0)
|
||
memory_increment = memory_usage * (execution_time_ms / 1000.0)
|
||
network_increment = int(network_io * 1024)
|
||
|
||
# 使用行锁保护并发更新（FOR UPDATE）
|
||
result = await db.execute(
|
||
select(ResourceUsage)
|
||
.where(
|
||
and_(
|
||
ResourceUsage.user_id == user_id,
|
||
ResourceUsage.period_start == period_start,
|
||
ResourceUsage.granularity == "hourly"
|
||
)
|
||
)
|
||
.with_for_update() # 行锁
|
||
)
|
||
usage = result.scalar_one_or_none()
|
||
|
||
if usage:
|
||
# 使用原子更新语句，而不是 ORM 属性修改
|
||
await db.execute(
|
||
update(ResourceUsage)
|
||
.where(ResourceUsage.id == usage.id)
|
||
.values(
|
||
cpu_seconds=ResourceUsage.cpu_seconds + cpu_increment,
|
||
memory_mb_seconds=ResourceUsage.memory_mb_seconds + memory_increment,
|
||
network_bytes=ResourceUsage.network_bytes + network_increment,
|
||
api_calls=ResourceUsage.api_calls + 1
|
||
)
|
||
)
|
||
else:
|
||
# 创建新记录
|
||
usage = ResourceUsage(
|
||
user_id=user_id,
|
||
agent_id=resource_id if resource_type == "agent" else None,
|
||
cpu_seconds=cpu_increment,
|
||
memory_mb_seconds=memory_increment,
|
||
network_bytes=network_increment,
|
||
storage_bytes=0,
|
||
api_calls=1,
|
||
period_start=period_start,
|
||
period_end=period_end,
|
||
granularity="hourly"
|
||
)
|
||
db.add(usage)
|
||
|
||
# 记录到 AgentBillingRecord 表（新计费表）
|
||
# 注意：这里只记录 Agent 类型的资源消耗
|
||
# 模型调用计费由 LiteLLM Callback 处理，存入 ModelBillingRecord
|
||
if resource_type == "agent":
|
||
# 查询用户的 channel_id
|
||
user_result = await db.execute(
|
||
select(User).where(User.id == user_id)
|
||
)
|
||
user = user_result.scalar_one_or_none()
|
||
user_channel_id = user.channel_id if user else None
|
||
|
||
# ✅ 确保 agent_name 不为空字符串
|
||
agent_name = resource_id if resource_id else "unknown"
|
||
|
||
billing_record = AgentBillingRecord(
|
||
user_id=user_id,
|
||
channel_id=user_channel_id, # 从用户表获取 channel_id
|
||
agent_name=agent_name, # ✅ 使用验证后的名称
|
||
agent_type="custom", # 默认为自定义 Agent
|
||
is_platform_agent=False,
|
||
duration_seconds=int(execution_time_ms / 1000),
|
||
cpu_seconds=cpu_usage * (execution_time_ms / 1000.0),
|
||
memory_gb_seconds=memory_usage / 1024 * (execution_time_ms / 1000.0),
|
||
request_count=1,
|
||
cost=cost,
|
||
eu_consumed=float(cost), # EU = Cost（1 EU = 1 美元）
|
||
start_time=now,
|
||
period_start=now,
|
||
)
|
||
db.add(billing_record)
|
||
|
||
await db.commit()
|
||
|
||
logger.info(
|
||
"resource_consumption_recorded",
|
||
user_id=user_id,
|
||
resource_type=resource_type,
|
||
cost=float(cost),
|
||
execution_time_ms=execution_time_ms
|
||
)
|
||
|
||
try:
|
||
# 使用重试机制处理数据库冲突
|
||
await with_retry(_do_record, max_retries=3)
|
||
|
||
except Exception as e:
|
||
logger.error(
|
||
"resource_consumption_recording_failed",
|
||
user_id=user_id,
|
||
error=str(e),
|
||
exc_info=True
|
||
)
|
||
await db.rollback()
|
||
|
||
|
||
# 全局资源控制器实例
|
||
resource_controller = ResourceController()
|
||
|
||
|
||
async def enforce_resource_control(
|
||
user_id: str,
|
||
resource_type: str,
|
||
resource_id: Optional[str] = None,
|
||
estimated_cost: Optional[Decimal] = None,
|
||
db: AsyncSession = None
|
||
) -> Dict:
|
||
"""
|
||
执行资源管控检查（供路由使用）
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
resource_type: 资源类型
|
||
resource_id: 资源ID
|
||
estimated_cost: 预估成本
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
资源管控信息字典
|
||
|
||
Raises:
|
||
HTTPException: 如果资源管控检查失败
|
||
"""
|
||
allowed, reason, info = await resource_controller.check_and_enforce(
|
||
user_id=user_id,
|
||
resource_type=resource_type,
|
||
resource_id=resource_id,
|
||
estimated_cost=estimated_cost,
|
||
db=db
|
||
)
|
||
|
||
if not allowed:
|
||
# 根据不同的拒绝原因返回不同的HTTP状态码
|
||
status_code_map = {
|
||
"insufficient_balance": status.HTTP_402_PAYMENT_REQUIRED,
|
||
"rate_limit_exceeded": status.HTTP_429_TOO_MANY_REQUESTS,
|
||
"channel_quota_exceeded": status.HTTP_403_FORBIDDEN,
|
||
"insufficient_funds_for_operation": status.HTTP_402_PAYMENT_REQUIRED,
|
||
"user_not_found": status.HTTP_404_NOT_FOUND,
|
||
}
|
||
|
||
status_code = status_code_map.get(reason, status.HTTP_403_FORBIDDEN)
|
||
|
||
raise HTTPException(
|
||
status_code=status_code,
|
||
detail={
|
||
"error": reason,
|
||
"message": info.get("message", "资源管控检查失败"),
|
||
**info
|
||
}
|
||
)
|
||
|
||
return info
|