"""报告生成的 celery 任务(name=``meida.report.generate``,cpu 队列,T-P3-07/09)。 注册说明:``apps/worker/celery_app.py`` 的 include 目前只含 apps.worker.tasks (本批禁改文件),worker 进程要注册本模块需要把 include 补上——由协调者在 合并阶段处理;API 侧用 send_task 按任务名投递,两侧只靠名字衔接,互不 import。 worker 侧 LLM 网关装配:apps.api.deps 属 HTTP 层,worker 不得 import apps.api (apps 互不依赖契约),这里按同一配方自建;改网关装配时两处同步。 成本归因约定:阶段经 gateway 调用时以 ``ref_type="report"``、 ``ref_id=str(report_id)`` 记账——成功结案按这个口径聚合写 ``rp_report.model_info`` / ``cost_cents``(ob_ai_call 是唯一数据源,架构 §13)。 """ from __future__ import annotations import asyncio from datetime import UTC, datetime from decimal import Decimal from typing import Any from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import flag_modified from apps.worker.celery_app import celery_app from packages.billing.service import BillingService from packages.domain.task.state import TERMINAL_STATUSES, TaskStatus from packages.llm.budget import DailyBudgetGuard, RedisCostCounter from packages.llm.circuit import CircuitBreakerRegistry from packages.llm.cost import ( CompositeAICallRecorder, ModelRate, PeakAwareRateTable, RateTable, ) from packages.llm.gateway import LLMGateway from packages.llm.prompts import PromptRegistry from packages.llm.providers.base import LLMProvider from packages.llm.providers.deepseek import DeepSeekProvider from packages.llm.providers.zhipu import ZhipuGLMProvider from packages.llm.recording import PrometheusAICallRecorder, SqlAlchemyAICallRecorder from packages.llm.routing import SceneRouter, build_scene_configs from packages.report import repo as report_repo from packages.report.pipeline import ( ReportContext, StageFailureError, assemble_stages, load_reservation, run_pipeline, ) from packages.shared.config import Settings, get_settings from packages.shared.context import bind_context from packages.shared.db import get_session_factory from packages.shared.errors import AppError, ErrorCode from packages.shared.logging import get_logger from packages.shared.models import AiCall, RpReport, Task from packages.shared.redis_client import get_redis from packages.task import repo as task_repo from packages.task.service import TaskService logger = get_logger(__name__) def build_llm_gateway(settings: Settings) -> LLMGateway: """worker 侧装配,与 apps.api.deps.build_llm_gateway 同构(见模块 docstring)。""" providers: dict[str, LLMProvider] = { "deepseek": DeepSeekProvider( api_key=settings.deepseek_api_key, base_url=settings.deepseek_base_url, default_model=settings.deepseek_model, ), "zhipu": ZhipuGLMProvider( api_key=settings.zhipu_api_key, base_url=settings.zhipu_base_url, default_model=settings.zhipu_model, ), } default_rate = ModelRate.from_mapping(settings.llm_default_cost_rate) rates = PeakAwareRateTable( RateTable( { model: ModelRate.from_mapping(rate) for model, rate in settings.llm_cost_rates.items() }, default=default_rate, ), RateTable( { model: ModelRate.from_mapping(rate) for model, rate in settings.llm_cost_rates_peak.items() }, default=default_rate, ) if settings.llm_cost_rates_peak else None, enabled=settings.llm_peak_pricing_enabled, ) return LLMGateway( providers=providers, router=SceneRouter( build_scene_configs( primary_provider=settings.llm_primary_provider, backup_provider=settings.llm_backup_provider, structured_max_retries=settings.llm_structured_max_retries, overrides=settings.llm_scene_overrides, ) ), prompts=PromptRegistry( settings.prompts_dir, cache_prefix_mode=settings.llm_cache_prefix_mode ), rates=rates, recorder=CompositeAICallRecorder( SqlAlchemyAICallRecorder(get_session_factory()), PrometheusAICallRecorder(), ), budget=DailyBudgetGuard( RedisCostCounter(get_redis()), daily_budget_cents=Decimal(settings.llm_daily_budget_cents), enabled=settings.llm_daily_budget_enabled, tz_offset_hours=settings.llm_budget_timezone_offset_hours, tiers=list(settings.llm_budget_tiers), ), breakers=CircuitBreakerRegistry( failure_threshold=settings.llm_circuit_failure_threshold, cooldown_seconds=settings.llm_circuit_cooldown_seconds, ), ) async def _aggregate_cost(session: AsyncSession, report_id: int) -> tuple[Decimal, int]: """聚合该报告的 AI 成本与调用次数(口径见模块 docstring 的归因约定)。""" result = await session.execute( select(func.coalesce(func.sum(AiCall.cost_cents), 0), func.count()).where( AiCall.ref_type == "report", AiCall.ref_id == str(report_id) ) ) row = result.one() return Decimal(row[0]), int(row[1]) async def _settle_success( session: AsyncSession, task: Task, report: RpReport, settings: Settings, stage_names: list[str], ) -> None: """成功结案:commit 配额(幂等)+ succeeded + completed_at + 成本聚合。""" cost_cents, ai_calls = await _aggregate_cost(session, report.report_id) report.status = "succeeded" report.completed_at = datetime.now(UTC) report.cost_cents = cost_cents report.model_info = { "stages": stage_names, "vision_provider": settings.report_vision_provider, "ai_calls": ai_calls, # Decimal 不能直接进 JSONB;字符串保留精度,展示层再格式化。 "cost_cents": str(cost_cents), } reservation = await load_reservation(session, task.task_id) if reservation is not None: await BillingService(session).commit(reservation) # 任务中心视图同步收口(前端轮询的就是 tk_task,不是 rp_report)。 await TaskService(session).transition( task_id=task.task_id, target=TaskStatus.SUCCEEDED, stage_text="已完成", result={"report_id": str(report.report_id), "version": report.version}, ) await session.commit() logger.info( "report_succeeded", report_id=report.report_id, task_id=task.task_id, cost_cents=str(cost_cents), ai_calls=ai_calls, ) async def _generate_report(task_id: int) -> dict[str, Any]: """执行一条报告流水线(celery 包装层的内芯;集成测试直调的就是它)。""" settings = get_settings() factory = get_session_factory() async with factory() as session: task = await task_repo.get_for_update(session, task_id) if task is None: raise AppError(ErrorCode.NOT_FOUND, message=f"任务不存在:{task_id}") current = TaskStatus(task.status) if current in TERMINAL_STATUSES: # 重复投递/幂等重放的守卫:已结案的任务不再执行,防重复扣配额。 logger.info("report_task_replay_skipped", task_id=task_id, status=task.status) return {"task_id": str(task_id), "status": task.status, "skipped": True} report_id_raw = task.payload.get("report_id") if not isinstance(report_id_raw, str | int): raise AppError(ErrorCode.STATE_CONFLICT, message="任务缺少 report_id") report_id = int(report_id_raw) report = await report_repo.get_report(session, report_id) report_input = await report_repo.get_input(session, report_id) if report is None or report_input is None: raise AppError(ErrorCode.STATE_CONFLICT, message="报告数据不完整") service = TaskService(session) # 首跑是 PENDING;重试入站可能是 RUNNING(worker 崩溃重投)或 FAILED # (上一轮失败后重试),先归一到 QUEUED(状态机允许,且会维护 # retry_count)再进 RUNNING。 if current is not TaskStatus.QUEUED: await service.transition(task_id=task_id, target=TaskStatus.QUEUED, stage_text="排队中") await service.transition(task_id=task_id, target=TaskStatus.RUNNING, stage_text="分析中") report.status = "running" await session.commit() # RUNNING 先落库:前端进度与崩溃恢复都读它 gateway = build_llm_gateway(settings) try: stages = assemble_stages(settings) ctx = ReportContext( report=report, report_input=report_input, session=session, settings=settings, gateway=gateway, task_id=task_id, ) await run_pipeline(ctx, stages) await _settle_success(session, task, report, settings, [s.name for s in stages]) finally: await gateway.aclose() return {"task_id": str(task_id), "report_id": str(report_id), "status": "succeeded"} async def _fail_report(task_id: int, message: str) -> None: """基础设施异常重试耗尽后的兜底结案:failed + 退款(尽力而为)。 与 run_pipeline 的失败结案不同,这里没有阶段名——异常发生在流水线之外 (装配、DB 连接等),错误摘要以 worker 为责任方落库。结案失败只记日志: 此时额度停在 reserved,靠流水对账兜底。 """ factory = get_session_factory() try: async with factory() as session: task = await task_repo.get_for_update(session, task_id) if task is None or TaskStatus(task.status) in TERMINAL_STATUSES: return report_id_raw = task.payload.get("report_id") if isinstance(report_id_raw, str | int): report = await report_repo.get_report(session, int(report_id_raw)) if report is not None and report.status != "succeeded": report.status = "failed" report.result["error"] = { "stage": "worker", "type": "infrastructure", "message": message[:200], } # JSONB 深变更要显式标脏,否则 flush 按值比较会跳过 UPDATE。 flag_modified(report, "result") await TaskService(session).transition( task_id=task_id, target=TaskStatus.FAILED, stage_text="执行失败", error_message=message, ) reservation = await load_reservation(session, task_id) if reservation is not None: await BillingService(session).refund(reservation, reason="report_worker_failed") await session.commit() except Exception: logger.exception("report_fail_settle_failed", task_id=task_id) # Celery 的 task 装饰器没有类型标注,会把被装饰函数擦成 untyped @celery_app.task(name="meida.report.generate", bind=True, max_retries=1) # type: ignore[untyped-decorator] def generate_report(self: Any, task_id: int) -> dict[str, Any]: """cpu 队列的报告生成入口。 重试语义:网络/基础设施类异常重试一次——重跑时 ``derived_features`` 已落库,视觉阶段自然跳过(T-P3-09);阶段业务失败(StageFailureError) 已在流水线内结案(退款已提交),重跑也不会好,不再重试。 """ bind_context(task_id=task_id) try: return asyncio.run(_generate_report(task_id)) except StageFailureError as exc: logger.warning( "report_stage_failed", task_id=task_id, stage=exc.stage_name, error=f"{type(exc.original).__name__}: {exc.original}", ) raise except Exception as exc: if self.request.retries < self.max_retries: logger.warning("report_generate_retry", task_id=task_id, error=str(exc)) raise self.retry(exc=exc, countdown=10) from exc logger.exception("report_generate_failed_after_retry", task_id=task_id) asyncio.run(_fail_report(task_id, f"{type(exc).__name__}: {exc}")) raise