Pine
479 lines
17 KiB
Python
479 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""算力多级折扣价格计算与余额扣费服务。
|
||
|
||
折扣体系（从高到低优先级）：
|
||
1. 载体端账号折扣（tenant_user_discounts）- 载体给特定账号的折扣
|
||
2. 载体端企业折扣（park_companies.compute_discount）- 载体给下属企业的折扣
|
||
3. 运营端载体折扣（tenant_discounts）- 运营给载体的折扣（载体进货价）
|
||
4. 运营端标准价 - 未绑定载体的用户使用
|
||
|
||
折扣表示：0-100，0=无折扣，100=全免费；实际价格=标准价×(1-discount/100)
|
||
|
||
扣费优先级：
|
||
1. 企业分配的余额（优先选择折扣最低的企业）
|
||
2. 个人余额
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import uuid
|
||
from datetime import datetime
|
||
from typing import Optional
|
||
|
||
from sqlalchemy import select, and_
|
||
|
||
from ..infrastructure.models import (
|
||
ParkTenant, ParkCompany, CompanyMember, User,
|
||
TenantDiscount, TenantUserDiscount,
|
||
ComputeRecharge, ComputeBalanceAllocation, ComputeUsageRecord,
|
||
)
|
||
from ..infrastructure.repositories import Database
|
||
|
||
MICRO_PER_FEN = 10000 # 1分 = 10000 micro（1元=100分=1,000,000 micro）
|
||
|
||
|
||
# ── 标准价配置（可从数据库或配置文件读取） ──
|
||
STANDARD_PRICES = {
|
||
# model: 每1000token价格（分）
|
||
"default": 2, # 默认模型
|
||
"gpt-4": 10, # 高价模型
|
||
"gpt-3.5-turbo": 2, # 基础模型
|
||
}
|
||
|
||
|
||
def get_standard_price(model: str, token_count: int) -> int:
|
||
"""获取标准价（分）。"""
|
||
per_1k = STANDARD_PRICES.get(model, STANDARD_PRICES["default"])
|
||
return max(1, int(per_1k * token_count / 1000))
|
||
|
||
|
||
def now_str() -> str:
|
||
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
|
||
|
||
def new_id() -> str:
|
||
return uuid.uuid4().hex[:24]
|
||
|
||
|
||
# ── 价格计算 ──
|
||
|
||
async def calculate_compute_price(db: Database, user_id: str, model: str, token_count: int) -> dict:
|
||
"""
|
||
计算用户使用算力的实际价格。
|
||
|
||
返回：
|
||
{
|
||
"standard_price": int, # 标准价（分）
|
||
"tenant_discount": int, # 载体折扣（0-100）
|
||
"discount": int, # 实际折扣（0-100）
|
||
"discount_source": str, # standard/tenant/company/tenant_user
|
||
"actual_price": int, # 实际价格（分）
|
||
"tenant_id": str, # 载体ID（如有）
|
||
"company_id": str, # 企业ID（如有）
|
||
}
|
||
"""
|
||
standard_price = get_standard_price(model, token_count)
|
||
result = {
|
||
"standard_price": standard_price,
|
||
"tenant_discount": 0,
|
||
"discount": 0,
|
||
"discount_source": "standard",
|
||
"actual_price": standard_price,
|
||
"tenant_id": "",
|
||
"company_id": "",
|
||
}
|
||
|
||
# 1. 检查用户是否绑定载体（通过企业成员关系或直接绑定）
|
||
tenant_id = await _get_user_tenant_id(db, user_id)
|
||
if not tenant_id:
|
||
return result
|
||
|
||
result["tenant_id"] = tenant_id
|
||
|
||
# 2. 获取运营端给载体的折扣（载体进货价）
|
||
tenant_discount = await _get_tenant_discount(db, tenant_id)
|
||
result["tenant_discount"] = tenant_discount
|
||
|
||
# 载体基础价 = 标准价 × (1 - 载体折扣/100)
|
||
tenant_base_price = int(standard_price * (1 - tenant_discount / 100))
|
||
|
||
# 3. 检查载体端是否给该账号设置了单独折扣（优先级最高）
|
||
user_discount = await _get_tenant_user_discount(db, tenant_id, user_id)
|
||
if user_discount > 0:
|
||
actual_price = int(tenant_base_price * (1 - user_discount / 100))
|
||
result.update({
|
||
"discount": user_discount,
|
||
"discount_source": "tenant_user",
|
||
"actual_price": actual_price,
|
||
})
|
||
return result
|
||
|
||
# 4. 检查用户是否属于载体下属企业
|
||
company = await _get_user_company(db, user_id, tenant_id)
|
||
if company:
|
||
company_discount = getattr(company, "compute_discount", 0) or 0
|
||
actual_price = int(tenant_base_price * (1 - company_discount / 100))
|
||
result.update({
|
||
"discount": company_discount,
|
||
"discount_source": "company",
|
||
"actual_price": actual_price,
|
||
"company_id": company.id,
|
||
})
|
||
return result
|
||
|
||
# 5. 使用载体默认账号折扣
|
||
tenant = await db.session.get(ParkTenant, tenant_id)
|
||
default_discount = getattr(tenant, "default_user_discount", 0) or 0 if tenant else 0
|
||
if default_discount > 0:
|
||
actual_price = int(tenant_base_price * (1 - default_discount / 100))
|
||
result.update({
|
||
"discount": default_discount,
|
||
"discount_source": "tenant_default",
|
||
"actual_price": actual_price,
|
||
})
|
||
return result
|
||
|
||
# 6. 仅使用载体折扣
|
||
result.update({
|
||
"discount": tenant_discount,
|
||
"discount_source": "tenant",
|
||
"actual_price": tenant_base_price,
|
||
})
|
||
return result
|
||
|
||
|
||
async def _get_user_tenant_id(db: Database, user_id: str) -> Optional[str]:
|
||
"""获取用户绑定的载体ID（通过企业成员关系）。"""
|
||
# 通过企业成员关系查找
|
||
result = await db.session.execute(
|
||
select(CompanyMember).where(
|
||
and_(CompanyMember.user_id == user_id, CompanyMember.status == "active")
|
||
)
|
||
)
|
||
members = result.scalars().all()
|
||
for m in members:
|
||
company = await db.session.get(ParkCompany, m.company_id)
|
||
if company and company.tenant_id:
|
||
return company.tenant_id
|
||
return None
|
||
|
||
|
||
async def _get_tenant_discount(db: Database, tenant_id: str) -> int:
|
||
"""获取运营端给载体的折扣。"""
|
||
result = await db.session.execute(
|
||
select(TenantDiscount).where(TenantDiscount.tenant_id == tenant_id)
|
||
)
|
||
td = result.scalars().first()
|
||
if td:
|
||
# 检查有效期
|
||
now = now_str()
|
||
if td.effective_date and td.effective_date > now:
|
||
return 0
|
||
if td.expire_date and td.expire_date < now:
|
||
return 0
|
||
return td.discount or 0
|
||
return 0
|
||
|
||
|
||
async def _get_tenant_user_discount(db: Database, tenant_id: str, user_id: str) -> int:
|
||
"""获取载体端给特定账号的折扣。"""
|
||
result = await db.session.execute(
|
||
select(TenantUserDiscount).where(
|
||
and_(TenantUserDiscount.tenant_id == tenant_id, TenantUserDiscount.user_id == user_id)
|
||
)
|
||
)
|
||
tud = result.scalars().first()
|
||
return tud.discount if tud else 0
|
||
|
||
|
||
async def _get_user_company(db: Database, user_id: str, tenant_id: str) -> Optional[ParkCompany]:
|
||
"""获取用户在指定载体下的企业。"""
|
||
result = await db.session.execute(
|
||
select(CompanyMember).where(
|
||
and_(CompanyMember.user_id == user_id, CompanyMember.status == "active")
|
||
)
|
||
)
|
||
members = result.scalars().all()
|
||
for m in members:
|
||
company = await db.session.get(ParkCompany, m.company_id)
|
||
if company and company.tenant_id == tenant_id:
|
||
return company
|
||
return None
|
||
|
||
|
||
# ── 余额扣费 ──
|
||
|
||
async def deduct_compute_balance(
|
||
db: Database, user_id: str, model: str, token_count: int
|
||
) -> dict:
|
||
"""
|
||
算力扣费。
|
||
|
||
扣费优先级：
|
||
1. 企业分配的余额（优先选择折扣最低的企业）
|
||
2. 个人余额
|
||
|
||
返回：
|
||
{
|
||
"success": bool,
|
||
"amount": int, # 实际扣费（分）
|
||
"source": str, # company/personal
|
||
"company_id": str, # 从哪家企业扣费（如有）
|
||
"price_info": dict, # 价格计算结果
|
||
"reason": str, # 失败原因
|
||
}
|
||
"""
|
||
price_info = await calculate_compute_price(db, user_id, model, token_count)
|
||
amount = price_info["actual_price"]
|
||
|
||
# 1. 尝试从企业余额扣费
|
||
company_balances = await _get_user_company_balances(db, user_id)
|
||
# 按折扣从低到高排序（折扣越低，价格越低，优先使用）
|
||
company_balances.sort(key=lambda x: x["discount"])
|
||
|
||
for cb in company_balances:
|
||
if cb["balance"] >= amount:
|
||
success = await _deduct_from_company_balance(db, cb["company_id"], user_id, amount)
|
||
if success:
|
||
await _record_usage(db, user_id, model, token_count, price_info, amount, "company", cb["company_id"])
|
||
return {
|
||
"success": True,
|
||
"amount": amount,
|
||
"source": "company",
|
||
"company_id": cb["company_id"],
|
||
"price_info": price_info,
|
||
}
|
||
|
||
# 2. 尝试从个人余额扣费
|
||
user_balance = await _get_user_balance(db, user_id)
|
||
if user_balance >= amount:
|
||
success = await _deduct_from_user_balance(db, user_id, amount)
|
||
if success:
|
||
await _record_usage(db, user_id, model, token_count, price_info, amount, "personal", "")
|
||
return {
|
||
"success": True,
|
||
"amount": amount,
|
||
"source": "personal",
|
||
"price_info": price_info,
|
||
}
|
||
|
||
# 3. 余额不足
|
||
return {
|
||
"success": False,
|
||
"reason": "insufficient_balance",
|
||
"required": amount,
|
||
"available": user_balance,
|
||
"price_info": price_info,
|
||
}
|
||
|
||
|
||
async def _get_user_company_balances(db: Database, user_id: str) -> list:
|
||
"""获取用户所有企业的余额和折扣。"""
|
||
result = await db.session.execute(
|
||
select(CompanyMember).where(
|
||
and_(CompanyMember.user_id == user_id, CompanyMember.status == "active")
|
||
)
|
||
)
|
||
members = result.scalars().all()
|
||
balances = []
|
||
for m in members:
|
||
company = await db.session.get(ParkCompany, m.company_id)
|
||
if company:
|
||
balances.append({
|
||
"company_id": m.company_id,
|
||
"company_name": company.name,
|
||
"balance": m.compute_balance or 0,
|
||
"discount": company.compute_discount or 0,
|
||
})
|
||
return balances
|
||
|
||
|
||
async def _get_user_balance(db: Database, user_id: str) -> int:
|
||
"""获取用户个人余额。"""
|
||
# TODO: 从用户表或余额表获取个人余额
|
||
# 暂时返回0，表示需要从企业余额扣费
|
||
return 0
|
||
|
||
|
||
async def _deduct_from_company_balance(db: Database, company_id: str, user_id: str, amount: int) -> bool:
|
||
"""从企业分配的成员余额扣费。"""
|
||
result = await db.session.execute(
|
||
select(CompanyMember).where(
|
||
and_(CompanyMember.company_id == company_id, CompanyMember.user_id == user_id)
|
||
)
|
||
)
|
||
member = result.scalars().first()
|
||
if not member or (member.compute_balance or 0) < amount:
|
||
return False
|
||
member.compute_balance = (member.compute_balance or 0) - amount
|
||
member.compute_balance_used = (member.compute_balance_used or 0) + amount
|
||
await db.session.commit()
|
||
return True
|
||
|
||
|
||
async def _deduct_from_user_balance(db: Database, user_id: str, amount: int) -> bool:
|
||
"""从用户个人余额扣费。"""
|
||
# TODO: 实现个人余额扣费
|
||
return False
|
||
|
||
|
||
async def _record_usage(
|
||
db: Database, user_id: str, model: str, token_count: int,
|
||
price_info: dict, amount: int, balance_source: str, company_id: str
|
||
):
|
||
"""记录算力使用。"""
|
||
record = ComputeUsageRecord(
|
||
id=new_id(),
|
||
user_id=user_id,
|
||
company_id=company_id,
|
||
model=model,
|
||
token_count=token_count,
|
||
standard_price=price_info["standard_price"],
|
||
discount=price_info["discount"],
|
||
discount_source=price_info["discount_source"],
|
||
actual_amount=amount,
|
||
balance_source=balance_source,
|
||
created_at=now_str(),
|
||
)
|
||
db.session.add(record)
|
||
await db.session.commit()
|
||
|
||
|
||
# ── 企业余额管理 ──
|
||
|
||
async def get_company_balance(db: Database, company_id: str) -> dict:
|
||
"""获取企业余额总览。"""
|
||
company = await db.session.get(ParkCompany, company_id)
|
||
if not company:
|
||
return {"error": "企业不存在"}
|
||
return {
|
||
"company_id": company.id,
|
||
"company_name": company.name,
|
||
"total_balance": company.compute_balance or 0,
|
||
"allocated_balance": company.compute_balance_used or 0,
|
||
"available_balance": (company.compute_balance or 0) - (company.compute_balance_used or 0),
|
||
}
|
||
|
||
|
||
async def allocate_balance_to_member(
|
||
db: Database, company_id: str, from_user_id: str, to_user_id: str, amount: int, reason: str = ""
|
||
) -> dict:
|
||
"""给成员分配余额。
|
||
企业扣（compute_balance_used 增加），成员个人算力增（company_members.compute_balance + compute-engine quota 同步）。
|
||
"""
|
||
company = await db.session.get(ParkCompany, company_id)
|
||
if not company:
|
||
return {"success": False, "reason": "企业不存在"}
|
||
|
||
available = (company.compute_balance or 0) - (company.compute_balance_used or 0)
|
||
if available < amount:
|
||
return {"success": False, "reason": "企业可用余额不足", "available": available}
|
||
|
||
# 查找成员关系
|
||
result = await db.session.execute(
|
||
select(CompanyMember).where(
|
||
and_(CompanyMember.company_id == company_id, CompanyMember.user_id == to_user_id)
|
||
)
|
||
)
|
||
member = result.scalars().first()
|
||
if not member:
|
||
return {"success": False, "reason": "用户不是该企业成员"}
|
||
|
||
# 获取成员 username，用于 compute-engine 同步
|
||
to_user = await db.session.get(User, to_user_id)
|
||
if not to_user:
|
||
return {"success": False, "reason": "成员用户不存在"}
|
||
|
||
# 先同步 compute-engine（增加成员 quota），成功后再提交数据库
|
||
from ..pay.service import _resolve_engine_user_id
|
||
from ..services import compute_client
|
||
total_micro = amount * MICRO_PER_FEN
|
||
try:
|
||
engine_user_id = await _resolve_engine_user_id(to_user.username or "")
|
||
if not engine_user_id:
|
||
return {"success": False, "reason": "成员算力账号未就绪"}
|
||
await compute_client.adjust_user_quota(engine_user_id, total_micro, "add")
|
||
await compute_client.sync_user_mirror(db, engine_user_id)
|
||
except Exception as exc:
|
||
return {"success": False, "reason": f"算力同步失败: {exc}"}
|
||
|
||
# 分配余额（数据库记录）
|
||
company.compute_balance_used = (company.compute_balance_used or 0) + amount
|
||
member.compute_balance = (member.compute_balance or 0) + amount
|
||
|
||
# 记录分配
|
||
allocation = ComputeBalanceAllocation(
|
||
id=new_id(),
|
||
company_id=company_id,
|
||
from_user_id=from_user_id,
|
||
to_user_id=to_user_id,
|
||
amount=amount,
|
||
type="allocate",
|
||
reason=reason,
|
||
created_at=now_str(),
|
||
)
|
||
db.session.add(allocation)
|
||
await db.session.commit()
|
||
|
||
return {"success": True, "allocation_id": allocation.id}
|
||
|
||
|
||
async def reclaim_balance_from_member(
|
||
db: Database, company_id: str, from_user_id: str, to_user_id: str, amount: int, reason: str = ""
|
||
) -> dict:
|
||
"""回收成员余额。
|
||
成员个人算力扣（company_members.compute_balance + compute-engine quota 同步），企业增（compute_balance_used 减少）。
|
||
"""
|
||
result = await db.session.execute(
|
||
select(CompanyMember).where(
|
||
and_(CompanyMember.company_id == company_id, CompanyMember.user_id == to_user_id)
|
||
)
|
||
)
|
||
member = result.scalars().first()
|
||
if not member:
|
||
return {"success": False, "reason": "用户不是该企业成员"}
|
||
|
||
if (member.compute_balance or 0) < amount:
|
||
return {"success": False, "reason": "成员余额不足", "available": member.compute_balance or 0}
|
||
|
||
company = await db.session.get(ParkCompany, company_id)
|
||
if not company:
|
||
return {"success": False, "reason": "企业不存在"}
|
||
|
||
# 获取成员 username，用于 compute-engine 同步
|
||
to_user = await db.session.get(User, to_user_id)
|
||
if not to_user:
|
||
return {"success": False, "reason": "成员用户不存在"}
|
||
|
||
# 先同步 compute-engine（扣减成员 quota），成功后再提交数据库
|
||
from ..pay.service import _resolve_engine_user_id
|
||
from ..services import compute_client
|
||
total_micro = amount * MICRO_PER_FEN
|
||
try:
|
||
engine_user_id = await _resolve_engine_user_id(to_user.username or "")
|
||
if not engine_user_id:
|
||
return {"success": False, "reason": "成员算力账号未就绪"}
|
||
await compute_client.adjust_user_quota(engine_user_id, total_micro, "subtract")
|
||
await compute_client.sync_user_mirror(db, engine_user_id)
|
||
except Exception as exc:
|
||
return {"success": False, "reason": f"算力同步失败: {exc}"}
|
||
|
||
# 回收余额（数据库记录）
|
||
member.compute_balance = (member.compute_balance or 0) - amount
|
||
company.compute_balance_used = max(0, (company.compute_balance_used or 0) - amount)
|
||
|
||
# 记录回收
|
||
allocation = ComputeBalanceAllocation(
|
||
id=new_id(),
|
||
company_id=company_id,
|
||
from_user_id=from_user_id,
|
||
to_user_id=to_user_id,
|
||
amount=amount,
|
||
type="reclaim",
|
||
reason=reason,
|
||
created_at=now_str(),
|
||
)
|
||
db.session.add(allocation)
|
||
await db.session.commit()
|
||
|
||
return {"success": True, "allocation_id": allocation.id}
|