272 lines
12 KiB
Python
272 lines
12 KiB
Python
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 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.utils.id_gen import generate_id
|
|
|
|
|
|
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.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,
|
|
) -> TokenUsage:
|
|
"""先持久化 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
|
|
|
|
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
|
|
return usage_row
|
|
|
|
|
|
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 = 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
|
|
record.token_usage_id = token_usage_id
|
|
ctx.token_usage_id = token_usage_id
|
|
|
|
|
|
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"
|