Files
video-gen/video-gen-api/app/services/generation_billing_service.py
T
2026-07-10 17:26:11 +08:00

599 lines
20 KiB
Python

from __future__ import annotations
import re
from dataclasses import asdict, dataclass
from typing import Any, Mapping
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.enums.credit_record import CreditRecordBillingScene, CreditRecordChargeKind, CreditRecordOwnerType, CreditRecordSourceModule
from app.models.credit_record import CreditRecord
from app.models.generation_record import GenerationRecord
from app.models.module_generation_step import ModuleGenerationStep
from app.models.token_usage import TokenUsage
from app.models.system_config import SystemConfig
from app.services.credit_record_meta_service import (
CreditRecordMeta,
build_generation_media_meta,
build_generation_record_prompt_meta,
build_module_step_prompt_meta,
build_shot_video_analysis_meta,
)
from app.services.credits import calc_image_credits, calc_text_credits, calc_video_credits, deduct_credits
CHARGE_TEXT_PROMPT = CreditRecordChargeKind.TEXT_PROMPT.value
CHARGE_FILE_PARSE = CreditRecordChargeKind.FILE_PARSE.value
CHARGE_VISION_INPUT = CreditRecordChargeKind.VISION_INPUT.value
CHARGE_MEDIA = CreditRecordChargeKind.MEDIA.value
CHARGE_VIDEO_ANALYSIS = CreditRecordChargeKind.VIDEO_ANALYSIS.value
OWNER_GENERATION_RECORD = CreditRecordOwnerType.GENERATION_RECORD.value
OWNER_CHAT_GENERATION_TASK = CreditRecordOwnerType.CHAT_GENERATION_TASK.value
OWNER_MODULE_GENERATION_STEP = CreditRecordOwnerType.MODULE_GENERATION_STEP.value
OWNER_SHOT_REPLICATE_TASK_SET = CreditRecordOwnerType.SHOT_REPLICATE_TASK_SET.value
OWNER_SHOT_REPLICATE_SEGMENT = CreditRecordOwnerType.SHOT_REPLICATE_SEGMENT.value
_BIZ_KEY_PATTERN = re.compile(
r"^(?P<owner_type>[^:]+):(?P<owner_id>[^:]+):attempt:(?P<attempt_no>\d+):(?P<charge_kind>[^:]+):(?P<action>charge|refund)$"
)
@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 and item.charged), 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 _safe_int(value: Any, default: int = 0) -> int:
try:
if value is None:
return default
return int(value)
except Exception:
return default
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", "refund"):
raise ValueError("action 仅支持 charge/refund")
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_config_float_or_none(db: AsyncSession, key: str) -> float | None:
result = await db.execute(select(SystemConfig).where(SystemConfig.key == key).limit(1))
config = result.scalar_one_or_none()
if not config:
return None
try:
return float(config.value)
except Exception:
return None
async def _calc_optional_token_credits(db: AsyncSession, tokens: int, config_key: str) -> float:
tokens = _safe_int(tokens)
if tokens <= 0:
return 0.0
rate = await _get_config_float_or_none(db, config_key)
if rate is None:
return 0.0
return round(tokens * rate / 1000, 2)
async def _find_existing_by_biz_key(db: AsyncSession, *, user_id: str, biz_key: str) -> CreditRecord | 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 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,
) -> BillingItem:
"""按 biz_key 做幂等扣费。
charge_key 只保留为业务分类;正式幂等以 biz_key 为准。
record_meta 负责把业务归属、模块、步骤、token、模型快照写入 CreditRecord。
"""
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)
if biz_key:
existing_charge = await _find_existing_by_biz_key(db, user_id=user_id, biz_key=biz_key)
if existing_charge:
return BillingItem(
charge_key=charge_key,
amount=abs(_round2(existing_charge.amount)),
charged=False,
skipped_reason="already_charged",
biz_key=biz_key,
attempt_no=attempt_no,
)
await deduct_credits(
db,
user_id=user_id,
amount=amount,
description=description,
related_id=related_id,
biz_key=biz_key,
record_meta=record_meta,
)
return BillingItem(charge_key=charge_key, amount=amount, charged=True, biz_key=biz_key, attempt_no=attempt_no)
async def charge_chatapi_prompt_usage(
db: AsyncSession,
*,
record: GenerationRecord,
usage: Mapping[str, Any],
project_name: str | None = None,
) -> BillingSummary:
"""项目记录提示词整理扣费;不参与生成失败媒体退款。"""
project_name = project_name or "AI生成任务"
items: list[BillingItem] = []
attempt_no = 1
owner_type = OWNER_GENERATION_RECORD
owner_id = record.id
input_tokens = _safe_int(usage.get("input_tokens"))
output_tokens = _safe_int(usage.get("output_tokens"))
text_credits = await calc_text_credits(db, input_tokens, output_tokens)
text_meta = await build_generation_record_prompt_meta(
db,
record_id=record.id,
attempt_no=attempt_no,
charge_kind=CHARGE_TEXT_PROMPT,
usage=usage,
media_references=record.media_references,
)
items.append(
await deduct_credits_locked_once(
db,
user_id=record.user_id,
amount=text_credits,
description=f"提示词优化 - {project_name}",
related_id=record.id,
charge_key=CHARGE_TEXT_PROMPT,
biz_key=build_credit_biz_key(
owner_type=owner_type,
owner_id=owner_id,
attempt_no=attempt_no,
charge_kind=CHARGE_TEXT_PROMPT,
action="charge",
),
attempt_no=attempt_no,
record_meta=text_meta,
)
)
file_tokens = usage.get("file_parse_tokens") or usage.get("file_tokens") or usage.get("document_tokens") or 0
file_parse_credits = await _calc_optional_token_credits(db, _safe_int(file_tokens), "file_parse_credits_per_1000_tokens")
file_meta = await build_generation_record_prompt_meta(
db,
record_id=record.id,
attempt_no=attempt_no,
charge_kind=CHARGE_FILE_PARSE,
usage={**dict(usage), "total_tokens": _safe_int(file_tokens), "input_tokens": _safe_int(file_tokens), "output_tokens": 0},
media_references=record.media_references,
)
items.append(
await deduct_credits_locked_once(
db,
user_id=record.user_id,
amount=file_parse_credits,
description="文件解析Token",
related_id=record.id,
charge_key=CHARGE_FILE_PARSE,
biz_key=build_credit_biz_key(
owner_type=owner_type,
owner_id=owner_id,
attempt_no=attempt_no,
charge_kind=CHARGE_FILE_PARSE,
action="charge",
),
attempt_no=attempt_no,
record_meta=file_meta,
)
)
vision_tokens = usage.get("vision_input_tokens") or usage.get("image_input_tokens") or usage.get("image_tokens") or 0
vision_input_credits = await _calc_optional_token_credits(db, _safe_int(vision_tokens), "vision_input_credits_per_1000_tokens")
vision_meta = await build_generation_record_prompt_meta(
db,
record_id=record.id,
attempt_no=attempt_no,
charge_kind=CHARGE_VISION_INPUT,
usage={**dict(usage), "total_tokens": _safe_int(vision_tokens), "input_tokens": _safe_int(vision_tokens), "output_tokens": 0},
media_references=record.media_references,
)
items.append(
await deduct_credits_locked_once(
db,
user_id=record.user_id,
amount=vision_input_credits,
description="图片理解Token",
related_id=record.id,
charge_key=CHARGE_VISION_INPUT,
biz_key=build_credit_biz_key(
owner_type=owner_type,
owner_id=owner_id,
attempt_no=attempt_no,
charge_kind=CHARGE_VISION_INPUT,
action="charge",
),
attempt_no=attempt_no,
record_meta=vision_meta,
)
)
if hasattr(record, "text_credits_cost"):
record.text_credits_cost = round(text_credits + file_parse_credits + vision_input_credits, 2)
if hasattr(record, "text_tokens_used"):
record.text_tokens_used = _safe_int(usage.get("total_tokens"), input_tokens + output_tokens)
return BillingSummary(record_id=record.id, user_id=record.user_id, items=items)
async def charge_module_prompt_usage(
db: AsyncSession,
*,
user_id: str,
step_id: str,
usage: Mapping[str, Any],
description: str,
attempt_no: int = 1,
) -> BillingSummary:
"""爆款开头复刻/拆镜复刻模块图片/视频 AI 提词扣文本积分。
文本提词属于已经发生的 LLM 消费:
- 调用成功后按 input_tokens + output_tokens 扣费。
- 不参与后续图片/视频媒体生成失败退款。
- 通过 module_generation_step:{step_id}:attempt:1:text_prompt:charge 幂等。
"""
input_tokens = _safe_int(usage.get("input_tokens"))
output_tokens = _safe_int(usage.get("output_tokens"))
text_credits = await calc_text_credits(db, input_tokens, output_tokens)
biz_key = build_credit_biz_key(
owner_type=OWNER_MODULE_GENERATION_STEP,
owner_id=step_id,
attempt_no=attempt_no,
charge_kind=CHARGE_TEXT_PROMPT,
action="charge",
)
record_meta = await build_module_step_prompt_meta(db, step_id=step_id, attempt_no=attempt_no, usage=usage)
item = await deduct_credits_locked_once(
db,
user_id=user_id,
amount=text_credits,
description=description,
related_id=step_id,
charge_key=CHARGE_TEXT_PROMPT,
biz_key=biz_key,
attempt_no=attempt_no,
record_meta=record_meta,
)
result = await db.execute(select(ModuleGenerationStep).where(ModuleGenerationStep.id == step_id).limit(1))
step = result.scalar_one_or_none()
if step:
step.token_usage_id = record_meta.token_usage_id
step.model_config_id = usage.get("model_config_id")
step.input_tokens = record_meta.input_tokens
step.output_tokens = record_meta.output_tokens
step.total_tokens = record_meta.total_tokens
step.text_credits_cost = text_credits
return BillingSummary(record_id=step_id, user_id=user_id, items=[item])
async def charge_shot_video_analysis_usage(
db: AsyncSession,
*,
user_id: str,
owner_type: str,
owner_id: str,
usage: Mapping[str, Any],
description: str,
billing_scene: str,
source_project_id: str | None = None,
source_step_id: str | None = None,
attempt_no: int | None = None,
) -> BillingSummary:
"""拆镜复刻视频分析扣分析积分。
视频分析属于“文字提示词 + 视频素材”的模型调用类消费,
按 input_tokens + output_tokens 参考文本积分规则计费,
但账务归类为 analysis/video_analysis,避免混入提词优化统计。
"""
attempt_no = attempt_no or await get_next_credit_attempt_no(
db,
owner_type=owner_type,
owner_id=owner_id,
charge_kind=CHARGE_VIDEO_ANALYSIS,
)
input_tokens = _safe_int(usage.get("input_tokens"))
output_tokens = _safe_int(usage.get("output_tokens"))
amount = await calc_text_credits(db, input_tokens, output_tokens)
biz_key = build_credit_biz_key(
owner_type=owner_type,
owner_id=owner_id,
attempt_no=attempt_no,
charge_kind=CHARGE_VIDEO_ANALYSIS,
action="charge",
)
record_meta = await build_shot_video_analysis_meta(
db,
owner_type=owner_type,
owner_id=owner_id,
attempt_no=attempt_no,
usage=usage,
billing_scene=billing_scene,
source_project_id=source_project_id,
source_step_id=source_step_id,
)
item = await deduct_credits_locked_once(
db,
user_id=user_id,
amount=amount,
description=description,
related_id=owner_id,
charge_key=CHARGE_VIDEO_ANALYSIS,
biz_key=biz_key,
attempt_no=attempt_no,
record_meta=record_meta,
)
if record_meta.token_usage_id:
result = await db.execute(select(TokenUsage).where(TokenUsage.id == record_meta.token_usage_id).limit(1))
token_usage = result.scalar_one_or_none()
if token_usage:
token_usage.owner_type = token_usage.owner_type or owner_type
token_usage.owner_id = token_usage.owner_id or owner_id
token_usage.biz_key = token_usage.biz_key or biz_key
token_usage.source_module = token_usage.source_module or CreditRecordSourceModule.SHOT_REPLICATE.value
token_usage.source_step_code = token_usage.source_step_code or "video_analysis"
return BillingSummary(record_id=owner_id, user_id=user_id, items=[item])
async def charge_generation_media_by_params(
db: AsyncSession,
*,
user_id: str,
record_id: str,
gen_type: str,
image_size: str | None = None,
image_px: str | None = None,
aspect_ratio: str | None = None,
duration: int | None = None,
resolution: str | None = None,
fps: int | float | None = None,
generate_audio: bool | None = None,
inference_mode: str | None = None,
engine_id: str | None = None,
input_video_duration: float | 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,
media_references: Any = None,
requested_output_count: int = 1,
provider_uses_media_references: bool | None = None,
) -> BillingSummary:
"""图片/视频媒体生成扣费。
正式幂等由 owner_type + record_id + attempt_no + media + charge 组成。
每次用户主动重试必须传入新的 attempt_no。
"""
project_name = project_name or "AI生成任务"
gen_type = (gen_type or "").lower().strip()
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,
media_references=media_references,
image_size=image_size,
image_px=image_px,
aspect_ratio=aspect_ratio,
duration=duration,
resolution=resolution,
fps=fps,
generate_audio=generate_audio,
inference_mode=inference_mode,
input_video_duration=input_video_duration,
requested_output_count=requested_output_count,
provider_uses_media_references=provider_uses_media_references,
)
if gen_type == "image":
size = image_size or "2K"
amount = await calc_image_credits(db, size, engine_id=engine_id)
items.append(
await deduct_credits_locked_once(
db,
user_id=user_id,
amount=amount,
description=f"{description_prefix}图片生成",
related_id=record_id,
charge_key=CHARGE_MEDIA,
biz_key=biz_key,
attempt_no=attempt_no,
record_meta=record_meta,
)
)
elif gen_type == "video":
amount = await calc_video_credits(
db, duration or 5, resolution or "720p",
engine_id=engine_id,
input_video_duration=input_video_duration,
)
items.append(
await deduct_credits_locked_once(
db,
user_id=user_id,
amount=amount,
description=f"{description_prefix}视频生成",
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,
) -> BillingSummary:
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,
image_px=record.image_px,
aspect_ratio=record.aspect_ratio or record.image_proportion,
duration=record.duration,
resolution=record.resolution,
project_name=project_name,
description_prefix=description_prefix,
owner_type=OWNER_GENERATION_RECORD,
attempt_no=attempt_no,
source_module=CreditRecordSourceModule.GENERATION_RECORD.value,
media_references=record.media_references,
provider_uses_media_references=False,
)