479 lines
16 KiB
Python
479 lines
16 KiB
Python
import math
|
|
from dataclasses import dataclass
|
|
|
|
from sqlalchemy import select, func
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.user import User
|
|
from app.models.credit_record import CreditRecord
|
|
from app.models.video_engine import VideoEngine
|
|
from app.models.image_engine import ImageEngine
|
|
from app.models.credit_ratio import CreditRatio
|
|
from app.utils.id_gen import generate_id
|
|
from app.utils.exceptions import InsufficientCreditsError
|
|
from app.enums.common import BillingBlockEventEnum
|
|
from app.services.operation_log_service import log_operation_event
|
|
from app.services.system_config_cache import get_system_config_value
|
|
from app.services.credit_record_meta_service import CreditRecordMeta, with_user_snapshot
|
|
|
|
|
|
async def calc_text_credits(db: AsyncSession, input_tokens: int, output_tokens: int) -> float:
|
|
"""Calculate text credits based on actual token usage and cached configurable rate."""
|
|
raw_rate = await get_system_config_value(db, "text_credits_per_1000_tokens")
|
|
try:
|
|
rate = float(raw_rate) if raw_rate not in (None, "") else 1.0
|
|
except (TypeError, ValueError):
|
|
rate = 1.0
|
|
total_tokens = input_tokens + output_tokens
|
|
return round(total_tokens * rate / 1000, 2)
|
|
|
|
|
|
async def _get_credit_ratio(
|
|
db: AsyncSession,
|
|
*,
|
|
gen_type: str,
|
|
resolution: str,
|
|
engine_id: str | None = None,
|
|
) -> CreditRatio | None:
|
|
"""按引擎精确规则优先获取积分规则;找不到时回退到同类型同参数最高规则。"""
|
|
gen_type = (gen_type or "").lower().strip()
|
|
resolution = (resolution or "").strip()
|
|
engine_id = (engine_id or "").strip() or None
|
|
|
|
if engine_id:
|
|
result = await db.execute(
|
|
select(CreditRatio)
|
|
.where(CreditRatio.gen_type == gen_type)
|
|
.where(CreditRatio.model_config_id == engine_id)
|
|
.where(CreditRatio.resolution == resolution)
|
|
.order_by(CreditRatio.base_credits.desc(), CreditRatio.per_second_credits.desc())
|
|
.limit(1)
|
|
)
|
|
ratio = result.scalar_one_or_none()
|
|
if ratio:
|
|
return ratio
|
|
|
|
result = await db.execute(
|
|
select(CreditRatio)
|
|
.where(CreditRatio.gen_type == gen_type)
|
|
.where(CreditRatio.resolution == resolution)
|
|
.order_by(CreditRatio.base_credits.desc(), CreditRatio.per_second_credits.desc())
|
|
.limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
async def calc_video_credits(
|
|
db: AsyncSession,
|
|
duration: int,
|
|
resolution: str,
|
|
engine_id: str | None = None,
|
|
input_video_duration: float | None = None,
|
|
input_image_count: int | None = None,
|
|
) -> float:
|
|
"""Calculate video credits using CreditRatio table, with fallback to hardcoded.
|
|
|
|
查询优先级:
|
|
1. gen_type=video + engine_id + resolution 精确规则;
|
|
2. gen_type=video + resolution 下 base_credits/per_second_credits 最高规则;
|
|
3. 原硬编码默认算法。
|
|
|
|
input_video_duration: 用户上传的参考视频总时长(秒),不为空时额外计费
|
|
input_image_count: 用户上传的参考图片数量,不为空时额外计费
|
|
"""
|
|
if not engine_id:
|
|
video_engines_result = await db.execute(
|
|
select(VideoEngine.id)
|
|
.where(VideoEngine.is_active == True, VideoEngine.deleted_at.is_(None))
|
|
.order_by(VideoEngine.priority.desc())
|
|
.limit(1)
|
|
)
|
|
engine_id = video_engines_result.scalar_one_or_none()
|
|
ratio = await _get_credit_ratio(
|
|
db,
|
|
gen_type="video",
|
|
resolution=resolution,
|
|
engine_id=engine_id,
|
|
)
|
|
if ratio:
|
|
base_cost = (ratio.base_credits + ratio.per_second_credits * duration) * ratio.ratio
|
|
if input_video_duration and input_video_duration > 0:
|
|
input_video_cost = (
|
|
ratio.input_video_base_credits + ratio.input_video_per_second_credits * input_video_duration
|
|
) * ratio.input_video_ratio
|
|
base_cost += input_video_cost
|
|
if input_image_count and input_image_count > 0:
|
|
input_image_cost = (
|
|
ratio.input_image_base_credits + ratio.input_image_per_image_credits * input_image_count
|
|
) * ratio.input_image_ratio
|
|
base_cost += input_image_cost
|
|
return round(base_cost, 2)
|
|
|
|
base = 60.0
|
|
duration_cost = duration * 2.0
|
|
multiplier = {"480p": 1, "1080p": 2, "720p": 1.5}.get(resolution, 1.0)
|
|
total = (base + duration_cost) * multiplier
|
|
if input_video_duration and input_video_duration > 0:
|
|
total += input_video_duration * 0.5 * multiplier
|
|
if input_image_count and input_image_count > 0:
|
|
total += input_image_count * 0.5 * multiplier
|
|
return round(total, 2)
|
|
|
|
|
|
def calc_credits(duration: int, resolution: str) -> float:
|
|
"""Legacy: hardcoded credit calculation. Prefer calc_video_credits for new code."""
|
|
base = 60.0
|
|
duration_cost = duration * 2.0
|
|
multiplier = {"480p": 2.5, "1080p": 1.5, "720p": 1.0}.get(resolution, 1.0)
|
|
return round((base + duration_cost) * multiplier, 2)
|
|
|
|
|
|
async def calc_image_credits(
|
|
db: AsyncSession,
|
|
image_size: str,
|
|
engine_id: str | None = None,
|
|
input_image_count: int | None = None,
|
|
) -> float:
|
|
"""Calculate image credits using CreditRatio table, with fallback to hardcoded.
|
|
|
|
查询优先级:
|
|
1. gen_type=image + engine_id + image_size 精确规则;
|
|
2. gen_type=image + image_size 下 base_credits/per_second_credits 最高规则;
|
|
3. 原硬编码默认算法。
|
|
|
|
input_image_count: 用户上传的参考图片数量,不为空时额外计费
|
|
"""
|
|
# 如果engine_id为空,默认查询权重最高的图片引擎积分规则
|
|
if not engine_id:
|
|
image_engines_result = await db.execute(
|
|
select(ImageEngine.id)
|
|
.where(ImageEngine.is_active == True, ImageEngine.deleted_at.is_(None))
|
|
.order_by(ImageEngine.priority.desc())
|
|
.limit(1)
|
|
)
|
|
engine_id = image_engines_result.scalar_one_or_none()
|
|
ratio = await _get_credit_ratio(
|
|
db,
|
|
gen_type="image",
|
|
resolution=image_size,
|
|
engine_id=engine_id,
|
|
)
|
|
if ratio:
|
|
base_cost = ratio.base_credits * ratio.ratio
|
|
if input_image_count and input_image_count > 0:
|
|
input_image_cost = (
|
|
ratio.input_image_base_credits + ratio.input_image_per_image_credits * input_image_count
|
|
) * ratio.input_image_ratio
|
|
base_cost += input_image_cost
|
|
return round(base_cost, 2)
|
|
|
|
# Fallback
|
|
multiplier = {"4K": 2.0, "2K": 1.0}.get(image_size, 1.0)
|
|
base_cost = 4.0
|
|
total = base_cost * multiplier
|
|
if input_image_count and input_image_count > 0:
|
|
total += input_image_count * 0.5 * multiplier
|
|
return round(total, 2)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CreditMutationResult:
|
|
user: User
|
|
record: CreditRecord | None
|
|
created: bool
|
|
amount: float
|
|
balance_before: float
|
|
balance_after: float
|
|
|
|
|
|
async def _get_existing_credit_record_by_biz_key(
|
|
db: AsyncSession,
|
|
*,
|
|
user_id: str,
|
|
biz_key: str | None,
|
|
) -> CreditRecord | None:
|
|
"""按正式业务幂等键查找已有积分流水。"""
|
|
if not biz_key:
|
|
return None
|
|
result = await db.execute(
|
|
select(CreditRecord)
|
|
.where(CreditRecord.user_id == user_id, CreditRecord.biz_key == biz_key)
|
|
.limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
async def deduct_credits_result(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
amount: float,
|
|
description: str,
|
|
related_id: str | None = None,
|
|
*,
|
|
biz_key: str | None = None,
|
|
refund_for_biz_key: str | None = None,
|
|
record_meta: CreditRecordMeta | dict | None = None,
|
|
record_type: str = "consume",
|
|
allow_negative: bool = False,
|
|
create_zero_record: bool = False,
|
|
) -> CreditMutationResult:
|
|
"""并发安全且可观察幂等结果的积分扣减。"""
|
|
amount = round(float(amount or 0), 2)
|
|
result = await db.execute(select(User).where(User.id == user_id).with_for_update().limit(1))
|
|
user = result.scalar_one_or_none()
|
|
if not user:
|
|
raise ValueError("User not found")
|
|
|
|
before_balance = round(float(user.credits or 0), 2)
|
|
if biz_key:
|
|
existing = await _get_existing_credit_record_by_biz_key(db, user_id=user_id, biz_key=biz_key)
|
|
if existing:
|
|
return CreditMutationResult(
|
|
user=user,
|
|
record=existing,
|
|
created=False,
|
|
amount=abs(round(float(existing.amount or 0), 2)),
|
|
balance_before=before_balance,
|
|
balance_after=before_balance,
|
|
)
|
|
|
|
if amount <= 0 and not create_zero_record:
|
|
return CreditMutationResult(
|
|
user=user,
|
|
record=None,
|
|
created=False,
|
|
amount=0.0,
|
|
balance_before=before_balance,
|
|
balance_after=before_balance,
|
|
)
|
|
|
|
if amount > 0 and not allow_negative and before_balance < amount:
|
|
event_type = (
|
|
BillingBlockEventEnum.NEGATIVE_BALANCE.value
|
|
if before_balance < 0
|
|
else BillingBlockEventEnum.INSUFFICIENT_CREDITS.value
|
|
)
|
|
log_operation_event(
|
|
domain="billing",
|
|
module="credits",
|
|
event_type=event_type,
|
|
event_status="failed",
|
|
source="app.services.credits.deduct_credits_result",
|
|
user_id=user_id,
|
|
task_id=related_id,
|
|
message="积分不足,已拦截新的扣费请求",
|
|
detail={
|
|
"user_id": user_id,
|
|
"amount": amount,
|
|
"before_balance": before_balance,
|
|
"allow_negative": allow_negative,
|
|
"biz_key": biz_key,
|
|
"refund_for_biz_key": refund_for_biz_key,
|
|
"description": description,
|
|
"record_type": record_type,
|
|
},
|
|
)
|
|
raise InsufficientCreditsError()
|
|
|
|
user.credits = round(before_balance - max(0.0, amount), 2)
|
|
meta_kwargs: dict = {}
|
|
if record_meta:
|
|
if isinstance(record_meta, CreditRecordMeta):
|
|
record_meta = await with_user_snapshot(db, record_meta, user_id, user=user)
|
|
meta_kwargs = record_meta.to_record_kwargs()
|
|
elif isinstance(record_meta, dict):
|
|
meta_kwargs = {k: v for k, v in record_meta.items() if v is not None}
|
|
record = CreditRecord(
|
|
id=generate_id(),
|
|
user_id=user_id,
|
|
type=record_type,
|
|
amount=-amount,
|
|
balance_after=user.credits,
|
|
description=description,
|
|
related_id=related_id,
|
|
biz_key=biz_key,
|
|
refund_for_biz_key=refund_for_biz_key,
|
|
**meta_kwargs,
|
|
)
|
|
db.add(record)
|
|
await db.flush()
|
|
return CreditMutationResult(
|
|
user=user,
|
|
record=record,
|
|
created=True,
|
|
amount=amount,
|
|
balance_before=before_balance,
|
|
balance_after=round(float(user.credits or 0), 2),
|
|
)
|
|
|
|
|
|
async def deduct_credits(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
amount: float,
|
|
description: str,
|
|
related_id: str | None = None,
|
|
*,
|
|
biz_key: str | None = None,
|
|
refund_for_biz_key: str | None = None,
|
|
record_meta: CreditRecordMeta | dict | None = None,
|
|
record_type: str = "consume",
|
|
allow_negative: bool = False,
|
|
create_zero_record: bool = False,
|
|
) -> User:
|
|
"""兼容旧调用:返回 User;精确幂等状态请使用 deduct_credits_result。"""
|
|
mutation = await deduct_credits_result(
|
|
db,
|
|
user_id=user_id,
|
|
amount=amount,
|
|
description=description,
|
|
related_id=related_id,
|
|
biz_key=biz_key,
|
|
refund_for_biz_key=refund_for_biz_key,
|
|
record_meta=record_meta,
|
|
record_type=record_type,
|
|
allow_negative=allow_negative,
|
|
create_zero_record=create_zero_record,
|
|
)
|
|
return mutation.user
|
|
|
|
|
|
async def add_credits_result(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
amount: float,
|
|
description: str,
|
|
related_id: str | None = None,
|
|
*,
|
|
record_type: str = "recharge",
|
|
biz_key: str | None = None,
|
|
refund_for_biz_key: str | None = None,
|
|
record_meta: CreditRecordMeta | dict | None = None,
|
|
) -> CreditMutationResult:
|
|
"""并发安全且可观察幂等结果的积分增加。"""
|
|
amount = round(float(amount or 0), 2)
|
|
result = await db.execute(select(User).where(User.id == user_id).with_for_update().limit(1))
|
|
user = result.scalar_one_or_none()
|
|
if not user:
|
|
raise ValueError("User not found")
|
|
|
|
before_balance = round(float(user.credits or 0), 2)
|
|
if biz_key:
|
|
existing = await _get_existing_credit_record_by_biz_key(db, user_id=user_id, biz_key=biz_key)
|
|
if existing:
|
|
return CreditMutationResult(
|
|
user=user,
|
|
record=existing,
|
|
created=False,
|
|
amount=abs(round(float(existing.amount or 0), 2)),
|
|
balance_before=before_balance,
|
|
balance_after=before_balance,
|
|
)
|
|
|
|
if amount <= 0:
|
|
return CreditMutationResult(
|
|
user=user,
|
|
record=None,
|
|
created=False,
|
|
amount=0.0,
|
|
balance_before=before_balance,
|
|
balance_after=before_balance,
|
|
)
|
|
|
|
user.credits = round(before_balance + amount, 2)
|
|
meta_kwargs: dict = {}
|
|
if record_meta:
|
|
if isinstance(record_meta, CreditRecordMeta):
|
|
record_meta = await with_user_snapshot(db, record_meta, user_id, user=user)
|
|
meta_kwargs = record_meta.to_record_kwargs()
|
|
elif isinstance(record_meta, dict):
|
|
meta_kwargs = {k: v for k, v in record_meta.items() if v is not None}
|
|
record = CreditRecord(
|
|
id=generate_id(),
|
|
user_id=user_id,
|
|
type=record_type,
|
|
amount=amount,
|
|
balance_after=user.credits,
|
|
description=description,
|
|
related_id=related_id,
|
|
biz_key=biz_key,
|
|
refund_for_biz_key=refund_for_biz_key,
|
|
**meta_kwargs,
|
|
)
|
|
db.add(record)
|
|
await db.flush()
|
|
return CreditMutationResult(
|
|
user=user,
|
|
record=record,
|
|
created=True,
|
|
amount=amount,
|
|
balance_before=before_balance,
|
|
balance_after=round(float(user.credits or 0), 2),
|
|
)
|
|
|
|
|
|
async def add_credits(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
amount: float,
|
|
description: str,
|
|
related_id: str | None = None,
|
|
*,
|
|
record_type: str = "recharge",
|
|
biz_key: str | None = None,
|
|
refund_for_biz_key: str | None = None,
|
|
record_meta: CreditRecordMeta | dict | None = None,
|
|
) -> User:
|
|
"""兼容旧调用:返回 User;精确幂等状态请使用 add_credits_result。"""
|
|
mutation = await add_credits_result(
|
|
db,
|
|
user_id=user_id,
|
|
amount=amount,
|
|
description=description,
|
|
related_id=related_id,
|
|
record_type=record_type,
|
|
biz_key=biz_key,
|
|
refund_for_biz_key=refund_for_biz_key,
|
|
record_meta=record_meta,
|
|
)
|
|
return mutation.user
|
|
|
|
|
|
async def refund_credits(
|
|
db: AsyncSession,
|
|
user_id: str,
|
|
amount: float,
|
|
description: str,
|
|
related_id: str | None = None,
|
|
*,
|
|
biz_key: str | None = None,
|
|
refund_for_biz_key: str | None = None,
|
|
record_meta: CreditRecordMeta | dict | None = None,
|
|
) -> User:
|
|
"""生成失败积分回退。"""
|
|
return await add_credits(
|
|
db,
|
|
user_id=user_id,
|
|
amount=amount,
|
|
description=description,
|
|
related_id=related_id,
|
|
record_type="refund",
|
|
biz_key=biz_key,
|
|
refund_for_biz_key=refund_for_biz_key,
|
|
record_meta=record_meta,
|
|
)
|
|
|
|
|
|
async def get_records(db: AsyncSession, user_id: str, page: int = 1, page_size: int = 20) -> tuple[list[CreditRecord], int]:
|
|
count_query = select(func.count(CreditRecord.id)).where(CreditRecord.user_id == user_id)
|
|
total = (await db.execute(count_query)).scalar() or 0
|
|
|
|
result = await db.execute(
|
|
select(CreditRecord)
|
|
.where(CreditRecord.user_id == user_id)
|
|
.order_by(CreditRecord.created_at.desc())
|
|
.offset((page - 1) * page_size)
|
|
.limit(page_size)
|
|
)
|
|
return list(result.scalars().all()), total
|