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[^:]+):(?P[^:]+):attempt:(?P\d+):(?P[^:]+):(?Pcharge|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, )