""" 资源管控模块 提供请求级别的资源管控、速率限制和配额检查 注意：计费数据已迁移到新表： - 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