""" 计费与资源管理逻辑 """ import math from datetime import datetime from decimal import Decimal from typing import Optional, Tuple from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from models import User, Channel, BillingRecord, Agent # ============= EU计算规则 ============= def calculate_eu(duration_seconds: int) -> int: """ 计算EU:1 EU = 10秒,不足10秒按1 EU计算 Args: duration_seconds: 调用时长(秒) Returns: EU数量 """ return math.ceil(duration_seconds / 10) def calculate_cost(eu: int, unit_price: Decimal = Decimal("0.01")) -> Decimal: """ 计算成本:1 EU = ¥0.01(默认) Args: eu: EU数量 unit_price: 单价(可配置) Returns: 成本金额 """ return Decimal(eu) * unit_price # ============= 余额与授信管理 ============= async def get_available_balance(user_id: str, db: AsyncSession) -> Tuple[Decimal, Decimal, Decimal]: """ 获取可用额度:账户余额 + 授信额度 Returns: (账户余额, 授信额度, 可用额度) """ result = await db.execute( select(User).where(User.id == user_id) ) user = result.scalar_one_or_none() if not user: return Decimal(0), Decimal(0), Decimal(0) balance = Decimal(str(user.balance)) credit_limit = Decimal(str(user.credit_limit)) available = balance + credit_limit return balance, credit_limit, available async def check_balance_sufficient( user_id: str, required_amount: Decimal, db: AsyncSession ) -> bool: """ 检查余额是否充足 Args: user_id: 用户ID required_amount: 所需金额 db: 数据库会话 Returns: 是否充足 """ balance, credit_limit, available = await get_available_balance(user_id, db) return available >= required_amount async def deduct_balance( user_id: str, amount: Decimal, db: AsyncSession, description: str = "消费" ) -> Tuple[bool, str]: """ 扣除余额(优先扣除账户余额,不足时使用授信额度) Args: user_id: 用户ID amount: 扣除金额 db: 数据库会话 description: 描述 Returns: (是否成功, 消息) """ result = await db.execute( select(User).where(User.id == user_id) ) user = result.scalar_one_or_none() if not user: return False, "用户不存在" balance = Decimal(str(user.balance)) credit_limit = Decimal(str(user.credit_limit)) available = balance + credit_limit if available < amount: return False, f"余额不足,当前可用额度: {available}, 需要: {amount}" # 优先扣除账户余额 if balance >= amount: user.balance = float(balance - amount) else: # 余额不足,使用授信额度 user.balance = 0 # 注意:授信额度是额度上限,不是实际金额,这里简化处理 # 实际应该有单独的授信使用记录表 await db.commit() return True, f"成功扣除 {amount} 元" async def add_balance( user_id: str, amount: Decimal, db: AsyncSession, description: str = "充值" ) -> Tuple[bool, str]: """ 增加余额 Args: user_id: 用户ID amount: 充值金额 db: 数据库会话 description: 描述 Returns: (是否成功, 消息) """ result = await db.execute( select(User).where(User.id == user_id) ) user = result.scalar_one_or_none() if not user: return False, "用户不存在" balance = Decimal(str(user.balance)) user.balance = float(balance + amount) await db.commit() return True, f"成功充值 {amount} 元,当前余额: {user.balance}" # ============= 计费记录创建 ============= async def create_billing_record( tenant_id: str, agent_id: str, agent_name: str, duration_seconds: int, db: AsyncSession, channel_id: Optional[str] = None, ) -> BillingRecord: """ 创建计费记录 Args: tenant_id: 租户ID agent_id: Agent ID agent_name: Agent名称 duration_seconds: 调用时长(秒) db: 数据库会话 channel_id: 渠道ID(可选) Returns: 计费记录 """ # 计算EU和成本 eu = calculate_eu(duration_seconds) cost = calculate_cost(eu) # 如果没有提供channel_id,从用户信息获取 if not channel_id: result = await db.execute( select(User.channel_id).where(User.id == tenant_id) ) row = result.first() if row: channel_id = str(row[0]) if row[0] else None # 创建计费记录 record = BillingRecord( timestamp=datetime.utcnow(), channel_id=channel_id, tenant_id=tenant_id, agent_id=agent_id, agent_name=agent_name, duration=duration_seconds, eu=eu, cost=cost, ) db.add(record) # 扣除用户余额 success, message = await deduct_balance(tenant_id, cost, db, f"Agent调用: {agent_name}") if not success: # 如果余额不足,记录但不扣除(实际应该阻止调用) # 这里简化处理,只记录 pass await db.commit() await db.refresh(record) return record # ============= 资源配额检查 ============= async def check_agent_quota( user_id: str, agent_id: str, db: AsyncSession ) -> Tuple[bool, str]: """ 检查Agent配额是否足够 Args: user_id: 用户ID agent_id: Agent ID db: 数据库会话 Returns: (是否有配额, 消息) """ # 简化实现:实际应该查询ResourceAllocation表 # 这里只是返回True,实际应该检查配额 return True, "配额充足" async def check_model_quota( user_id: str, model_name: str, db: AsyncSession ) -> Tuple[bool, int, int]: """ 检查模型配额(RPM/TPM) Args: user_id: 用户ID model_name: 模型名称 db: 数据库会话 Returns: (是否有配额, 当前RPM, 当前TPM) """ # 简化实现:实际应该查询ResourceAllocation和使用情况 # 这里只是返回默认值 return True, 60, 60000 # ============= 资源分配层级验证 ============= async def validate_resource_allocation( parent_type: str, # channel | tenant parent_id: str, resource_type: str, # agent | model resource_id: str, quantity: int, db: AsyncSession ) -> Tuple[bool, str]: """ 验证资源分配是否超过上级分配的资源 Args: parent_type: 父级类型 parent_id: 父级ID resource_type: 资源类型 resource_id: 资源ID quantity: 数量 db: 数据库会话 Returns: (是否有效, 消息) """ # 简化实现:实际应该递归检查上级的配额 # 这里只是返回True return True, "资源分配有效" # ============= 工作流限制验证 ============= def validate_workflow_nodes(nodes: list) -> Tuple[bool, str]: """ 验证工作流节点数量限制 Args: nodes: 节点列表 Returns: (是否有效, 消息) """ if len(nodes) > 3: return False, "工作流最多支持3个Agent节点" return True, "节点数量有效" # ============= 平台Agent资源固定配置 ============= PLATFORM_AGENT_CONFIG = { "cpu": 2.0, # 2核 "memory": 4.0, # 4GB } def get_platform_agent_resources() -> dict: """ 获取平台Agent的固定资源配置 Returns: 资源配置字典 """ return PLATFORM_AGENT_CONFIG.copy() # ============= 统计辅助函数 ============= async def calculate_monthly_cost(user_id: str, db: AsyncSession) -> Decimal: """ 计算用户本月消费 Args: user_id: 用户ID db: 数据库会话 Returns: 本月消费金额 """ from sqlalchemy import func # 获取本月第一天 now = datetime.utcnow() month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) result = await db.execute( select(func.sum(BillingRecord.cost)) .where(BillingRecord.tenant_id == user_id) .where(BillingRecord.timestamp >= month_start) ) total = result.scalar() return Decimal(str(total)) if total else Decimal(0) async def calculate_channel_commission( channel_id: str, start_date: datetime, end_date: datetime, db: AsyncSession ) -> Tuple[Decimal, Decimal]: """ 计算渠道佣金 Args: channel_id: 渠道ID start_date: 开始日期 end_date: 结束日期 db: 数据库会话 Returns: (总收入, 佣金金额) """ from sqlalchemy import func # 查询渠道佣金率 result = await db.execute( select(Channel.commission_rate).where(Channel.id == channel_id) ) row = result.first() commission_rate = Decimal(str(row[0])) / 100 if row else Decimal(0) # 查询总收入 result = await db.execute( select(func.sum(BillingRecord.cost)) .where(BillingRecord.channel_id == channel_id) .where(BillingRecord.timestamp >= start_date) .where(BillingRecord.timestamp <= end_date) ) total_revenue = result.scalar() total_revenue = Decimal(str(total_revenue)) if total_revenue else Decimal(0) commission = total_revenue * commission_rate return total_revenue, commission