322 lines
10 KiB
Python
322 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, asdict
|
|
from typing import Any, Mapping
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.credit_record import CreditRecord
|
|
from app.models.generation_record import GenerationRecord
|
|
from app.models.system_config import SystemConfig
|
|
from app.models.user import User
|
|
from app.services.credits import calc_image_credits, calc_text_credits, calc_video_credits
|
|
from app.utils.exceptions import InsufficientCreditsError
|
|
from app.utils.id_gen import generate_id
|
|
|
|
|
|
CHARGE_TEXT_PROMPT = "CHATAPI_TEXT_PROMPT"
|
|
CHARGE_FILE_PARSE = "CHATAPI_FILE_PARSE"
|
|
CHARGE_VISION_INPUT = "CHATAPI_VISION_INPUT"
|
|
CHARGE_MEDIA_IMAGE = "CHATAPI_MEDIA_IMAGE"
|
|
CHARGE_MEDIA_VIDEO = "CHATAPI_MEDIA_VIDEO"
|
|
|
|
|
|
@dataclass
|
|
class BillingItem:
|
|
charge_key: str
|
|
amount: float
|
|
charged: bool
|
|
skipped_reason: str | 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
|
|
|
|
|
|
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:
|
|
"""Calculate optional token billing. Missing config means do not charge.
|
|
|
|
This prevents double-charging existing projects where uploaded file/OCR/vision
|
|
content is already included in the LLM provider's input_tokens.
|
|
"""
|
|
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)
|
|
|
|
|
|
def _legacy_description_keywords(charge_key: str) -> list[str]:
|
|
# Compatibility with old patch/original project records that were inserted
|
|
# before this safe billing service added [CHARGE_KEY] prefixes.
|
|
if charge_key == CHARGE_TEXT_PROMPT:
|
|
return ["ChatAPI提示词整理", "提示词优化"]
|
|
if charge_key == CHARGE_MEDIA_IMAGE:
|
|
return ["ChatAPI异步图片生成", "图片生成"]
|
|
if charge_key == CHARGE_MEDIA_VIDEO:
|
|
return ["ChatAPI异步视频生成", "视频生成"]
|
|
if charge_key == CHARGE_FILE_PARSE:
|
|
return ["文件解析Token"]
|
|
if charge_key == CHARGE_VISION_INPUT:
|
|
return ["图片理解Token"]
|
|
return []
|
|
|
|
|
|
async def _find_existing_charge(db: AsyncSession, user_id: str, related_id: str, charge_key: str) -> CreditRecord | None:
|
|
base = (
|
|
select(CreditRecord)
|
|
.where(CreditRecord.user_id == user_id)
|
|
.where(CreditRecord.related_id == related_id)
|
|
.where(CreditRecord.type == "consume")
|
|
)
|
|
|
|
result = await db.execute(base.where(CreditRecord.description.like(f"[{charge_key}]%")).limit(1))
|
|
existing = result.scalar_one_or_none()
|
|
if existing:
|
|
return existing
|
|
|
|
for keyword in _legacy_description_keywords(charge_key):
|
|
result = await db.execute(base.where(CreditRecord.description.like(f"%{keyword}%")).limit(1))
|
|
existing = result.scalar_one_or_none()
|
|
if existing:
|
|
return existing
|
|
return None
|
|
|
|
|
|
async def deduct_credits_locked_once(
|
|
db: AsyncSession,
|
|
*,
|
|
user_id: str,
|
|
amount: float,
|
|
description: str,
|
|
related_id: str,
|
|
charge_key: str,
|
|
) -> BillingItem:
|
|
"""Deduct credits with row lock and idempotency.
|
|
|
|
- User row is locked by SELECT ... FOR UPDATE, so concurrent deductions for
|
|
the same user are serialized in PostgreSQL/MySQL.
|
|
- CreditRecord description prefix + related_id is used as an idempotency key
|
|
without changing existing table structures.
|
|
"""
|
|
amount = _round2(amount)
|
|
if amount <= 0:
|
|
return BillingItem(charge_key=charge_key, amount=0.0, charged=False, skipped_reason="amount_lte_zero")
|
|
|
|
result = await db.execute(select(User).where(User.id == user_id).with_for_update())
|
|
user = result.scalar_one_or_none()
|
|
if not user:
|
|
raise ValueError("User not found")
|
|
|
|
existing_charge = await _find_existing_charge(db, user_id, related_id, charge_key)
|
|
if existing_charge:
|
|
return BillingItem(
|
|
charge_key=charge_key,
|
|
amount=abs(_round2(existing_charge.amount)),
|
|
charged=False,
|
|
skipped_reason="already_charged",
|
|
)
|
|
|
|
if float(user.credits or 0) < amount:
|
|
raise InsufficientCreditsError()
|
|
|
|
user.credits = round(float(user.credits or 0) - amount, 2)
|
|
db.add(
|
|
CreditRecord(
|
|
id=generate_id(),
|
|
user_id=user_id,
|
|
type="consume",
|
|
amount=-amount,
|
|
balance_after=user.credits,
|
|
description=f"[{charge_key}] {description}",
|
|
related_id=related_id,
|
|
)
|
|
)
|
|
await db.flush()
|
|
return BillingItem(charge_key=charge_key, amount=amount, charged=True)
|
|
|
|
|
|
async def charge_chatapi_prompt_usage(
|
|
db: AsyncSession,
|
|
*,
|
|
record: GenerationRecord,
|
|
usage: Mapping[str, Any],
|
|
project_name: str | None = None,
|
|
) -> BillingSummary:
|
|
"""Charge ChatAPI prompt optimization and optional uploaded-file/vision tokens.
|
|
|
|
file_parse_credits / vision_input_credits are optional and disabled unless
|
|
SystemConfig contains these keys:
|
|
- file_parse_credits_per_1000_tokens
|
|
- vision_input_credits_per_1000_tokens
|
|
"""
|
|
project_name = project_name or "AI生成任务"
|
|
items: list[BillingItem] = []
|
|
|
|
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)
|
|
items.append(
|
|
await deduct_credits_locked_once(
|
|
db,
|
|
user_id=record.user_id,
|
|
amount=text_credits,
|
|
description=f"ChatAPI提示词整理 - {project_name}",
|
|
related_id=record.id,
|
|
charge_key=CHARGE_TEXT_PROMPT,
|
|
)
|
|
)
|
|
|
|
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",
|
|
)
|
|
items.append(
|
|
await deduct_credits_locked_once(
|
|
db,
|
|
user_id=record.user_id,
|
|
amount=file_parse_credits,
|
|
description=f"文件解析Token - {project_name}",
|
|
related_id=record.id,
|
|
charge_key=CHARGE_FILE_PARSE,
|
|
)
|
|
)
|
|
|
|
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",
|
|
)
|
|
items.append(
|
|
await deduct_credits_locked_once(
|
|
db,
|
|
user_id=record.user_id,
|
|
amount=vision_input_credits,
|
|
description=f"图片理解Token - {project_name}",
|
|
related_id=record.id,
|
|
charge_key=CHARGE_VISION_INPUT,
|
|
)
|
|
)
|
|
|
|
if hasattr(record, "text_credits_cost"):
|
|
# Store expected text-side cost, even when this task is a retry and the
|
|
# actual CreditRecord was already written by an earlier attempt.
|
|
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_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,
|
|
project_name: str | None = None,
|
|
description_prefix: str = "ChatAPI异步",
|
|
) -> BillingSummary:
|
|
"""Charge image/video generation fee safely before creating provider task."""
|
|
project_name = project_name or "AI生成任务"
|
|
gen_type = (gen_type or "").lower().strip()
|
|
items: list[BillingItem] = []
|
|
|
|
if gen_type == "image":
|
|
size = image_size or "2K"
|
|
amount = await calc_image_credits(db, size)
|
|
items.append(
|
|
await deduct_credits_locked_once(
|
|
db,
|
|
user_id=user_id,
|
|
amount=amount,
|
|
description=f"{description_prefix}图片生成 - {project_name}",
|
|
related_id=record_id,
|
|
charge_key=CHARGE_MEDIA_IMAGE,
|
|
)
|
|
)
|
|
elif gen_type == "video":
|
|
amount = await calc_video_credits(db, duration or 5, resolution or "720p")
|
|
items.append(
|
|
await deduct_credits_locked_once(
|
|
db,
|
|
user_id=user_id,
|
|
amount=amount,
|
|
description=f"{description_prefix}视频生成 - {project_name}",
|
|
related_id=record_id,
|
|
charge_key=CHARGE_MEDIA_VIDEO,
|
|
)
|
|
)
|
|
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 = "ChatAPI异步",
|
|
) -> 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,
|
|
duration=record.duration,
|
|
resolution=record.resolution,
|
|
project_name=project_name,
|
|
description_prefix=description_prefix,
|
|
)
|