254 lines
10 KiB
Python
254 lines
10 KiB
Python
"""报告流水线的阶段契约与编排（第四批 T15/T16/T17 的共同接口）。
|
||
|
||
阶段契约（``ReportContext`` / ``Stage`` / ``DEFAULT_STAGE_PATHS``）由协调提交
|
||
定死，T15 实现编排（``assemble_stages`` / ``run_pipeline``）。T16 提供
|
||
vision_stage / diagnosis_stage，T17 提供 copy_stage / validation_stage——
|
||
三个任务互不改对方文件，靠 Settings.report_pipeline_stages 里的 dotted
|
||
path 动态装配。
|
||
|
||
设计动机（为什么动态 import）：三个阶段模块由三个并行分支各自创建，
|
||
静态 import 会在任一分支缺文件时炸掉编排；按配置动态装配让每个分支
|
||
独立可测，合并后自动咬合。代价是 import-linter 看不到这条边——阶段模块
|
||
自身的依赖仍受契约约束，编排层注入的 session/gateway 都是既有合法依赖。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import importlib
|
||
from collections.abc import Sequence
|
||
from typing import Any, Protocol, runtime_checkable
|
||
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy.orm.attributes import flag_modified
|
||
|
||
from packages.billing import repo as billing_repo
|
||
from packages.billing.service import BillingService, Reservation
|
||
from packages.domain.task.state import TaskStatus
|
||
from packages.llm.gateway import LLMGateway
|
||
from packages.shared.config import Settings
|
||
from packages.shared.errors import AppError, ErrorCode
|
||
from packages.shared.logging import get_logger
|
||
from packages.shared.models import Ledger, RpInput, RpReport
|
||
from packages.task import repo as task_repo
|
||
from packages.task.service import TaskService
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class ReportContext:
|
||
"""一个阶段的执行上下文。
|
||
|
||
不是 frozen：阶段要写 ``ctx.report.result[stage.name]`` 与
|
||
``ctx.rp_input.derived_features``（中间态落库由编排层 flush）。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
report: RpReport,
|
||
report_input: RpInput,
|
||
session: AsyncSession,
|
||
settings: Settings,
|
||
gateway: LLMGateway | None = None,
|
||
task_id: int | None = None,
|
||
) -> None:
|
||
self.report = report
|
||
self.report_input = report_input
|
||
self.session = session
|
||
self.settings = settings
|
||
self.gateway = gateway
|
||
# 编排层专用：失败结案要推动 tk_task。阶段实现不需要感知它，
|
||
# 保持可选以保证旧构造调用（T16/T17 的测试）不受影响。
|
||
self.task_id = task_id
|
||
|
||
|
||
@runtime_checkable
|
||
class Stage(Protocol):
|
||
"""流水线阶段。name 同时是 result JSONB 里的产物键。
|
||
|
||
约定：阶段把自己的全部产物写进 ``ctx.report.result[self.name]``；
|
||
抛异常即整条流水线失败（编排层负责 tk_task→failed + refund）。
|
||
"""
|
||
|
||
name: str
|
||
|
||
async def run(self, ctx: ReportContext) -> None: ...
|
||
|
||
|
||
#: T15 的 assemble_stages 按此清单 importlib 动态导入；缺失即 CONFIG_ERROR。
|
||
DEFAULT_STAGE_PATHS: tuple[str, ...] = (
|
||
"packages.report.vision_stage:VisionStage",
|
||
"packages.report.diagnosis_stage:DiagnosisStage",
|
||
"packages.report.copy_stage:CopyStage",
|
||
"packages.report.validation_stage:ValidationStage",
|
||
)
|
||
|
||
|
||
def stage_names(result: dict[str, Any]) -> list[str]:
|
||
"""result JSONB 里已出现的阶段键（诊断/调试用的小工具）。"""
|
||
return [k for k in result if isinstance(result[k], dict)]
|
||
|
||
|
||
class StageFailureError(Exception):
|
||
"""阶段异常且已按失败结案（rp_report failed + tk_task failed + 退款已提交）。
|
||
|
||
celery 包装层据此区分两类异常：``StageFailureError`` 是业务性失败，重跑
|
||
也不会好，不再重试；其它异常（DB/网络/装配）发生在结算之外，值得用
|
||
重试换恢复。
|
||
"""
|
||
|
||
def __init__(self, stage_name: str, original: BaseException) -> None:
|
||
self.stage_name = stage_name
|
||
self.original = original
|
||
super().__init__(f"stage={stage_name} error={type(original).__name__}: {original}")
|
||
|
||
|
||
def assemble_stages(settings: Settings, *, paths: Sequence[str] | None = None) -> list[Stage]:
|
||
"""按阶段配置（dotted path 清单）动态装配阶段列表。
|
||
|
||
默认读 ``settings.report_pipeline_stages``（basic 报告）；tier=full 的
|
||
深度诊断由调用方显式传 ``settings.report_full_pipeline_stages``——两条
|
||
流水线共用装配与失败语义，只是阶段清单不同。
|
||
|
||
动态 import 是刻意的插件式装配（见模块 docstring）：T16/T17 各自建阶段
|
||
模块，不碰本文件；import-linter 看不到这条边，阶段模块自身的依赖仍受
|
||
分层契约约束，这里只做装配不做业务。
|
||
|
||
模块/类不存在、实例化失败、或不符合 Stage 协议时，收集**全部**缺失项
|
||
抛 ``CONFIG_ERROR`` —— 绝不静默跳过：静默等于流水线空跑还扣用户配额。
|
||
"""
|
||
stage_paths = list(paths) if paths is not None else list(settings.report_pipeline_stages)
|
||
if not stage_paths:
|
||
raise AppError(
|
||
ErrorCode.CONFIG_ERROR,
|
||
message="报告流水线阶段配置为空",
|
||
data={"missing": list(DEFAULT_STAGE_PATHS)},
|
||
)
|
||
|
||
missing: list[str] = []
|
||
stages: list[Stage] = []
|
||
for dotted in stage_paths:
|
||
module_path, sep, attr = dotted.partition(":")
|
||
stage: Any = None
|
||
if sep and module_path and attr:
|
||
try:
|
||
# 阶段构造约定无参。__init__ 抛错视同配置坏了（比如阶段自身
|
||
# 校验它的依赖），归入 CONFIG_ERROR 而不是让 worker 崩溃。
|
||
stage = getattr(importlib.import_module(module_path), attr)()
|
||
except Exception:
|
||
stage = None
|
||
if stage is None or not isinstance(stage, Stage):
|
||
missing.append(dotted)
|
||
else:
|
||
stages.append(stage)
|
||
|
||
if missing:
|
||
raise AppError(
|
||
ErrorCode.CONFIG_ERROR,
|
||
message="报告流水线阶段装配失败：缺失或不可用",
|
||
data={"missing": missing},
|
||
)
|
||
return stages
|
||
|
||
|
||
def reservation_from_ledger(ledger: Ledger) -> Reservation:
|
||
"""从 reserved 流水重建 Reservation（worker 成功 commit / 失败 refund 都要它）。"""
|
||
return Reservation(
|
||
tx_id=ledger.tx_id,
|
||
user_id=ledger.user_id,
|
||
quota_key=ledger.quota_key or "",
|
||
period_key=ledger.period_key or "",
|
||
amount=abs(ledger.delta),
|
||
remaining=ledger.balance_after,
|
||
idempotency_key=ledger.idempotency_key,
|
||
)
|
||
|
||
|
||
async def load_reservation(session: AsyncSession, task_id: int) -> Reservation | None:
|
||
"""按 ``tk_task.reserve_tx_id`` 取回预扣；无预扣或流水缺失返回 None。"""
|
||
task = await task_repo.get(session, task_id)
|
||
if task is None or task.reserve_tx_id is None:
|
||
return None
|
||
ledger = await billing_repo.get_ledger(session, task.reserve_tx_id)
|
||
if ledger is None:
|
||
return None
|
||
return reservation_from_ledger(ledger)
|
||
|
||
|
||
async def run_pipeline(ctx: ReportContext, stages: list[Stage]) -> None:
|
||
"""逐阶段执行，每阶段产物即刻 flush 落库（T-P3-09 中间态持久化）。
|
||
|
||
中间态的意义：阶段半程失败或 worker 崩溃后重试，已完成阶段（尤其
|
||
``rp_input.derived_features`` 里的视觉结果）不必重算——「失败不需重新
|
||
上传照片」的结构保证就落在这几次 flush 里。
|
||
|
||
阶段异常 → ``_settle_stage_failure``（事务纪律见其 docstring）后包成
|
||
:class:`StageFailureError` 上抛；正常返回即全部阶段成功，成功结案归调用方
|
||
（worker：commit 配额 + succeeded + 成本聚合）。
|
||
"""
|
||
# 防御：default=dict 只在 INSERT 时生效，测试/手工构造的瞬态对象可能为 None。
|
||
if ctx.report.result is None:
|
||
ctx.report.result = {}
|
||
if ctx.report_input.derived_features is None:
|
||
ctx.report_input.derived_features = {}
|
||
|
||
for stage in stages:
|
||
try:
|
||
await stage.run(ctx)
|
||
except Exception as exc:
|
||
await _settle_stage_failure(ctx, stage.name, exc)
|
||
raise StageFailureError(stage.name, exc) from exc
|
||
# JSONB 列没有 Mutable 追踪：阶段对 dict 的深变更 ORM 看不见。且不能
|
||
# 靠「重新赋一个相等的新 dict」——flush 按值比较会认为没变而跳过
|
||
# UPDATE。flag_modified 强制该列进入下一条 UPDATE（中间态逐段落库
|
||
# 的关键）。derived_features 同理。
|
||
ctx.report.result.setdefault(stage.name, {})
|
||
flag_modified(ctx.report, "result")
|
||
flag_modified(ctx.report_input, "derived_features")
|
||
await ctx.session.flush()
|
||
logger.info(
|
||
"report_stage_done",
|
||
report_id=ctx.report.report_id,
|
||
task_id=ctx.task_id,
|
||
stage=stage.name,
|
||
)
|
||
|
||
|
||
async def _settle_stage_failure(ctx: ReportContext, stage_name: str, exc: Exception) -> None:
|
||
"""阶段失败的结案：report failed + tk_task failed + 退款，提交后才返回。
|
||
|
||
事务纪律（对齐 AgentService）：退款必须先于异常上抛提交，否则外层回滚
|
||
会把退款一起卷走（FR-BILL-05）。结案本身失败只记日志、不吞原始异常——
|
||
此时额度挂在 reserved 状态，靠流水对账兜底。
|
||
"""
|
||
summary = {"stage": stage_name, "type": type(exc).__name__, "message": str(exc)[:200]}
|
||
try:
|
||
ctx.report.status = "failed"
|
||
ctx.report.result["error"] = summary
|
||
flag_modified(ctx.report, "result") # 深变更必须显式标脏，见 run_pipeline 的注释
|
||
if ctx.task_id is not None:
|
||
await TaskService(ctx.session).transition(
|
||
task_id=ctx.task_id,
|
||
target=TaskStatus.FAILED,
|
||
stage_text=f"阶段失败：{stage_name}",
|
||
error_message=f"{type(exc).__name__}: {exc}",
|
||
)
|
||
reservation = await load_reservation(ctx.session, ctx.task_id)
|
||
if reservation is not None:
|
||
await BillingService(ctx.session).refund(
|
||
reservation, reason=f"report_stage_failed:{stage_name}"
|
||
)
|
||
await ctx.session.commit()
|
||
logger.warning(
|
||
"report_stage_failed",
|
||
report_id=ctx.report.report_id,
|
||
task_id=ctx.task_id,
|
||
stage=stage_name,
|
||
error_type=type(exc).__name__,
|
||
)
|
||
except Exception:
|
||
await ctx.session.rollback()
|
||
logger.exception(
|
||
"report_failure_settle_failed", report_id=ctx.report.report_id, stage=stage_name
|
||
)
|