forked from xiaohei/taiji-AI-PAD
996 lines
28 KiB
Python
996 lines
28 KiB
Python
"""
|
||
计费与资源管理逻辑
|
||
|
||
EU（执行单元）计算说明：
|
||
=======================
|
||
|
||
系统中有两种EU计算方式：
|
||
|
||
1. Agent 运行时长计费（本模块）
|
||
- 计算公式：EU = ceil(duration_seconds / 10)
|
||
- 即：1 EU = 10秒运行时间，不足10秒按1 EU计算
|
||
- 适用于：平台Agent、自定义Agent的Pod运行时长
|
||
- 周期性统计：每小时更新一次计费记录（periodic_billing.py）
|
||
|
||
2. 模型 Token 计费（billing_webhook.py）
|
||
- 计算公式：EU = total_tokens * MODEL_EU_RATE[model_name]
|
||
- 不同模型有不同的转换率：
|
||
- gpt-4: 0.0001 EU/token (1000 tokens = 0.1 EU)
|
||
- gpt-3.5-turbo: 0.00005 EU/token (1000 tokens = 0.05 EU)
|
||
- claude: 0.0001 EU/token
|
||
- 默认: 0.0001 EU/token
|
||
- 适用于：LiteLLM模型调用的Token消耗
|
||
|
||
数据存储：
|
||
- Agent计费记录：agent_billing_records 表
|
||
- 模型计费记录：model_billing_records 表
|
||
- 用户余额：balances 表（Balance模型）
|
||
|
||
注意：User.eu_balance 和 User.balance 字段已废弃，请使用 Balance 表
|
||
"""
|
||
|
||
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, Balance, AgentBillingRecord
|
||
|
||
|
||
# ============= EU定价配置 =============
|
||
|
||
# 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单价
|
||
|
||
Args:
|
||
subscription_tier: 订阅等级（free/starter/pro/enterprise）
|
||
|
||
Returns:
|
||
EU单价（USD）
|
||
|
||
定价说明：
|
||
- 入门级（Starter/free）：$0.015 / EU - 拉新、试用、轻 Agent
|
||
- 专业级（Pro）：$0.02 / EU - 主力商业用户
|
||
- 企业级（Enterprise）：$0.03 / EU - 高复杂度 / 高 SLA
|
||
"""
|
||
tier = subscription_tier.lower() if subscription_tier else "pro"
|
||
return EU_PRICING.get(tier, DEFAULT_EU_PRICE)
|
||
|
||
|
||
# ============= 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 = None, subscription_tier: str = None) -> Decimal:
|
||
"""
|
||
计算成本
|
||
|
||
Args:
|
||
eu: EU数量
|
||
unit_price: 单价（可直接指定，优先级最高）
|
||
subscription_tier: 订阅等级（用于自动获取单价）
|
||
|
||
Returns:
|
||
成本金额（USD）
|
||
|
||
定价说明：
|
||
- 入门级（Starter/free）：$0.015 / EU
|
||
- 专业级（Pro）：$0.02 / EU
|
||
- 企业级（Enterprise）：$0.03 / EU
|
||
"""
|
||
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 = False
|
||
) -> Tuple[bool, str]:
|
||
"""
|
||
扣除余额（优先扣除账户余额，不足时使用授信额度）
|
||
|
||
使用行锁保护并发扣款操作，防止超扣。
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
amount: 扣除金额
|
||
db: 数据库会话
|
||
description: 描述
|
||
auto_commit: 是否自动提交（默认False，由调用者管理事务）
|
||
|
||
Returns:
|
||
(是否成功, 消息)
|
||
|
||
Note:
|
||
默认不会 commit，由调用者统一管理事务。
|
||
如需独立提交，请设置 auto_commit=True。
|
||
"""
|
||
# 使用 FOR UPDATE 锁定余额行，防止并发扣款
|
||
balance_result = await db.execute(
|
||
select(Balance)
|
||
.where(Balance.user_id == user_id)
|
||
.with_for_update()
|
||
)
|
||
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 False, "用户不存在"
|
||
|
||
if balance_obj is None:
|
||
# 如果余额记录不存在，创建一个新的（初始余额为0）
|
||
balance_obj = Balance(user_id=user_id, eu_balance=0.0)
|
||
db.add(balance_obj)
|
||
await db.flush() # 确保记录创建后再继续
|
||
|
||
balance = Decimal(str(balance_obj.eu_balance))
|
||
credit_limit = Decimal(str(user.credit_limit))
|
||
available = balance + credit_limit
|
||
|
||
if available < amount:
|
||
return False, f"余额不足，当前可用额度: {available}, 需要: {amount}"
|
||
|
||
# 优先扣除账户余额，允许透支到授信额度
|
||
new_balance = balance - amount
|
||
balance_obj.eu_balance = float(new_balance)
|
||
|
||
# 可选：自动提交
|
||
if auto_commit:
|
||
await db.commit()
|
||
|
||
return True, f"成功扣除 {amount} 元"
|
||
|
||
|
||
async def add_balance(
|
||
user_id: str,
|
||
amount: Decimal,
|
||
db: AsyncSession,
|
||
description: str = "充值",
|
||
auto_commit: bool = False
|
||
) -> Tuple[bool, str]:
|
||
"""
|
||
增加余额
|
||
|
||
使用行锁保护并发充值操作，确保余额准确。
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
amount: 充值金额
|
||
db: 数据库会话
|
||
description: 描述
|
||
auto_commit: 是否自动提交（默认False，由调用者管理事务）
|
||
|
||
Returns:
|
||
(是否成功, 消息)
|
||
|
||
Note:
|
||
默认不会 commit，由调用者统一管理事务。
|
||
如需独立提交，请设置 auto_commit=True。
|
||
"""
|
||
# 检查用户是否存在
|
||
user_result = await db.execute(
|
||
select(User).where(User.id == user_id)
|
||
)
|
||
user = user_result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
return False, "用户不存在"
|
||
|
||
# 使用 FOR UPDATE 锁定余额行，防止并发更新
|
||
balance_result = await db.execute(
|
||
select(Balance)
|
||
.where(Balance.user_id == user_id)
|
||
.with_for_update()
|
||
)
|
||
balance_obj = balance_result.scalar_one_or_none()
|
||
|
||
if balance_obj is None:
|
||
# 如果余额记录不存在，创建一个新的
|
||
balance_obj = Balance(user_id=user_id, eu_balance=0.0)
|
||
db.add(balance_obj)
|
||
await db.flush() # 确保记录创建后再继续
|
||
|
||
old_balance = Decimal(str(balance_obj.eu_balance))
|
||
new_balance = old_balance + amount
|
||
balance_obj.eu_balance = float(new_balance)
|
||
|
||
# 可选：自动提交
|
||
if auto_commit:
|
||
await db.commit()
|
||
|
||
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
|
||
}
|
||
|
||
|
||
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
|
||
|
||
|
||
# ============= 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
|
||
eu = calculate_eu(duration_seconds)
|
||
|
||
# 创建记录
|
||
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,
|
||
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)
|
||
record.eu_consumed = calculate_eu(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))
|
||
|
||
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)
|
||
|