309 lines
9.9 KiB
Python
309 lines
9.9 KiB
Python
from __future__ import annotations
|
||
|
||
import re
|
||
from dataclasses import asdict, dataclass
|
||
from typing import Any
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.enums.credit_record import CreditRecordChargeKind, CreditRecordOwnerType, CreditRecordSourceModule
|
||
from app.models.credit_record import CreditRecord
|
||
from app.models.generation_record import GenerationRecord
|
||
from app.services.generation.media_reference_service import calculate_media_reference_usage
|
||
from app.services.credit_record_meta_service import (
|
||
CreditRecordMeta,
|
||
build_generation_media_meta,
|
||
)
|
||
from app.services.credits import calc_image_credits, calc_video_credits, deduct_credits_result
|
||
|
||
|
||
CHARGE_MEDIA = CreditRecordChargeKind.MEDIA.value
|
||
|
||
OWNER_GENERATION_RECORD = CreditRecordOwnerType.GENERATION_RECORD.value
|
||
OWNER_CHAT_GENERATION_TASK = CreditRecordOwnerType.CHAT_GENERATION_TASK.value
|
||
|
||
_BIZ_KEY_PATTERN = re.compile(
|
||
r"^(?P<owner_type>[^:]+):(?P<owner_id>[^:]+):attempt:(?P<attempt_no>\d+):(?P<charge_kind>[^:]+):(?P<action>charge|llm_charge|refund|pre_deduct|hold|hold_release)$"
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class BillingItem:
|
||
charge_key: str
|
||
amount: float
|
||
charged: bool
|
||
skipped_reason: str | None = None
|
||
biz_key: str | None = None
|
||
attempt_no: int | None = None
|
||
|
||
|
||
@dataclass
|
||
class BillingSummary:
|
||
record_id: str
|
||
user_id: str
|
||
items: list[BillingItem]
|
||
|
||
@property
|
||
def total_charged(self) -> float:
|
||
return round(sum(item.amount for item in self.items if item.charged), 2)
|
||
|
||
def get_amount(self, charge_key: str) -> float:
|
||
"""返回该业务动作的已记录金额；幂等重放命中旧流水时也返回真实金额。"""
|
||
return round(sum(item.amount for item in self.items if item.charge_key == charge_key), 2)
|
||
|
||
def to_dict(self) -> dict[str, Any]:
|
||
data = asdict(self)
|
||
data["total_charged"] = self.total_charged
|
||
return data
|
||
|
||
|
||
def _round2(value: float | int | None) -> float:
|
||
return round(float(value or 0), 2)
|
||
|
||
|
||
def build_credit_biz_key(
|
||
*,
|
||
owner_type: str,
|
||
owner_id: str,
|
||
attempt_no: int,
|
||
charge_kind: str,
|
||
action: str,
|
||
) -> str:
|
||
"""生成正式积分流水幂等键。
|
||
|
||
示例：generation_record:xxx:attempt:2:media:charge
|
||
"""
|
||
owner_type = owner_type.strip()
|
||
owner_id = owner_id.strip()
|
||
charge_kind = charge_kind.strip()
|
||
action = action.strip()
|
||
if action not in ("charge", "llm_charge", "refund", "pre_deduct", "hold", "hold_release"):
|
||
raise ValueError("action 仅支持 charge/llm_charge/refund/pre_deduct/hold/hold_release")
|
||
if attempt_no <= 0:
|
||
raise ValueError("attempt_no 必须大于 0")
|
||
return f"{owner_type}:{owner_id}:attempt:{attempt_no}:{charge_kind}:{action}"
|
||
|
||
|
||
def parse_credit_biz_key(biz_key: str | None) -> dict[str, Any] | None:
|
||
if not biz_key:
|
||
return None
|
||
match = _BIZ_KEY_PATTERN.match(biz_key)
|
||
if not match:
|
||
return None
|
||
data = match.groupdict()
|
||
data["attempt_no"] = int(data["attempt_no"])
|
||
return data
|
||
|
||
|
||
async def get_next_credit_attempt_no(
|
||
db: AsyncSession,
|
||
*,
|
||
owner_type: str,
|
||
owner_id: str,
|
||
charge_kind: str = CHARGE_MEDIA,
|
||
) -> int:
|
||
"""根据已有正式 biz_key 计算下一轮扣费 attempt_no。
|
||
|
||
不依赖任务表 retry_count，防止 worker 崩溃/重复消息造成状态字段不可信。
|
||
"""
|
||
prefix = f"{owner_type}:{owner_id}:attempt:%:{charge_kind}:charge"
|
||
result = await db.execute(
|
||
select(CreditRecord.biz_key)
|
||
.where(CreditRecord.related_id == owner_id)
|
||
.where(CreditRecord.type == "consume")
|
||
.where(CreditRecord.biz_key.like(prefix))
|
||
)
|
||
max_attempt = 0
|
||
for (biz_key,) in result.all():
|
||
parsed = parse_credit_biz_key(biz_key)
|
||
if parsed and parsed.get("owner_type") == owner_type and parsed.get("owner_id") == owner_id:
|
||
max_attempt = max(max_attempt, int(parsed.get("attempt_no") or 0))
|
||
return max_attempt + 1
|
||
|
||
|
||
async def deduct_credits_locked_once(
|
||
db: AsyncSession,
|
||
*,
|
||
user_id: str,
|
||
amount: float,
|
||
description: str,
|
||
related_id: str,
|
||
charge_key: str,
|
||
biz_key: str | None = None,
|
||
attempt_no: int | None = None,
|
||
record_meta: CreditRecordMeta | dict | None = None,
|
||
allow_negative: bool = False,
|
||
) -> BillingItem:
|
||
"""按 biz_key 做幂等扣费。
|
||
|
||
幂等判断、用户行锁、余额更新和流水写入由 deduct_credits_result 在同一短事务内完成，
|
||
避免先查一次 biz_key、加锁后再查一次的重复 SQL 和竞态窗口。
|
||
"""
|
||
amount = _round2(amount)
|
||
if amount <= 0:
|
||
return BillingItem(
|
||
charge_key=charge_key,
|
||
amount=0.0,
|
||
charged=False,
|
||
skipped_reason="amount_lte_zero",
|
||
biz_key=biz_key,
|
||
attempt_no=attempt_no,
|
||
)
|
||
|
||
mutation = await deduct_credits_result(
|
||
db,
|
||
user_id=user_id,
|
||
amount=amount,
|
||
description=description,
|
||
related_id=related_id,
|
||
biz_key=biz_key,
|
||
record_meta=record_meta,
|
||
allow_negative=allow_negative,
|
||
)
|
||
return BillingItem(
|
||
charge_key=charge_key,
|
||
amount=mutation.amount if not mutation.created else amount,
|
||
charged=mutation.created,
|
||
skipped_reason=None if mutation.created else "already_charged",
|
||
biz_key=biz_key,
|
||
attempt_no=attempt_no,
|
||
)
|
||
|
||
|
||
async def charge_generation_media_by_params(
|
||
db: AsyncSession,
|
||
*,
|
||
user_id: str,
|
||
record_id: str,
|
||
gen_type: str,
|
||
image_size: str | None = None,
|
||
duration: int | None = None,
|
||
resolution: str | None = None,
|
||
engine_id: str | None = None,
|
||
input_video_duration: float | None = None,
|
||
input_image_count: int | None = None,
|
||
project_name: str | None = None,
|
||
description_prefix: str = "AI创作-",
|
||
owner_type: str = OWNER_CHAT_GENERATION_TASK,
|
||
attempt_no: int | None = None,
|
||
source_module: str | None = None,
|
||
source_project_id: str | None = None,
|
||
source_step_id: str | None = None,
|
||
source_step_code: str | None = None,
|
||
billing_scene: str | None = None,
|
||
quantity: int = 1,
|
||
) -> BillingSummary:
|
||
"""图片/视频媒体生成扣费。
|
||
|
||
正式幂等由 owner_type + record_id + attempt_no + media + charge 组成。
|
||
每次用户主动重试必须传入新的 attempt_no。
|
||
"""
|
||
project_name = project_name or "AI生成任务"
|
||
gen_type = (gen_type or "").lower().strip()
|
||
quantity = max(1, int(quantity or 1))
|
||
attempt_no = attempt_no or await get_next_credit_attempt_no(
|
||
db,
|
||
owner_type=owner_type,
|
||
owner_id=record_id,
|
||
charge_kind=CHARGE_MEDIA,
|
||
)
|
||
biz_key = build_credit_biz_key(
|
||
owner_type=owner_type,
|
||
owner_id=record_id,
|
||
attempt_no=attempt_no,
|
||
charge_kind=CHARGE_MEDIA,
|
||
action="charge",
|
||
)
|
||
items: list[BillingItem] = []
|
||
record_meta = await build_generation_media_meta(
|
||
db,
|
||
owner_type=owner_type,
|
||
owner_id=record_id,
|
||
attempt_no=attempt_no,
|
||
gen_type=gen_type,
|
||
engine_id=engine_id,
|
||
source_module=source_module,
|
||
source_project_id=source_project_id,
|
||
source_step_id=source_step_id,
|
||
source_step_code=source_step_code,
|
||
billing_scene=billing_scene,
|
||
)
|
||
|
||
if gen_type == "image":
|
||
size = image_size or "2K"
|
||
unit_amount = await calc_image_credits(
|
||
db, size, engine_id=engine_id, input_image_count=input_image_count
|
||
)
|
||
amount = round(unit_amount * quantity, 2)
|
||
items.append(
|
||
await deduct_credits_locked_once(
|
||
db,
|
||
user_id=user_id,
|
||
amount=amount,
|
||
description=f"{description_prefix}图片生成" + (f"×{quantity}" if quantity > 1 else ""),
|
||
related_id=record_id,
|
||
charge_key=CHARGE_MEDIA,
|
||
biz_key=biz_key,
|
||
attempt_no=attempt_no,
|
||
record_meta=record_meta,
|
||
)
|
||
)
|
||
elif gen_type == "video":
|
||
unit_amount = await calc_video_credits(
|
||
db, duration or 5, resolution or "720p",
|
||
engine_id=engine_id,
|
||
input_video_duration=input_video_duration,
|
||
input_image_count=input_image_count,
|
||
)
|
||
amount = round(unit_amount * quantity, 2)
|
||
items.append(
|
||
await deduct_credits_locked_once(
|
||
db,
|
||
user_id=user_id,
|
||
amount=amount,
|
||
description=f"{description_prefix}视频生成" + (f"×{quantity}" if quantity > 1 else ""),
|
||
related_id=record_id,
|
||
charge_key=CHARGE_MEDIA,
|
||
biz_key=biz_key,
|
||
attempt_no=attempt_no,
|
||
record_meta=record_meta,
|
||
)
|
||
)
|
||
else:
|
||
raise ValueError(f"不支持的生成类型: {gen_type}")
|
||
|
||
return BillingSummary(record_id=record_id, user_id=user_id, items=items)
|
||
|
||
|
||
async def charge_generation_media_for_record(
|
||
db: AsyncSession,
|
||
*,
|
||
record: GenerationRecord,
|
||
project_name: str | None = None,
|
||
description_prefix: str = "AI创作-",
|
||
attempt_no: int | None = None,
|
||
engine_id: str | None = None,
|
||
) -> BillingSummary:
|
||
reference_usage = calculate_media_reference_usage(
|
||
record.media_references,
|
||
include=bool(record.include_media_references),
|
||
)
|
||
return await charge_generation_media_by_params(
|
||
db,
|
||
user_id=record.user_id,
|
||
record_id=record.id,
|
||
gen_type=record.gen_type,
|
||
image_size=record.image_size,
|
||
duration=record.duration,
|
||
resolution=record.resolution,
|
||
engine_id=engine_id or getattr(record, "engine_id", None),
|
||
input_video_duration=reference_usage.input_video_duration or None,
|
||
input_image_count=reference_usage.image_count or None,
|
||
project_name=project_name,
|
||
description_prefix=description_prefix,
|
||
owner_type=OWNER_GENERATION_RECORD,
|
||
attempt_no=attempt_no,
|
||
source_module=CreditRecordSourceModule.GENERATION_RECORD.value,
|
||
)
|