153 lines
6.4 KiB
Python
153 lines
6.4 KiB
Python
"""圆桌后台作业执行器（Phase 3a）。
|
||
|
||
把 Phase 2 的编排引擎（``run_orchestration``）作为**后台 asyncio 任务**跑起来，
|
||
survive 发起请求的生命周期。每个任务：
|
||
|
||
1. ``set_current_user``：用作业的 user_id 设进程内用户上下文，让引擎里的进度落库
|
||
（``JobProgressWriter`` → ``RoundtableJobRepository.update_progress(user_id=AUTO)``）
|
||
正确归属（与 scheduled task 后台执行同一套做法）；
|
||
2. 用 ``gateway_factory`` 构造一个 ``RoundtableGateway``（Phase 3a 默认模拟网关，
|
||
Phase 3b 换成进程内真网关），跑完整编排（循环 → 共识 → 报告）。
|
||
|
||
任务句柄登记在 ``_tasks`` 里，供 ``cancel`` 用；完成后自动摘除。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from typing import Callable
|
||
|
||
from deerflow.agents.roundtable_orchestrator import (
|
||
JobProgressWriter,
|
||
SeatRef,
|
||
run_orchestration,
|
||
)
|
||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class _JobUser:
|
||
"""最小 CurrentUser：满足 user_context 的 ``.id`` 结构协议。"""
|
||
|
||
def __init__(self, user_id: str) -> None:
|
||
self.id = user_id
|
||
|
||
|
||
@dataclass
|
||
class StartParams:
|
||
job_id: str
|
||
user_id: str
|
||
agents: list[SeatRef]
|
||
seed_message: str
|
||
model: str | None = None
|
||
mode: str = "recommend"
|
||
intent_text: str = ""
|
||
thread_ids: dict[str, str] | None = None
|
||
coordinator_name: str = "roundtable-coordinator"
|
||
# 关联草稿：作业完成后把 step3 报告写回该草稿，使「查看报告」可用。
|
||
draft_id: str | None = None
|
||
|
||
|
||
# 网关工厂：给定一次启动参数，返回一个 RoundtableGateway 实例。
|
||
GatewayFactory = Callable[[StartParams], object]
|
||
|
||
|
||
class RoundtableJobExecutor:
|
||
def __init__(self, store, *, gateway_factory: GatewayFactory, draft_store=None) -> None:
|
||
self._store = store
|
||
self._gateway_factory = gateway_factory
|
||
# 草稿仓储（可选）：作业完成后把 step3 写回草稿，供前端「查看报告」。
|
||
self._draft_store = draft_store
|
||
self._tasks: dict[str, asyncio.Task] = {}
|
||
|
||
def start_job(self, params: StartParams) -> bool:
|
||
"""启动一个后台编排任务。同一 job 已在跑则忽略（幂等）。返回是否新启动。"""
|
||
if params.job_id in self._tasks and not self._tasks[params.job_id].done():
|
||
return False
|
||
task = asyncio.create_task(self._run(params))
|
||
self._tasks[params.job_id] = task
|
||
task.add_done_callback(lambda t, jid=params.job_id: self._tasks.pop(jid, None))
|
||
return True
|
||
|
||
async def _run(self, params: StartParams) -> None:
|
||
token = set_current_user(_JobUser(params.user_id))
|
||
try:
|
||
gateway = self._gateway_factory(params)
|
||
thread_ids = params.thread_ids or {}
|
||
coordinator_name = params.coordinator_name
|
||
# 真网关需要先进程内建好 leader/seat 线程；模拟网关无 init_threads。
|
||
# 仅首启建线程（resume 时 params.thread_ids 已带，复用、不重建）。
|
||
init = getattr(gateway, "init_threads", None)
|
||
if init is not None and not thread_ids:
|
||
thread_ids, coordinator_name = await init()
|
||
try:
|
||
await self._store.update_progress(
|
||
params.job_id, thread_ids=thread_ids, coordinator_name=coordinator_name
|
||
)
|
||
except Exception: # noqa: BLE001
|
||
logger.warning("persist thread_ids for job %s failed", params.job_id, exc_info=True)
|
||
|
||
progress = JobProgressWriter(self._store, job_id=params.job_id, agents=params.agents)
|
||
result = await run_orchestration(
|
||
gateway=gateway,
|
||
progress=progress,
|
||
agents=params.agents,
|
||
thread_ids=thread_ids,
|
||
coordinator_name=coordinator_name,
|
||
model=params.model,
|
||
seed_message=params.seed_message,
|
||
mode=params.mode,
|
||
intent_text=params.intent_text,
|
||
)
|
||
# 完成且有报告：写回关联草稿的 step3 + 推进 furthest_step=3，
|
||
# 这样用户从历史「查看报告」加载草稿即可看到最终方案（沿用既有 hydrateStep3 路径）。
|
||
if (
|
||
result.status == "done"
|
||
and result.step3
|
||
and params.draft_id
|
||
and self._draft_store is not None
|
||
):
|
||
try:
|
||
await self._draft_store.update_draft(
|
||
params.draft_id,
|
||
params.user_id,
|
||
step3=result.step3,
|
||
furthest_step=3,
|
||
)
|
||
except Exception: # noqa: BLE001 — 写报告失败不影响作业本身已 done
|
||
logger.warning("persist step3 to draft %s failed", params.draft_id, exc_info=True)
|
||
except asyncio.CancelledError:
|
||
# 干净停止即可，**不在被取消的任务里写 DB**：取消态下二次 await 会再抛
|
||
# CancelledError、并可能打断进行中的 sqlite 连接。``cancelled`` 终态由
|
||
# 路由的 cancel 端点（executor.cancel + store.update_progress）权威写入。
|
||
logger.info("roundtable job %s cancelled", params.job_id)
|
||
raise
|
||
except Exception: # noqa: BLE001 — 顶层兜底：run_orchestration 内部已尽量标 error
|
||
logger.exception("roundtable job %s crashed", params.job_id)
|
||
finally:
|
||
reset_current_user(token)
|
||
|
||
def cancel(self, job_id: str) -> bool:
|
||
"""请求取消一个在跑的作业。返回是否命中在跑任务。"""
|
||
task = self._tasks.get(job_id)
|
||
if task is None or task.done():
|
||
return False
|
||
task.cancel()
|
||
return True
|
||
|
||
def is_running(self, job_id: str) -> bool:
|
||
task = self._tasks.get(job_id)
|
||
return task is not None and not task.done()
|
||
|
||
async def await_job(self, job_id: str) -> None:
|
||
"""等待某作业任务结束（测试用；生产不调用）。"""
|
||
task = self._tasks.get(job_id)
|
||
if task is not None:
|
||
try:
|
||
await task
|
||
except (Exception, asyncio.CancelledError): # noqa: BLE001 — 取消是 BaseException
|
||
pass
|