from __future__ import annotations from typing import Any, Mapping from sqlalchemy import func, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from app.enums.llm_billing import LlmBillingDomain, LlmBillingEvent, LlmBillingExecutionStatus, LlmCallAttemptStatus 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.token_usage import TokenUsage from app.services.credit.utils import utc_now from app.services.llm_billing.context import LlmBillingContext, LlmBillingStateError from app.services.operation_log_service import log_operation_event from app.utils.id_gen import generate_id def _audit_log(ctx: LlmBillingContext, event: LlmBillingEvent, **detail: Any) -> None: log_operation_event( domain=LlmBillingDomain.LLM_BILLING.value, module=ctx.source_module or "llm", event_type=event.value, event_status="success", 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=ctx.description_prefix or ctx.billing_scene or "LLM调用审计", detail={ "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, "call_attempt_id": ctx.current_call_attempt_id, **detail, }, ) def _usage_int(usage: Mapping[str, Any] | None, *keys: str) -> int: for key in keys: value = (usage or {}).get(key) if value is not None: try: return max(0, int(value)) except (TypeError, ValueError): pass return 0 async def create_call_attempt( ctx: LlmBillingContext, *, detail: Mapping[str, Any] | None = None, ) -> str: async with AsyncSessionLocal() as db: async with db.begin(): result = await db.execute( select(LlmBillingExecution) .where(LlmBillingExecution.id == ctx.billing_execution_id) .limit(1) .with_for_update() ) execution = result.scalar_one_or_none() if execution is None: raise LlmBillingStateError("LLM积分消费执行记录不存在，禁止调用模型") if execution.model_config_id and ctx.model_config_id and execution.model_config_id != ctx.model_config_id: raise LlmBillingStateError("自动重试模型与首次选定模型不一致，已拦截降级/切换") if execution.status not in { LlmBillingExecutionStatus.CHARGED.value, LlmBillingExecutionStatus.PRE_DEDUCTED.value, # 历史兼容 LlmBillingExecutionStatus.PROCESSING.value, }: raise LlmBillingStateError(f"LLM执行状态不允许调用模型：{execution.status}") max_result = await db.execute( select(func.coalesce(func.max(LlmCallAttempt.call_sequence), 0)).where( LlmCallAttempt.billing_execution_id == execution.id ) ) sequence = int(max_result.scalar_one() or 0) + 1 attempt = LlmCallAttempt( id=generate_id(), billing_execution_id=execution.id, call_sequence=sequence, retry_sequence=max(0, sequence - 1), model_config_id=execution.model_config_id or ctx.model_config_id, model_name_snapshot=execution.model_name_snapshot or ctx.model_name, provider_snapshot=execution.provider_snapshot or ctx.provider, request_started_at=utc_now(), status=LlmCallAttemptStatus.STARTED.value, postprocess_status="pending", ) db.add(attempt) execution.status = LlmBillingExecutionStatus.PROCESSING.value execution.total_call_count += 1 ctx.current_call_attempt_id = attempt.id # 同一个业务 execution 可能有多次供应商调用；每个新调用必须清空上一调用的 # 成功事实和用量快照，避免下一次真实 provider failure 被误判为后处理失败。 ctx.provider_call_succeeded = False ctx.provider_usage_snapshot = None ctx.token_usage_id = None return attempt.id async def _get_or_create_token_usage( db: AsyncSession, *, ctx: LlmBillingContext, attempt: LlmCallAttempt, usage: Mapping[str, Any] | None, input_tokens: int, output_tokens: int, total_tokens: int, ) -> tuple[TokenUsage, bool]: """先持久化 TokenUsage，再允许任何外键引用它。 LlmCallAttempt 仅保存 token_usage_id 字符串，没有 ORM relationship。SQLAlchemy 无法仅凭字符串赋值推导 INSERT/UPDATE 顺序，所以必须显式 flush TokenUsage。 同时按 user_id + biz_key 复用记录，保证审计重试幂等。 """ biz_key = f"llm-call:{attempt.id}" existing_result = await db.execute( select(TokenUsage) .where(TokenUsage.user_id == ctx.user_id, TokenUsage.biz_key == biz_key) .limit(1) ) existing = existing_result.scalar_one_or_none() if existing is not None: return existing, False usage_row = TokenUsage( id=generate_id(), model_config_id=(usage or {}).get("model_config_id") or attempt.model_config_id, user_id=ctx.user_id, input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=total_tokens, owner_type=ctx.owner_type, owner_id=ctx.owner_id, biz_key=biz_key, source_module=ctx.source_module, source_step_code=ctx.source_step_code, ) try: async with db.begin_nested(): db.add(usage_row) # 关键：先确保 token_usage 已 INSERT，后续 attempt.token_usage_id UPDATE # 才不会违反 llm_call_attempts_token_usage_id_fkey。 await db.flush([usage_row]) except IntegrityError: # 并发或审计重试可能已经创建同一 biz_key，回查复用。 existing_result = await db.execute( select(TokenUsage) .where(TokenUsage.user_id == ctx.user_id, TokenUsage.biz_key == biz_key) .limit(1) ) existing = existing_result.scalar_one_or_none() if existing is None: raise return existing, False return usage_row, True async def finish_call_success( ctx: LlmBillingContext, *, usage: Mapping[str, Any] | None = None, ) -> None: if not ctx.current_call_attempt_id: await create_call_attempt(ctx) now = utc_now() input_tokens = _usage_int(usage, "input_tokens", "prompt_tokens") output_tokens = _usage_int(usage, "output_tokens", "completion_tokens") total_tokens = _usage_int(usage, "total_tokens") or input_tokens + output_tokens token_usage_id: str | None = None async with AsyncSessionLocal() as db: async with db.begin(): result = await db.execute( select(LlmCallAttempt) .where(LlmCallAttempt.id == ctx.current_call_attempt_id) .limit(1) .with_for_update() ) attempt = result.scalar_one_or_none() if attempt is None: raise LlmBillingStateError("LLM调用审计记录不存在") if attempt.status == LlmCallAttemptStatus.SUCCEEDED.value: ctx.token_usage_id = attempt.token_usage_id return previous_status = attempt.status usage_row, token_usage_created = await _get_or_create_token_usage( db, ctx=ctx, attempt=attempt, usage=usage, input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=total_tokens, ) token_usage_id = usage_row.id attempt.status = LlmCallAttemptStatus.SUCCEEDED.value attempt.response_received_at = now attempt.duration_ms = max(0, int((now - attempt.request_started_at).total_seconds() * 1000)) attempt.input_tokens = input_tokens attempt.output_tokens = output_tokens attempt.total_tokens = total_tokens attempt.token_usage_id = token_usage_id attempt.provider_request_id = (usage or {}).get("provider_request_id") or (usage or {}).get("request_id") try: attempt.http_status = int((usage or {}).get("http_status")) if (usage or {}).get("http_status") is not None else None except (TypeError, ValueError): attempt.http_status = None # 支持把此前因审计异常误记的 failed 调用恢复为真实 succeeded。 attempt.error_message = None attempt.token_unavailable_reason = None execution_result = await db.execute( select(LlmBillingExecution) .where(LlmBillingExecution.id == attempt.billing_execution_id) .limit(1) .with_for_update() ) execution = execution_result.scalar_one() if previous_status in {LlmCallAttemptStatus.FAILED.value, LlmCallAttemptStatus.TIMEOUT.value}: execution.failed_call_count = max(0, int(execution.failed_call_count or 0) - 1) execution.successful_call_count += 1 execution.total_input_tokens += input_tokens execution.total_output_tokens += output_tokens execution.total_tokens += total_tokens 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 ctx.token_usage_id = token_usage_id _audit_log( ctx, LlmBillingEvent.TOKEN_USAGE_CREATED if token_usage_created else LlmBillingEvent.TOKEN_USAGE_REUSED, token_usage_id=token_usage_id, input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=total_tokens, ) async def finish_call_failure(ctx: LlmBillingContext, *, error: str, status: str = "failed") -> str: if not ctx.current_call_attempt_id: await create_call_attempt(ctx) now = utc_now() async with AsyncSessionLocal() as db: async with db.begin(): result = await db.execute( select(LlmCallAttempt) .where(LlmCallAttempt.id == ctx.current_call_attempt_id) .limit(1) .with_for_update() ) attempt = result.scalar_one_or_none() if attempt is None: return "missing" if attempt.status == LlmCallAttemptStatus.SUCCEEDED.value: # 供应商已成功并已记录 Token；后续解析/校验/落库失败只能记为后处理失败， # 不能覆盖真实供应商成功事实，也不能重复累计失败调用次数。 attempt.postprocess_status = "failed" attempt.postprocess_error = str(error)[:5000] return "postprocess_failure" if attempt.status != LlmCallAttemptStatus.STARTED.value: return "noop" attempt.status = LlmCallAttemptStatus.TIMEOUT.value if status == "timeout" else LlmCallAttemptStatus.FAILED.value attempt.response_received_at = now attempt.duration_ms = max(0, int((now - attempt.request_started_at).total_seconds() * 1000)) attempt.error_message = str(error)[:5000] attempt.token_unavailable_reason = "供应商调用失败，未返回Token使用量" execution_result = await db.execute( select(LlmBillingExecution) .where(LlmBillingExecution.id == attempt.billing_execution_id) .limit(1) .with_for_update() ) execution = execution_result.scalar_one() execution.failed_call_count += 1 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.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 return "provider_failure"