forked from xiaohei/taiji-AI-PAD
415 lines
9.7 KiB
Python
415 lines
9.7 KiB
Python
"""
|
||
计费与资源管理逻辑
|
||
"""
|
||
|
||
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
|
||
|
||
|