303 lines
13 KiB
Python
303 lines
13 KiB
Python
"""报告生成的 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
|
||
|
||
from datetime import UTC, datetime
|
||
from decimal import Decimal
|
||
from typing import Any, cast
|
||
|
||
from sqlalchemy import func, select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy.orm.attributes import flag_modified
|
||
|
||
from apps.worker import loop as worker_loop
|
||
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:
|
||
if report.tier == "full":
|
||
# 深度诊断：复用 basic 报告的 derived_features，无视觉阶段，
|
||
# 文案/校验换 full 专用阶段（产出 result.full 四大块）。
|
||
stages = assemble_stages(settings, paths=settings.report_full_pipeline_stages)
|
||
else:
|
||
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 cast("dict[str, Any]", worker_loop.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)
|
||
worker_loop.run(_fail_report(task_id, f"{type(exc).__name__}: {exc}"))
|
||
raise
|