from __future__ import annotations from dataclasses import replace from typing import Any, Iterable, Mapping from sqlalchemy import select, tuple_ from sqlalchemy.ext.asyncio import AsyncSession from app.enums.credit_record import ( CREDIT_RECORD_BILLING_SCENE_LABELS, CreditRecordAction, CreditRecordChargeKind, CreditRecordSubject, ) from app.enums.llm_billing import ( LlmBillingDomain, LlmBillingEvent, LlmBillingExecutionStatus, LlmBillingLedgerState, ) from app.models.base import AsyncSessionLocal from app.models.credit_record import CreditRecord from app.models.llm_billing.call_attempt import LlmCallAttempt from app.models.llm_billing.execution import LlmBillingExecution from app.models.model_config import ModelConfig from app.services.credit.locking import acquire_user_credit_lock from app.services.credit.utils import to_float, utc_now from app.services.credits import deduct_credits_result from app.services.generation.billing_service import BillingItem, BillingSummary from app.services.llm_billing.call_audit_service import ( create_call_attempt, finish_call_failure, finish_call_success, ) from app.services.llm_billing.config import get_llm_billing_policy from app.services.llm_billing.context import ( LlmBillingConfigurationError, LlmBillingContext, LlmBillingStateError, LlmProviderPostprocessError, LlmChargeResult, LlmChargeValidation, ) from app.services.operation_log_service import log_operation_event from app.utils.exceptions import InsufficientCreditsError from app.utils.id_gen import generate_id def _credit_subject(ctx: LlmBillingContext) -> str: if ctx.charge_kind == CreditRecordChargeKind.VIDEO_ANALYSIS.value: return CreditRecordSubject.ANALYSIS.value return CreditRecordSubject.TEXT.value def _scene_name(ctx: LlmBillingContext, policy_name: str | None = None) -> str: return ( policy_name or CREDIT_RECORD_BILLING_SCENE_LABELS.get(ctx.billing_scene or "") or ctx.description_prefix or ctx.billing_scene or "LLM功能" ) def _detail(ctx: LlmBillingContext, **extra: Any) -> dict[str, Any]: return { "scene_code": ctx.billing_scene, "owner_type": ctx.owner_type, "owner_id": ctx.owner_id, "attempt_no": ctx.attempt_no, "billing_execution_id": ctx.billing_execution_id, "model_config_id": ctx.model_config_id, "model_name": ctx.model_name, "provider": ctx.provider, **extra, } def _log(ctx: LlmBillingContext, event: LlmBillingEvent, *, status: str = "success", error: str | None = None, **detail: Any) -> None: log_operation_event( domain=LlmBillingDomain.LLM_BILLING.value, module=ctx.source_module or "llm", event_type=event.value, event_status=status, trace_id=ctx.trace_id, request_id=ctx.request_id, user_id=ctx.user_id, project_id=ctx.source_project_id, task_id=ctx.celery_task_id, step_id=ctx.source_step_id, message=_scene_name(ctx), detail=_detail(ctx, **detail), error=error, ) async def _select_fixed_model(db: AsyncSession, ctx: LlmBillingContext) -> None: if ctx.model_config_id: result = await db.execute(select(ModelConfig).where(ModelConfig.id == ctx.model_config_id).limit(1)) model = result.scalar_one_or_none() else: result = await db.execute( select(ModelConfig) .where(ModelConfig.is_active.is_(True), ModelConfig.deleted_at.is_(None), ModelConfig.provider != "mock") .order_by(ModelConfig.priority.desc(), ModelConfig.id.asc()) .limit(1) ) model = result.scalar_one_or_none() if model is None: raise LlmBillingConfigurationError("没有可用的真实LLM模型，且本版本禁止Mock或自动降级") ctx.model_config_id = model.id ctx.model_name = model.model_name ctx.provider = model.provider ctx.model_parameters_snapshot = { "model_config_id": model.id, "name": model.name, "provider": model.provider, "api_base": model.api_base, "model_name": model.model_name, "max_tokens": model.max_tokens, "temperature": model.temperature, } async def _find_execution(db: AsyncSession, ctx: LlmBillingContext, *, for_update: bool = False) -> LlmBillingExecution | None: stmt = ( select(LlmBillingExecution) .where( LlmBillingExecution.user_id == ctx.user_id, LlmBillingExecution.scene_code == (ctx.billing_scene or ""), LlmBillingExecution.owner_type == ctx.owner_type, LlmBillingExecution.owner_id == ctx.owner_id, LlmBillingExecution.business_attempt_no == ctx.attempt_no, ) .limit(1) # CallAttempt/TokenUsage 使用独立短事务更新 execution 聚合值；主 Session # 可能在 commit 前后仍持有同一对象。强制回填数据库最新标量，避免 # expire_on_commit=False 下 identity map 返回旧 Token/调用次数。 .execution_options(populate_existing=True) ) if for_update: stmt = stmt.with_for_update() result = await db.execute(stmt) return result.scalar_one_or_none() def _state_for_execution(execution: LlmBillingExecution | None) -> LlmBillingLedgerState: if execution is None: return LlmBillingLedgerState.MISSING mapping = { LlmBillingExecutionStatus.CHARGED.value: LlmBillingLedgerState.ACTIVE, LlmBillingExecutionStatus.PRE_DEDUCTED.value: LlmBillingLedgerState.ACTIVE, LlmBillingExecutionStatus.PROCESSING.value: LlmBillingLedgerState.ACTIVE, LlmBillingExecutionStatus.SUCCEEDED.value: LlmBillingLedgerState.SUCCEEDED, LlmBillingExecutionStatus.FINAL_FAILED.value: LlmBillingLedgerState.FINAL_FAILED, LlmBillingExecutionStatus.REFUNDED.value: LlmBillingLedgerState.REFUNDED, LlmBillingExecutionStatus.REFUND_FAILED.value: LlmBillingLedgerState.INVALID, } return mapping.get(execution.status, LlmBillingLedgerState.INVALID) async def charge_llm_credits(db: AsyncSession, ctx: LlmBillingContext) -> LlmChargeResult: """在业务 API 事务内按后台场景配置同步真实消费积分。""" _log(ctx, LlmBillingEvent.CHARGE_START, status="started") if not ctx.billing_scene: _log(ctx, LlmBillingEvent.CHARGE_CONFIG_INVALID, status="failed", error="billing_scene_missing") raise LlmBillingConfigurationError("LLM积分消费必须指定billing_scene") # 同一 user + scene + owner + attempt 的幂等判断必须位于用户积分锁之后， # 避免并发 API 请求同时看到 missing 后重复消费或撞 execution 唯一键。 await acquire_user_credit_lock(db, ctx.user_id) existing = await _find_execution(db, ctx, for_update=True) if existing: ctx.billing_execution_id = existing.id ctx.model_config_id = existing.model_config_id ctx.model_name = existing.model_name_snapshot ctx.provider = existing.provider_snapshot ctx.model_parameters_snapshot = existing.model_parameters_snapshot state = _state_for_execution(existing) _log( ctx, LlmBillingEvent.CHARGE_IDEMPOTENT_HIT, amount=to_float(existing.pre_deduct_credits), credit_record_id=existing.credit_record_id, execution_status=existing.status, ledger_state=state.value, ) return LlmChargeResult( amount=to_float(existing.pre_deduct_credits), state=state, created=False, record_id=existing.credit_record_id, execution_id=existing.id, ) await _select_fixed_model(db, ctx) policy = await get_llm_billing_policy(db, scene_code=ctx.billing_scene) if not policy.valid: _log(ctx, LlmBillingEvent.CHARGE_CONFIG_INVALID, status="failed", error=policy.error) raise LlmBillingConfigurationError(policy.error or "LLM场景积分配置无效") scene_name = _scene_name(ctx, policy.scene_name) meta = { "owner_type": ctx.owner_type, "owner_id": ctx.owner_id, "attempt_no": ctx.attempt_no, "charge_kind": ctx.charge_kind, "charge_action": CreditRecordAction.CHARGE.value, "credit_subject": _credit_subject(ctx), "billing_scene": ctx.billing_scene, "scene_name_snapshot": scene_name, "source_module": ctx.source_module, "source_project_id": ctx.source_project_id, "source_step_id": ctx.source_step_id, "source_step_code": ctx.source_step_code, "engine_type": "model", "engine_id": ctx.model_config_id, "engine_provider": ctx.provider, "engine_model_name": ctx.model_name, } try: mutation = await deduct_credits_result( db, user_id=ctx.user_id, amount=policy.charge_credits, description=f"{scene_name}积分消费", related_id=ctx.related_id or ctx.owner_id, biz_key=ctx.charge_biz_key, record_meta=meta, record_type="consume", request_time=ctx.request_time or utc_now(), ) except InsufficientCreditsError: _log( ctx, LlmBillingEvent.CHARGE_INSUFFICIENT, status="failed", error="insufficient_credits", amount=policy.charge_credits, ) raise if mutation.record is None: raise LlmBillingStateError("LLM积分消费流水创建失败") execution = LlmBillingExecution( id=generate_id(), user_id=ctx.user_id, scene_code=ctx.billing_scene, scene_name_snapshot=scene_name, owner_type=ctx.owner_type, owner_id=ctx.owner_id, business_attempt_no=ctx.attempt_no, model_config_id=ctx.model_config_id, model_name_snapshot=ctx.model_name, provider_snapshot=ctx.provider, model_parameters_snapshot=ctx.model_parameters_snapshot, billing_policy_id=policy.policy_id, billing_policy_version=policy.version, request_time=ctx.request_time or utc_now(), # 数据库物理字段仍沿用历史 pre_deduct_credits 命名；语义已经是本 attempt 真实消费积分。 pre_deduct_credits=policy.charge_credits, credit_record_id=mutation.record.id, status=LlmBillingExecutionStatus.CHARGED.value, ) db.add(execution) await db.flush() # flush 后只使用已加载标量，不访问 relationship，避免 commit 后懒加载/MissingGreenlet 风险。 mutation.record.llm_billing_execution_id = execution.id mutation.record.scene_name_snapshot = scene_name ctx.billing_execution_id = execution.id _log( ctx, LlmBillingEvent.CHARGE_SUCCESS, amount=policy.charge_credits, credit_record_id=mutation.record.id, created=mutation.created, ) return LlmChargeResult( amount=policy.charge_credits, state=LlmBillingLedgerState.ACTIVE, created=mutation.created, record_id=mutation.record.id, execution_id=execution.id, ) async def ensure_llm_charged(db: AsyncSession, ctx: LlmBillingContext) -> LlmChargeValidation: _log(ctx, LlmBillingEvent.EXECUTION_VALIDATE_START, status="started") execution = await _find_execution(db, ctx) if execution is None: _log( ctx, LlmBillingEvent.EXECUTION_BLOCKED, status="failed", error="charge_missing", ledger_state=LlmBillingLedgerState.MISSING.value, ) return LlmChargeValidation(False, 0.0, LlmBillingLedgerState.MISSING, "charge_missing") ctx.billing_execution_id = execution.id ctx.model_config_id = execution.model_config_id ctx.model_name = execution.model_name_snapshot ctx.provider = execution.provider_snapshot ctx.model_parameters_snapshot = execution.model_parameters_snapshot state = _state_for_execution(execution) can_execute = state == LlmBillingLedgerState.ACTIVE if not can_execute: _log(ctx, LlmBillingEvent.EXECUTION_BLOCKED, status="failed", execution_status=execution.status) else: _log(ctx, LlmBillingEvent.EXECUTION_VALIDATE_SUCCESS, execution_status=execution.status) return LlmChargeValidation( can_execute=can_execute, amount=to_float(execution.pre_deduct_credits), state=state, reason=None if can_execute else execution.status, charge_record_id=execution.credit_record_id, execution_id=execution.id, ) async def get_llm_ledger_states(db: AsyncSession, contexts: Iterable[LlmBillingContext]) -> dict[str, LlmChargeValidation]: items = list(contexts) if not items: return {} keys = list(dict.fromkeys( (ctx.user_id, ctx.billing_scene or "", ctx.owner_type, ctx.owner_id, ctx.attempt_no) for ctx in items )) result = await db.execute( select(LlmBillingExecution).where( tuple_( LlmBillingExecution.user_id, LlmBillingExecution.scene_code, LlmBillingExecution.owner_type, LlmBillingExecution.owner_id, LlmBillingExecution.business_attempt_no, ).in_(keys) ) ) by_key = { (row.user_id, row.scene_code, row.owner_type, row.owner_id, row.business_attempt_no): row for row in result.scalars().all() } output: dict[str, LlmChargeValidation] = {} for ctx in items: row = by_key.get((ctx.user_id, ctx.billing_scene or "", ctx.owner_type, ctx.owner_id, ctx.attempt_no)) state = _state_for_execution(row) output[ctx.charge_biz_key] = LlmChargeValidation( can_execute=state == LlmBillingLedgerState.ACTIVE, amount=to_float(row.pre_deduct_credits) if row else 0.0, state=state, reason=None if state == LlmBillingLedgerState.ACTIVE else (row.status if row else "missing"), charge_record_id=row.credit_record_id if row else None, execution_id=row.id if row else None, ) return output async def validate_retryable_previous_attempt(db: AsyncSession, ctx: LlmBillingContext) -> LlmChargeValidation: previous = replace(ctx, attempt_no=max(1, ctx.attempt_no - 1)) if ctx.attempt_no > 1 else ctx validation = (await get_llm_ledger_states(db, [previous]))[previous.charge_biz_key] if ctx.attempt_no <= 1: return LlmChargeValidation(True, 0.0, LlmBillingLedgerState.MISSING, "first_attempt") allowed = validation.state in { LlmBillingLedgerState.SUCCEEDED, LlmBillingLedgerState.REFUNDED, LlmBillingLedgerState.FINAL_FAILED, } return LlmChargeValidation(allowed, validation.amount, validation.state, None if allowed else "previous_attempt_not_final", validation.charge_record_id, validation.execution_id) async def log_provider_start(db: AsyncSession, ctx: LlmBillingContext, *, detail: Mapping[str, Any] | None = None) -> None: if not ctx.billing_execution_id: validation = await ensure_llm_charged(db, ctx) if not validation.can_execute: raise LlmBillingStateError("缺少有效LLM积分消费记录，禁止调用LLM") await create_call_attempt(ctx, detail=detail) _log(ctx, LlmBillingEvent.PROVIDER_START, status="started", call_attempt_id=ctx.current_call_attempt_id, detail=dict(detail or {})) async def log_provider_success(db: AsyncSession, ctx: LlmBillingContext, *, usage: Mapping[str, Any] | None = None) -> None: # 先记录内存中的供应商成功事实。即使独立审计事务异常，外层异常处理也 # 必须按“供应商成功、后处理失败”处理，不能覆盖成 provider failed。 ctx.provider_call_succeeded = True ctx.provider_usage_snapshot = dict(usage or {}) await finish_call_success(ctx, usage=usage) _log(ctx, LlmBillingEvent.PROVIDER_SUCCESS, call_attempt_id=ctx.current_call_attempt_id, usage=dict(usage or {})) async def log_provider_failure(db: AsyncSession, ctx: LlmBillingContext, *, error: str) -> None: if ctx.provider_call_succeeded: # 可能是首次成功审计落库失败后进入外层异常处理。先幂等补记真实成功， # 再把当前异常记为 postprocess failure。 await finish_call_success(ctx, usage=ctx.provider_usage_snapshot or {}) outcome = await finish_call_failure(ctx, error=error) event = LlmBillingEvent.POSTPROCESS_FAILURE if outcome == "postprocess_failure" else LlmBillingEvent.PROVIDER_FAILURE _log( ctx, event, status="failed", call_attempt_id=ctx.current_call_attempt_id, failure_stage="postprocess" if outcome == "postprocess_failure" else "provider", error=error, ) async def record_provider_exception( db: AsyncSession, ctx: LlmBillingContext, exc: Exception, ) -> tuple[bool, dict[str, Any]]: """记录供应商失败，或“供应商成功但后处理失败”的真实调用事实。""" if ctx.provider_call_succeeded: usage = dict(ctx.provider_usage_snapshot or {}) await log_provider_failure(db, ctx, error=str(exc)) return True, usage if isinstance(exc, LlmProviderPostprocessError): usage = dict(exc.usage or {}) await log_provider_success(db, ctx, usage=usage) await log_provider_failure(db, ctx, error=str(exc)) return True, usage await log_provider_failure(db, ctx, error=str(exc)) return False, {} def log_celery_dispatch_start(ctx: LlmBillingContext) -> None: _log(ctx, LlmBillingEvent.CELERY_DISPATCH_START, status="started") def log_celery_dispatch_success(ctx: LlmBillingContext) -> None: _log(ctx, LlmBillingEvent.CELERY_DISPATCH_SUCCESS) def log_celery_dispatch_failure(ctx: LlmBillingContext, *, error: str) -> None: _log(ctx, LlmBillingEvent.CELERY_DISPATCH_FAILURE, status="failed", error=error) def log_celery_dispatch_compensated(ctx: LlmBillingContext, *, error: str) -> None: _log(ctx, LlmBillingEvent.CELERY_DISPATCH_COMPENSATED, status="warning", error=error) async def mark_business_success( db: AsyncSession, ctx: LlmBillingContext, *, usage: Mapping[str, Any], description: str | None = None, ) -> BillingSummary: # 调用审计使用独立短事务，必须在当前 Session 锁定 execution 之前完成。 # 若先锁 execution 再等待独立 Session 更新同一行，会形成自锁死等待。 if ctx.current_call_attempt_id: await finish_call_success(ctx, usage=usage) execution = await _find_execution(db, ctx, for_update=True) if execution is None: raise LlmBillingStateError("LLM积分消费执行记录不存在") ctx.billing_execution_id = execution.id created = execution.status != LlmBillingExecutionStatus.SUCCEEDED.value if execution.status in {LlmBillingExecutionStatus.REFUNDED.value, LlmBillingExecutionStatus.FINAL_FAILED.value}: raise LlmBillingStateError("LLM业务已最终失败并退款，不能再标记成功") execution.status = LlmBillingExecutionStatus.SUCCEEDED.value execution.completed_at = utc_now() call_result = await db.execute( select(LlmCallAttempt) .where(LlmCallAttempt.billing_execution_id == execution.id) .order_by(LlmCallAttempt.call_sequence.desc()) .limit(1) ) last_call = call_result.scalar_one_or_none() if last_call: last_call.postprocess_status = "succeeded" last_call.postprocess_error = None record_result = await db.execute(select(CreditRecord).where(CreditRecord.id == execution.credit_record_id).limit(1).with_for_update()) record = record_result.scalar_one_or_none() if record: record.input_tokens = execution.total_input_tokens record.output_tokens = execution.total_output_tokens record.total_tokens = execution.total_tokens record.llm_call_count = execution.total_call_count record.llm_success_call_count = execution.successful_call_count record.llm_failed_call_count = execution.failed_call_count _log(ctx, LlmBillingEvent.BUSINESS_SUCCESS, charge_credits=to_float(execution.pre_deduct_credits), idempotent=not created) return BillingSummary( record_id=ctx.owner_id, user_id=ctx.user_id, items=[ BillingItem( charge_key=ctx.charge_kind, amount=to_float(execution.pre_deduct_credits), charged=created, skipped_reason=None if created else "already_succeeded", biz_key=ctx.billing_biz_key, attempt_no=ctx.attempt_no, ) ], ) async def refund_on_final_failure(db: AsyncSession, ctx: LlmBillingContext, *, error: str | None = None) -> BillingSummary: from app.services.credit.ledger_service import refund_consumption # 与真实消费保持 user advisory -> execution -> original record -> balance 的统一锁顺序。 await acquire_user_credit_lock(db, ctx.user_id) execution = await _find_execution(db, ctx, for_update=True) if execution is None: return BillingSummary(ctx.owner_id, ctx.user_id, []) ctx.billing_execution_id = execution.id if execution.status == LlmBillingExecutionStatus.REFUNDED.value: return BillingSummary( ctx.owner_id, ctx.user_id, [BillingItem(ctx.charge_kind, to_float(execution.pre_deduct_credits), False, "already_refunded", ctx.refund_biz_key, ctx.attempt_no)], ) original_result = await db.execute( select(CreditRecord) .where(CreditRecord.id == execution.credit_record_id, CreditRecord.user_id == ctx.user_id) .limit(1) .with_for_update() ) original_record = original_result.scalar_one_or_none() if original_record is None or not original_record.biz_key: execution.status = LlmBillingExecutionStatus.REFUND_FAILED.value execution.final_error_message = "原LLM积分消费流水不存在或缺少BizKey" execution.completed_at = utc_now() _log( ctx, LlmBillingEvent.FINAL_FAILURE_REFUND_FAILED, status="failed", error=execution.final_error_message, credit_record_id=execution.credit_record_id, ) raise LlmBillingStateError(execution.final_error_message) execution.status = LlmBillingExecutionStatus.FINAL_FAILED.value execution.final_error_message = str(error or "业务最终失败")[:5000] execution.completed_at = utc_now() original_biz_key = original_record.biz_key _log( ctx, LlmBillingEvent.FINAL_FAILURE_START, status="started", error=error, credit_record_id=original_record.id, original_biz_key=original_biz_key, ) try: refunded = await refund_consumption( db, user_id=ctx.user_id, refund_for_biz_key=original_biz_key, description=f"{execution.scene_name_snapshot}最终失败", related_id=ctx.related_id or ctx.owner_id, biz_key=f"{original_biz_key}:business-final-refund", record_meta={ "owner_type": ctx.owner_type, "owner_id": ctx.owner_id, "attempt_no": ctx.attempt_no, "charge_kind": ctx.charge_kind, "charge_action": CreditRecordAction.REFUND.value, "credit_subject": CreditRecordSubject.REFUND.value, "billing_scene": ctx.billing_scene, "scene_name_snapshot": execution.scene_name_snapshot, "source_module": ctx.source_module, "source_project_id": ctx.source_project_id, "source_step_id": ctx.source_step_id, "source_step_code": ctx.source_step_code, "llm_billing_execution_id": execution.id, }, refund_time=utc_now(), ) execution.refund_available_credits = refunded.available_amount execution.refund_expired_credits = refunded.expired_amount execution.status = LlmBillingExecutionStatus.REFUNDED.value execution.refunded_at = utc_now() _log( ctx, LlmBillingEvent.FINAL_FAILURE_REFUND_SUCCESS, refund_available=to_float(refunded.available_amount), refund_expired=to_float(refunded.expired_amount), credit_record_id=original_record.id, ) return BillingSummary( ctx.owner_id, ctx.user_id, [BillingItem(ctx.charge_kind, to_float(refunded.total_amount), True, None, ctx.refund_biz_key, ctx.attempt_no)], ) except Exception as exc: execution.status = LlmBillingExecutionStatus.REFUND_FAILED.value _log(ctx, LlmBillingEvent.FINAL_FAILURE_REFUND_FAILED, status="failed", error=str(exc), credit_record_id=original_record.id) raise async def finalize_llm_business_failure( ctx: LlmBillingContext, *, error: str | None = None, ) -> BillingSummary: """在独立短事务中完成最终失败退款，并确保退款失败状态不会被业务 rollback 吞掉。""" try: async with AsyncSessionLocal() as session: async with session.begin(): return await refund_on_final_failure(session, ctx, error=error) except Exception as exc: try: async with AsyncSessionLocal() as audit_session: async with audit_session.begin(): execution = await _find_execution(audit_session, ctx, for_update=True) if execution is not None and execution.status != LlmBillingExecutionStatus.REFUNDED.value: execution.status = LlmBillingExecutionStatus.REFUND_FAILED.value execution.final_error_message = str(error or exc)[:5000] execution.completed_at = utc_now() except Exception: _log(ctx, LlmBillingEvent.FINAL_FAILURE_REFUND_FAILED, status="failed", error=str(exc), persist_failed_status=False) raise