forked from xiaohei/taiji-AI-PAD
1116 lines
32 KiB
Python
1116 lines
32 KiB
Python
"""
|
||
计费与资源管理逻辑
|
||
|
||
EU（执行单元）计算说明：
|
||
=======================
|
||
|
||
**重要：1 EU = 1 美元**
|
||
|
||
EU 是系统的计费单位，直接等于成本（美元）。
|
||
eu_consumed 字段的值应该等于 cost 字段的值。
|
||
|
||
计费方式：
|
||
1. Agent 运行时长计费
|
||
- 平台 Agent：按小时固定费率（$0.10-$0.15/小时）
|
||
- 自定义 Agent：按资源使用量（CPU $0.05/核/小时 + 内存 $0.01/GB/小时）
|
||
- EU = Cost（美元）
|
||
|
||
2. 模型 Token 计费
|
||
- 按模型定价计算成本
|
||
- EU = Cost（美元）
|
||
|
||
数据存储：
|
||
- Agent计费记录：agent_billing_records 表
|
||
- 模型计费记录：model_billing_records 表
|
||
- 用户余额：balances 表（Balance模型）
|
||
|
||
注意：User.eu_balance 和 User.balance 字段已废弃，请使用 Balance 表
|
||
"""
|
||
|
||
import math
|
||
import logging
|
||
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, Balance, AgentBillingRecord
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ============= EU定价配置（已废弃） =============
|
||
#
|
||
# 重要：1 EU = 1 美元
|
||
# EU 现在直接等于成本（美元），不再需要单独的定价配置
|
||
# 以下配置仅用于向后兼容旧的 BillingRecord 表
|
||
|
||
# EU单价配置（按订阅等级）- 已废弃，仅用于旧表兼容
|
||
EU_PRICING = {
|
||
"free": Decimal("0.015"), # 入门级（Starter）：$0.015 / EU
|
||
"starter": Decimal("0.015"), # 入门级别名
|
||
"pro": Decimal("0.02"), # 专业级（Pro）：$0.02 / EU
|
||
"enterprise": Decimal("0.03"), # 企业级（Enterprise）：$0.03 / EU
|
||
}
|
||
|
||
# 默认单价（未知等级时使用专业级价格）- 已废弃
|
||
DEFAULT_EU_PRICE = Decimal("0.02")
|
||
|
||
|
||
def get_eu_price_by_tier(subscription_tier: str) -> Decimal:
|
||
"""
|
||
根据订阅等级获取EU单价
|
||
|
||
⚠️ 已废弃：此函数仅用于旧的 BillingRecord 表兼容
|
||
新的计费逻辑中，1 EU = 1 美元，不需要单独的定价
|
||
|
||
Args:
|
||
subscription_tier: 订阅等级（free/starter/pro/enterprise）
|
||
|
||
Returns:
|
||
EU单价（USD）
|
||
"""
|
||
import warnings
|
||
warnings.warn(
|
||
"get_eu_price_by_tier() 已废弃，新计费逻辑中 1 EU = 1 美元",
|
||
DeprecationWarning,
|
||
stacklevel=2
|
||
)
|
||
tier = subscription_tier.lower() if subscription_tier else "pro"
|
||
return EU_PRICING.get(tier, DEFAULT_EU_PRICE)
|
||
|
||
|
||
# ============= EU计算规则 =============
|
||
#
|
||
# 重要：1 EU = 1 美元
|
||
# EU 应该等于 Cost（成本），不再基于时间计算
|
||
|
||
def calculate_eu(duration_seconds: int) -> int:
|
||
"""
|
||
⚠️ 已废弃：基于时间的 EU 计算
|
||
|
||
此函数使用旧的计算方式（1 EU = 10秒），已废弃。
|
||
新的计费逻辑中，EU = Cost（美元），即 1 EU = 1 美元。
|
||
|
||
请使用 calculate_agent_cost_by_resources() 或 calculate_platform_agent_cost()
|
||
计算成本，然后将成本值作为 EU 值。
|
||
|
||
Args:
|
||
duration_seconds: 调用时长（秒）
|
||
|
||
Returns:
|
||
EU数量（已废弃的计算方式）
|
||
"""
|
||
import warnings
|
||
warnings.warn(
|
||
"calculate_eu() 已废弃，新计费逻辑中 EU = Cost（美元）",
|
||
DeprecationWarning,
|
||
stacklevel=2
|
||
)
|
||
return math.ceil(duration_seconds / 10)
|
||
|
||
|
||
def calculate_cost(eu: int, unit_price: Decimal = None, subscription_tier: str = None) -> Decimal:
|
||
"""
|
||
⚠️ 已废弃：基于 EU 数量计算成本
|
||
|
||
此函数使用旧的计算方式，已废弃。
|
||
新的计费逻辑中，EU = Cost（美元），即 1 EU = 1 美元。
|
||
|
||
Args:
|
||
eu: EU数量
|
||
unit_price: 单价（可直接指定，优先级最高）
|
||
subscription_tier: 订阅等级（用于自动获取单价）
|
||
|
||
Returns:
|
||
成本金额（USD）
|
||
"""
|
||
import warnings
|
||
warnings.warn(
|
||
"calculate_cost() 已废弃，新计费逻辑中 EU = Cost（美元）",
|
||
DeprecationWarning,
|
||
stacklevel=2
|
||
)
|
||
if unit_price is None:
|
||
unit_price = get_eu_price_by_tier(subscription_tier) if subscription_tier else DEFAULT_EU_PRICE
|
||
return Decimal(eu) * unit_price
|
||
|
||
|
||
# ============= 余额与授信管理 =============
|
||
|
||
async def get_available_balance(user_id: str, db: AsyncSession) -> Tuple[Decimal, Decimal, Decimal]:
|
||
"""
|
||
获取可用额度：账户余额 + 授信额度
|
||
|
||
Returns:
|
||
(账户余额, 授信额度, 可用额度)
|
||
"""
|
||
# 从 Balance 表获取余额
|
||
balance_result = await db.execute(
|
||
select(Balance).where(Balance.user_id == user_id)
|
||
)
|
||
balance_obj = balance_result.scalar_one_or_none()
|
||
|
||
# 从 User 表获取授信额度
|
||
user_result = await db.execute(
|
||
select(User).where(User.id == user_id)
|
||
)
|
||
user = user_result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
return Decimal(0), Decimal(0), Decimal(0)
|
||
|
||
balance = Decimal(str(balance_obj.eu_balance)) if balance_obj else Decimal(0)
|
||
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 = "消费",
|
||
auto_commit: bool = True # 🔧 修改：默认自动提交，确保扣款立即生效
|
||
) -> Tuple[bool, str]:
|
||
"""
|
||
扣除余额（优先扣除账户余额，不足时使用授信额度）
|
||
|
||
🔧 安全修复：使用真正的原子操作，在一条SQL中完成余额检查和扣款，
|
||
彻底避免竞态条件。
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
amount: 扣除金额（自动统一精度到6位小数）
|
||
db: 数据库会话
|
||
description: 描述
|
||
auto_commit: 是否自动提交（默认True，确保扣款立即生效）
|
||
|
||
Returns:
|
||
(是否成功, 消息)
|
||
|
||
Note:
|
||
修复后默认 auto_commit=True，确保扣款立即生效，避免长时间事务导致的问题。
|
||
如需批量操作，请明确设置 auto_commit=False。
|
||
"""
|
||
from sqlalchemy import text
|
||
|
||
# 🔧 修复1: 统一精度到6位小数，避免累积误差
|
||
amount = amount.quantize(Decimal('0.000001'))
|
||
|
||
# 从 User 表获取授信额度
|
||
user_result = await db.execute(
|
||
select(User).where(User.id == user_id)
|
||
)
|
||
user = user_result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
return False, "用户不存在"
|
||
|
||
credit_limit = Decimal(str(user.credit_limit))
|
||
amount_str = str(amount)
|
||
credit_str = str(credit_limit)
|
||
|
||
# 🔧 修复2: 真正的原子操作 - 分步执行但在同一事务中
|
||
# Step 1: 确保余额记录存在（生成UUID）
|
||
await db.execute(
|
||
text("""
|
||
INSERT INTO balances (id, user_id, eu_balance, created_at, updated_at)
|
||
VALUES (gen_random_uuid(), :uid, 0.0, NOW(), NOW())
|
||
ON CONFLICT (user_id) DO NOTHING
|
||
"""),
|
||
{"uid": user_id}
|
||
)
|
||
|
||
# Step 2: 原子更新 - 只有在余额充足时才扣款
|
||
result = await db.execute(
|
||
text("""
|
||
WITH current AS (
|
||
SELECT eu_balance FROM balances WHERE user_id = :uid FOR UPDATE
|
||
)
|
||
UPDATE balances
|
||
SET eu_balance = eu_balance - CAST(:amount AS NUMERIC(15, 6)),
|
||
updated_at = NOW()
|
||
WHERE user_id = :uid
|
||
AND (SELECT eu_balance FROM current) + CAST(:credit AS NUMERIC(15, 6)) >= CAST(:amount AS NUMERIC(15, 6))
|
||
RETURNING
|
||
eu_balance as new_balance,
|
||
eu_balance + CAST(:amount AS NUMERIC(15, 6)) as old_balance
|
||
"""),
|
||
{"uid": user_id, "amount": amount_str, "credit": credit_str}
|
||
)
|
||
|
||
row = result.fetchone()
|
||
|
||
if row is None:
|
||
# UPDATE 影响了0行，说明余额不足
|
||
# 查询当前余额用于错误信息
|
||
balance_result = await db.execute(
|
||
text("SELECT eu_balance FROM balances WHERE user_id = :uid"),
|
||
{"uid": user_id}
|
||
)
|
||
balance_row = balance_result.fetchone()
|
||
current = Decimal(str(balance_row[0])) if balance_row else Decimal("0")
|
||
available = current + credit_limit
|
||
|
||
if auto_commit:
|
||
await db.commit() # 即使失败也提交（没有实际修改）
|
||
|
||
return False, f"余额不足，当前可用: {available:.6f}, 需要: {amount:.6f}"
|
||
|
||
new_balance = Decimal(str(row[0]))
|
||
old_balance = Decimal(str(row[1]))
|
||
|
||
# 可选：自动提交
|
||
if auto_commit:
|
||
await db.commit()
|
||
logger.info(
|
||
f"💰 扣款成功并已提交: user={user_id[:8]}, "
|
||
f"金额={amount:.6f}, 余额: {old_balance:.6f} → {new_balance:.6f}"
|
||
)
|
||
else:
|
||
logger.info(
|
||
f"💰 扣款成功待提交: user={user_id[:8]}, "
|
||
f"金额={amount:.6f}, 余额: {old_balance:.6f} → {new_balance:.6f}"
|
||
)
|
||
|
||
return True, f"成功扣除 {amount:.6f} 元，余额: {old_balance:.6f} → {new_balance:.6f}"
|
||
|
||
|
||
async def add_balance(
|
||
user_id: str,
|
||
amount: Decimal,
|
||
db: AsyncSession,
|
||
description: str = "充值",
|
||
auto_commit: bool = True # 🔧 修改：默认自动提交，与deduct_balance保持一致
|
||
) -> Tuple[bool, str]:
|
||
"""
|
||
增加余额
|
||
|
||
🔧 安全修复：使用原子操作，统一精度处理
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
amount: 充值金额（自动统一精度到6位小数）
|
||
db: 数据库会话
|
||
description: 描述
|
||
auto_commit: 是否自动提交（默认True）
|
||
|
||
Returns:
|
||
(是否成功, 消息)
|
||
|
||
Note:
|
||
修复后默认 auto_commit=True，与 deduct_balance 保持一致。
|
||
"""
|
||
from sqlalchemy import text
|
||
|
||
# 🔧 修复：统一精度到6位小数
|
||
amount = amount.quantize(Decimal('0.000001'))
|
||
|
||
# 检查用户是否存在
|
||
user_result = await db.execute(
|
||
select(User).where(User.id == user_id)
|
||
)
|
||
user = user_result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
return False, "用户不存在"
|
||
|
||
amount_str = str(amount)
|
||
|
||
# 🔧 修复：使用原子操作 - 分步执行
|
||
# Step 1: 确保余额记录存在（生成UUID）
|
||
await db.execute(
|
||
text("""
|
||
INSERT INTO balances (id, user_id, eu_balance, created_at, updated_at)
|
||
VALUES (gen_random_uuid(), :uid, 0.0, NOW(), NOW())
|
||
ON CONFLICT (user_id) DO NOTHING
|
||
"""),
|
||
{"uid": user_id}
|
||
)
|
||
|
||
# Step 2: 原子更新 - 增加余额
|
||
result = await db.execute(
|
||
text("""
|
||
UPDATE balances
|
||
SET eu_balance = eu_balance + CAST(:amount AS NUMERIC(15, 6)),
|
||
updated_at = NOW()
|
||
WHERE user_id = :uid
|
||
RETURNING
|
||
eu_balance as new_balance,
|
||
eu_balance - CAST(:amount AS NUMERIC(15, 6)) as old_balance
|
||
"""),
|
||
{"uid": user_id, "amount": amount_str}
|
||
)
|
||
|
||
row = result.fetchone()
|
||
new_balance = Decimal(str(row[0])) if row else amount
|
||
old_balance = Decimal(str(row[1])) if row else Decimal("0")
|
||
|
||
# 可选：自动提交
|
||
if auto_commit:
|
||
await db.commit()
|
||
logger.info(
|
||
f"💰 充值成功并已提交: user={user_id[:8]}, "
|
||
f"金额={amount:.6f}, 余额: {old_balance:.6f} → {new_balance:.6f}"
|
||
)
|
||
else:
|
||
logger.info(
|
||
f"💰 充值成功待提交: user={user_id[:8]}, "
|
||
f"金额={amount:.6f}, 余额: {old_balance:.6f} → {new_balance:.6f}"
|
||
)
|
||
|
||
return True, f"成功充值 {amount:.6f} 元，余额: {old_balance:.6f} → {new_balance:.6f}"
|
||
|
||
return True, f"成功充值 {amount} 元，当前余额: {new_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:
|
||
"""
|
||
创建计费记录
|
||
|
||
⚠️ 已废弃 (DEPRECATED)
|
||
此函数使用旧的 BillingRecord 表，已废弃。
|
||
请使用 create_agent_billing_record() 函数替代。
|
||
|
||
Args:
|
||
tenant_id: 租户ID
|
||
agent_id: Agent ID
|
||
agent_name: Agent名称
|
||
duration_seconds: 调用时长（秒）
|
||
db: 数据库会话
|
||
channel_id: 渠道ID（可选）
|
||
|
||
Returns:
|
||
计费记录
|
||
|
||
计费说明：
|
||
根据租户订阅等级自动计算EU单价：
|
||
- 入门级（Starter/free）：$0.015 / EU
|
||
- 专业级（Pro）：$0.02 / EU
|
||
- 企业级（Enterprise）：$0.03 / EU
|
||
"""
|
||
import warnings
|
||
warnings.warn(
|
||
"create_billing_record() 已废弃，请使用 create_agent_billing_record() 替代",
|
||
DeprecationWarning,
|
||
stacklevel=2
|
||
)
|
||
# 获取用户信息（包括订阅等级和channel_id）
|
||
result = await db.execute(
|
||
select(User.channel_id, User.subscription_tier).where(User.id == tenant_id)
|
||
)
|
||
row = result.first()
|
||
|
||
subscription_tier = "pro" # 默认专业级
|
||
if row:
|
||
if not channel_id and row[0]:
|
||
channel_id = str(row[0])
|
||
if row[1]:
|
||
subscription_tier = row[1]
|
||
|
||
# 计算EU和成本（根据订阅等级）
|
||
eu = calculate_eu(duration_seconds)
|
||
cost = calculate_cost(eu, subscription_tier=subscription_tier)
|
||
|
||
# 创建计费记录
|
||
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:
|
||
"""
|
||
计算用户本月消费（从 AgentBillingRecord 和 ModelBillingRecord 合并统计）
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
本月消费金额
|
||
"""
|
||
from sqlalchemy import func
|
||
from models import ModelBillingRecord
|
||
|
||
# 获取本月第一天
|
||
now = datetime.utcnow()
|
||
month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
|
||
|
||
# Agent 计费
|
||
agent_result = await db.execute(
|
||
select(func.sum(AgentBillingRecord.cost))
|
||
.where(AgentBillingRecord.user_id == user_id)
|
||
.where(AgentBillingRecord.start_time >= month_start)
|
||
)
|
||
agent_cost = agent_result.scalar() or 0
|
||
|
||
# 模型调用计费
|
||
model_result = await db.execute(
|
||
select(func.sum(ModelBillingRecord.total_cost))
|
||
.where(ModelBillingRecord.tenant_id == user_id)
|
||
.where(ModelBillingRecord.created_at >= month_start)
|
||
)
|
||
model_cost = model_result.scalar() or 0
|
||
|
||
total = Decimal(str(agent_cost)) + Decimal(str(model_cost))
|
||
return total
|
||
|
||
|
||
async def calculate_channel_commission(
|
||
channel_id: str,
|
||
start_date: datetime,
|
||
end_date: datetime,
|
||
db: AsyncSession
|
||
) -> Tuple[Decimal, Decimal]:
|
||
"""
|
||
计算渠道佣金（从 AgentBillingRecord 和 ModelBillingRecord 合并统计）
|
||
|
||
Args:
|
||
channel_id: 渠道ID
|
||
start_date: 开始日期
|
||
end_date: 结束日期
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
(总收入, 佣金金额)
|
||
"""
|
||
from sqlalchemy import func
|
||
from models import ModelBillingRecord
|
||
|
||
# 查询渠道佣金率
|
||
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)
|
||
|
||
# Agent 计费收入
|
||
agent_result = await db.execute(
|
||
select(func.sum(AgentBillingRecord.cost))
|
||
.where(AgentBillingRecord.channel_id == channel_id)
|
||
.where(AgentBillingRecord.start_time >= start_date)
|
||
.where(AgentBillingRecord.start_time <= end_date)
|
||
)
|
||
agent_revenue = agent_result.scalar() or 0
|
||
|
||
# 模型调用计费收入
|
||
model_result = await db.execute(
|
||
select(func.sum(ModelBillingRecord.total_cost))
|
||
.where(ModelBillingRecord.channel_id == channel_id)
|
||
.where(ModelBillingRecord.created_at >= start_date)
|
||
.where(ModelBillingRecord.created_at <= end_date)
|
||
)
|
||
model_revenue = model_result.scalar() or 0
|
||
|
||
total_revenue = Decimal(str(agent_revenue)) + Decimal(str(model_revenue))
|
||
commission = total_revenue * commission_rate
|
||
|
||
return total_revenue, commission
|
||
|
||
|
||
# ============= Agent 计费配置 =============
|
||
|
||
# Agent 资源单价配置（按资源类型）
|
||
AGENT_RESOURCE_PRICING = {
|
||
"cpu_per_hour": Decimal("0.05"), # CPU 每核每小时 $0.05
|
||
"memory_per_gb_hour": Decimal("0.01"), # 内存每 GB 每小时 $0.01
|
||
}
|
||
|
||
# 平台 Agent 固定价格（按模板）
|
||
PLATFORM_AGENT_PRICING = {
|
||
"gpt-assistant": Decimal("0.10"), # GPT 助手每小时 $0.10
|
||
"code-reviewer": Decimal("0.15"), # 代码审查每小时 $0.15
|
||
"data-analyst": Decimal("0.12"), # 数据分析每小时 $0.12
|
||
"default": Decimal("0.10"), # 默认每小时 $0.10
|
||
}
|
||
|
||
# API 调用计费配置（Agent Manager 回调）
|
||
API_CALL_PRICING = {
|
||
"per_call": Decimal("0.01"), # 每次 API 调用 $0.01（0.01 EU/call）
|
||
}
|
||
|
||
|
||
def get_platform_agent_hourly_price(template_name: str) -> Decimal:
|
||
"""
|
||
获取平台 Agent 每小时价格
|
||
|
||
Args:
|
||
template_name: 模板名称
|
||
|
||
Returns:
|
||
每小时价格（USD）
|
||
"""
|
||
return PLATFORM_AGENT_PRICING.get(template_name, PLATFORM_AGENT_PRICING["default"])
|
||
|
||
|
||
def calculate_agent_cost_by_resources(
|
||
cpu_cores: float,
|
||
memory_gb: float,
|
||
duration_seconds: int
|
||
) -> Decimal:
|
||
"""
|
||
根据资源使用量计算成本（用于自定义 Agent）
|
||
|
||
Args:
|
||
cpu_cores: CPU 核心数
|
||
memory_gb: 内存 GB 数
|
||
duration_seconds: 运行时长（秒）
|
||
|
||
Returns:
|
||
成本金额（USD）
|
||
"""
|
||
hours = Decimal(str(duration_seconds)) / Decimal("3600")
|
||
|
||
cpu_cost = Decimal(str(cpu_cores)) * hours * AGENT_RESOURCE_PRICING["cpu_per_hour"]
|
||
memory_cost = Decimal(str(memory_gb)) * hours * AGENT_RESOURCE_PRICING["memory_per_gb_hour"]
|
||
|
||
return cpu_cost + memory_cost
|
||
|
||
|
||
def calculate_platform_agent_cost(
|
||
template_name: str,
|
||
duration_seconds: int
|
||
) -> Decimal:
|
||
"""
|
||
计算平台 Agent 成本（固定价格）
|
||
|
||
Args:
|
||
template_name: 模板名称
|
||
duration_seconds: 运行时长（秒）
|
||
|
||
Returns:
|
||
成本金额（USD）
|
||
"""
|
||
hours = Decimal(str(duration_seconds)) / Decimal("3600")
|
||
hourly_price = get_platform_agent_hourly_price(template_name)
|
||
|
||
return hours * hourly_price
|
||
|
||
|
||
def calculate_api_call_cost() -> Decimal:
|
||
"""
|
||
计算 API 调用成本（固定费用）
|
||
|
||
用于 Agent Manager 回调的 API 调用计费。
|
||
每次调用固定收费 0.001 EU（= $0.001），与运行时间无关。
|
||
|
||
Returns:
|
||
成本金额（USD）：0.001
|
||
"""
|
||
return API_CALL_PRICING["per_call"]
|
||
|
||
|
||
# ============= Agent 计费记录创建 =============
|
||
|
||
async def create_agent_billing_record(
|
||
user_id: str,
|
||
agent_name: str,
|
||
agent_type: str,
|
||
is_platform_agent: bool,
|
||
duration_seconds: int,
|
||
db: AsyncSession,
|
||
channel_id: Optional[str] = None,
|
||
cpu_used: Optional[str] = None,
|
||
memory_used: Optional[str] = None,
|
||
tools_used: Optional[list] = None,
|
||
request_id: Optional[str] = None,
|
||
) -> "AgentBillingRecord":
|
||
"""
|
||
创建 Agent 计费记录
|
||
|
||
Args:
|
||
user_id: 用户 ID
|
||
agent_name: Agent 名称
|
||
agent_type: Agent 类型/模板名称
|
||
is_platform_agent: 是否为平台 Agent
|
||
duration_seconds: 运行时长（秒）
|
||
db: 数据库会话
|
||
channel_id: 渠道 ID（可选）
|
||
cpu_used: CPU 使用量（K8s 格式，如 "100m"）
|
||
memory_used: 内存使用量（K8s 格式，如 "128Mi"）
|
||
tools_used: 使用的工具列表（可选）
|
||
request_id: 请求 ID（可选）
|
||
|
||
Returns:
|
||
计费记录
|
||
"""
|
||
from models import AgentBillingRecord
|
||
|
||
# 计算成本
|
||
if is_platform_agent:
|
||
cost = calculate_platform_agent_cost(agent_type, duration_seconds)
|
||
else:
|
||
# 解析资源量
|
||
cpu_cores = _parse_cpu_to_cores(cpu_used) if cpu_used else 0.1
|
||
memory_gb = _parse_memory_to_gb(memory_used) if memory_used else 0.125
|
||
cost = calculate_agent_cost_by_resources(cpu_cores, memory_gb, duration_seconds)
|
||
|
||
# EU = Cost（1 EU = 1 美元）
|
||
eu = float(cost)
|
||
|
||
# 创建记录
|
||
record = AgentBillingRecord(
|
||
user_id=user_id,
|
||
channel_id=channel_id,
|
||
agent_name=agent_name,
|
||
agent_type=agent_type,
|
||
is_platform_agent=is_platform_agent,
|
||
period_start=datetime.utcnow(),
|
||
duration_seconds=duration_seconds,
|
||
eu_consumed=eu, # EU = Cost（美元）
|
||
cpu_used=cpu_used,
|
||
memory_used=memory_used,
|
||
cost=cost,
|
||
tools_used=tools_used,
|
||
request_id=request_id,
|
||
)
|
||
|
||
db.add(record)
|
||
|
||
# 扣除用户余额
|
||
success, message = await deduct_balance(user_id, cost, db, f"Agent 使用: {agent_name}")
|
||
|
||
await db.commit()
|
||
await db.refresh(record)
|
||
|
||
return record
|
||
|
||
|
||
async def finalize_agent_billing(
|
||
agent_name: str,
|
||
user_id: str,
|
||
db: AsyncSession
|
||
) -> Optional["AgentBillingRecord"]:
|
||
"""
|
||
结算 Agent 计费（Agent 停止时调用）
|
||
|
||
Args:
|
||
agent_name: Agent 名称
|
||
user_id: 用户 ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
更新后的计费记录
|
||
"""
|
||
from models import AgentBillingRecord
|
||
from sqlalchemy import and_
|
||
|
||
# 查找未结束的计费记录
|
||
result = await db.execute(
|
||
select(AgentBillingRecord).where(
|
||
and_(
|
||
AgentBillingRecord.agent_name == agent_name,
|
||
AgentBillingRecord.user_id == user_id,
|
||
AgentBillingRecord.period_end == None
|
||
)
|
||
)
|
||
)
|
||
record = result.scalar_one_or_none()
|
||
|
||
if not record:
|
||
return None
|
||
|
||
# 计算实际运行时长
|
||
record.period_end = datetime.utcnow()
|
||
duration = (record.period_end - record.period_start).total_seconds()
|
||
record.duration_seconds = int(duration)
|
||
|
||
# 重新计算成本
|
||
if record.is_platform_agent:
|
||
record.cost = calculate_platform_agent_cost(record.agent_type, int(duration))
|
||
else:
|
||
cpu_cores = _parse_cpu_to_cores(record.cpu_used) if record.cpu_used else 0.1
|
||
memory_gb = _parse_memory_to_gb(record.memory_used) if record.memory_used else 0.125
|
||
record.cost = calculate_agent_cost_by_resources(cpu_cores, memory_gb, int(duration))
|
||
|
||
# EU = Cost（1 EU = 1 美元）
|
||
record.eu_consumed = float(record.cost)
|
||
|
||
await db.commit()
|
||
await db.refresh(record)
|
||
|
||
return record
|
||
|
||
|
||
# ============= Agent 计费统计 =============
|
||
|
||
async def get_agent_billing_stats(
|
||
user_id: str,
|
||
start_date: datetime,
|
||
end_date: datetime,
|
||
db: AsyncSession
|
||
) -> dict:
|
||
"""
|
||
获取 Agent 计费统计
|
||
|
||
Args:
|
||
user_id: 用户 ID
|
||
start_date: 开始日期
|
||
end_date: 结束日期
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
统计数据
|
||
"""
|
||
from models import AgentBillingRecord
|
||
from sqlalchemy import func, and_
|
||
|
||
# 总体统计
|
||
result = await db.execute(
|
||
select(
|
||
func.sum(AgentBillingRecord.cost).label("total_cost"),
|
||
func.sum(AgentBillingRecord.duration_seconds).label("total_duration"),
|
||
func.count(AgentBillingRecord.id).label("total_count")
|
||
).where(
|
||
and_(
|
||
AgentBillingRecord.user_id == user_id,
|
||
AgentBillingRecord.period_start >= start_date,
|
||
AgentBillingRecord.period_start <= end_date
|
||
)
|
||
)
|
||
)
|
||
row = result.first()
|
||
|
||
total_cost = float(row.total_cost) if row.total_cost else 0
|
||
total_duration = int(row.total_duration) if row.total_duration else 0
|
||
total_count = int(row.total_count) if row.total_count else 0
|
||
|
||
# 按 Agent 类型统计
|
||
result = await db.execute(
|
||
select(
|
||
AgentBillingRecord.is_platform_agent,
|
||
func.sum(AgentBillingRecord.cost).label("cost")
|
||
).where(
|
||
and_(
|
||
AgentBillingRecord.user_id == user_id,
|
||
AgentBillingRecord.period_start >= start_date,
|
||
AgentBillingRecord.period_start <= end_date
|
||
)
|
||
).group_by(AgentBillingRecord.is_platform_agent)
|
||
)
|
||
|
||
by_agent_type = {}
|
||
for row in result.all():
|
||
agent_type = "platform" if row.is_platform_agent else "custom"
|
||
by_agent_type[agent_type] = float(row.cost) if row.cost else 0
|
||
|
||
# 按模板统计
|
||
result = await db.execute(
|
||
select(
|
||
AgentBillingRecord.agent_type,
|
||
func.sum(AgentBillingRecord.cost).label("cost")
|
||
).where(
|
||
and_(
|
||
AgentBillingRecord.user_id == user_id,
|
||
AgentBillingRecord.period_start >= start_date,
|
||
AgentBillingRecord.period_start <= end_date
|
||
)
|
||
).group_by(AgentBillingRecord.agent_type)
|
||
)
|
||
|
||
by_template = {}
|
||
for row in result.all():
|
||
by_template[row.agent_type] = float(row.cost) if row.cost else 0
|
||
|
||
return {
|
||
"totalCost": total_cost,
|
||
"totalDurationSeconds": total_duration,
|
||
"totalRequests": total_count,
|
||
"byAgentType": by_agent_type,
|
||
"byTemplate": by_template,
|
||
}
|
||
|
||
|
||
async def get_channel_agent_billing_stats(
|
||
channel_id: str,
|
||
start_date: datetime,
|
||
end_date: datetime,
|
||
db: AsyncSession
|
||
) -> dict:
|
||
"""
|
||
获取渠道 Agent 计费统计
|
||
|
||
Args:
|
||
channel_id: 渠道 ID
|
||
start_date: 开始日期
|
||
end_date: 结束日期
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
统计数据
|
||
"""
|
||
from models import AgentBillingRecord
|
||
from sqlalchemy import func, and_
|
||
|
||
# 总体统计
|
||
result = await db.execute(
|
||
select(
|
||
func.sum(AgentBillingRecord.cost).label("total_cost"),
|
||
func.sum(AgentBillingRecord.duration_seconds).label("total_duration"),
|
||
func.count(AgentBillingRecord.id).label("total_count")
|
||
).where(
|
||
and_(
|
||
AgentBillingRecord.channel_id == channel_id,
|
||
AgentBillingRecord.period_start >= start_date,
|
||
AgentBillingRecord.period_start <= end_date
|
||
)
|
||
)
|
||
)
|
||
row = result.first()
|
||
|
||
total_cost = float(row.total_cost) if row.total_cost else 0
|
||
total_duration = int(row.total_duration) if row.total_duration else 0
|
||
total_count = int(row.total_count) if row.total_count else 0
|
||
|
||
# 按用户统计
|
||
result = await db.execute(
|
||
select(
|
||
AgentBillingRecord.user_id,
|
||
func.sum(AgentBillingRecord.cost).label("cost"),
|
||
func.count(AgentBillingRecord.id).label("count")
|
||
).where(
|
||
and_(
|
||
AgentBillingRecord.channel_id == channel_id,
|
||
AgentBillingRecord.period_start >= start_date,
|
||
AgentBillingRecord.period_start <= end_date
|
||
)
|
||
).group_by(AgentBillingRecord.user_id)
|
||
)
|
||
|
||
by_user = []
|
||
for row in result.all():
|
||
by_user.append({
|
||
"userId": row.user_id,
|
||
"cost": float(row.cost) if row.cost else 0,
|
||
"count": int(row.count) if row.count else 0,
|
||
})
|
||
|
||
return {
|
||
"totalCost": total_cost,
|
||
"totalDurationSeconds": total_duration,
|
||
"totalRequests": total_count,
|
||
"byUser": by_user,
|
||
}
|
||
|
||
|
||
# ============= 资源解析辅助函数 =============
|
||
|
||
def _parse_cpu_to_cores(cpu_str: str) -> float:
|
||
"""
|
||
解析 CPU 字符串为核心数
|
||
|
||
支持格式：
|
||
- "100m" -> 0.1 核
|
||
- "1" -> 1 核
|
||
- "1.5" -> 1.5 核
|
||
"""
|
||
if not cpu_str:
|
||
return 0.0
|
||
|
||
cpu_str = cpu_str.strip().lower()
|
||
|
||
if cpu_str.endswith("m"):
|
||
return float(cpu_str[:-1]) / 1000
|
||
else:
|
||
return float(cpu_str)
|
||
|
||
|
||
def _parse_memory_to_gb(memory_str: str) -> float:
|
||
"""
|
||
解析内存字符串为 GB
|
||
|
||
支持格式：
|
||
- "128Mi" -> 0.125 GB
|
||
- "1Gi" -> 1 GB
|
||
- "512M" -> 0.5 GB
|
||
- "2G" -> 2 GB
|
||
"""
|
||
if not memory_str:
|
||
return 0.0
|
||
|
||
memory_str = memory_str.strip()
|
||
|
||
# 处理 Kubernetes 格式
|
||
if memory_str.endswith("Gi"):
|
||
return float(memory_str[:-2])
|
||
elif memory_str.endswith("Mi"):
|
||
return float(memory_str[:-2]) / 1024
|
||
elif memory_str.endswith("Ki"):
|
||
return float(memory_str[:-2]) / (1024 * 1024)
|
||
# 处理简化格式
|
||
elif memory_str.endswith("G"):
|
||
return float(memory_str[:-1])
|
||
elif memory_str.endswith("M"):
|
||
return float(memory_str[:-1]) / 1024
|
||
elif memory_str.endswith("K"):
|
||
return float(memory_str[:-1]) / (1024 * 1024)
|
||
else:
|
||
# 假设是字节
|
||
return float(memory_str) / (1024 * 1024 * 1024)
|
||
|