182 lines
7.4 KiB
Python
182 lines
7.4 KiB
Python
"""会商 Step2 性能改造引擎测试：研讨前并行『取数』阶段（gather_first）+ recommend 批内并行。
|
||
|
||
覆盖：
|
||
- gather_first=True：研讨前先对**全席位并行**跑一轮「只取数」，产出作为共享资料注入后续研讨轮席位任务；
|
||
- gather_first=False（默认）：不跑取数阶段，行为不变；
|
||
- recommend 一轮派多席位 → 批内**并行**执行（_run_recommend 由串行改 asyncio.gather，峰值并发>1）。
|
||
|
||
引擎纯逻辑，用 JobProgressWriter(store=None) 干跑 + 录制网关，不需要 DB / 真模型。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
|
||
import pytest
|
||
|
||
from deerflow.agents.roundtable_orchestrator import (
|
||
JobProgressWriter,
|
||
SeatRef,
|
||
run_orchestration,
|
||
)
|
||
from deerflow.agents.roundtable_orchestrator.types import LeaderTurn, SeatTurn
|
||
|
||
_AGENTS = [SeatRef("seat-a", "情报"), SeatRef("seat-b", "环境"), SeatRef("seat-c", "方案")]
|
||
# _GATHER_TASK 里的稳定标记（取数阶段）/ _seat_message 里注入资料块的稳定标记（研讨阶段）。
|
||
_GATHER_MARK = "资料收集阶段"
|
||
_GATHERED_BLOCK_MARK = "研讨前各席位已并行收集"
|
||
|
||
|
||
def _writer() -> JobProgressWriter:
|
||
return JobProgressWriter(None, job_id="job-x", agents=_AGENTS)
|
||
|
||
|
||
def _thread_ids() -> dict[str, str]:
|
||
return {a.agent_id: f"tid-{a.agent_id}" for a in _AGENTS} | {"roundtable-coordinator": "tid-co"}
|
||
|
||
|
||
class _RecGateway:
|
||
"""录制网关：run_leader 按预设各轮派活；run_seat 记录 (agent_id, task) 并追踪峰值并发。"""
|
||
|
||
def __init__(self, dispatch_cycles: list[list[tuple[str, str]]]) -> None:
|
||
self._cycles = dispatch_cycles
|
||
self._leader_calls = 0
|
||
self.leader_requests: list[tuple[str, bool]] = []
|
||
self.seat_tasks: list[tuple[str, str]] = []
|
||
self.active = 0
|
||
self.max_active = 0
|
||
|
||
async def run_leader(
|
||
self,
|
||
*,
|
||
thread_ids,
|
||
coordinator_name,
|
||
message,
|
||
model,
|
||
finalize_only: bool = False,
|
||
) -> LeaderTurn:
|
||
self.leader_requests.append((message, finalize_only))
|
||
i = self._leader_calls
|
||
self._leader_calls += 1
|
||
disp = self._cycles[i] if i < len(self._cycles) else []
|
||
return LeaderTurn(final_text=f"leader-{i}", dispatched=disp)
|
||
|
||
async def run_seat(self, *, thread_ids, agent_id, task, model) -> SeatTurn:
|
||
self.active += 1
|
||
self.max_active = max(self.max_active, self.active)
|
||
self.seat_tasks.append((agent_id, task))
|
||
await asyncio.sleep(0.01) # 留出并发重叠窗口
|
||
self.active -= 1
|
||
return SeatTurn(final_text=f"{agent_id} 的产出")
|
||
|
||
# Step3 关掉，不应被调用
|
||
async def create_report_thread(self) -> str: # pragma: no cover - 不应触发
|
||
raise AssertionError("enable_report=False 不应绘制报告")
|
||
|
||
|
||
async def _run(gw: _RecGateway, *, gather_first: bool, max_cycles: int = 8):
|
||
return await run_orchestration(
|
||
gateway=gw,
|
||
progress=_writer(),
|
||
agents=_AGENTS,
|
||
thread_ids=_thread_ids(),
|
||
coordinator_name="roundtable-coordinator",
|
||
model=None,
|
||
seed_message="背景X",
|
||
mode="recommend",
|
||
intent_text="背景X",
|
||
enable_report=False,
|
||
gather_first=gather_first,
|
||
max_cycles=max_cycles,
|
||
)
|
||
|
||
|
||
# 一轮派全部 3 席位 → 下一轮共识收口。
|
||
_DISPATCH_ALL_THEN_CONSENSUS = [[("seat-a", "taskA"), ("seat-b", "taskB"), ("seat-c", "taskC")], []]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_recommend_dispatch_runs_in_parallel():
|
||
"""recommend 一轮派多席位 → 批内并行（峰值并发 = 席位数，证明已由串行改并行）。"""
|
||
gw = _RecGateway(_DISPATCH_ALL_THEN_CONSENSUS)
|
||
result = await _run(gw, gather_first=False)
|
||
assert result.status == "done"
|
||
assert gw.max_active == 3 # 三席位真并行；若仍串行则恒为 1
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_gather_first_false_skips_gather_phase():
|
||
"""默认 gather_first=False：不跑取数阶段，run_seat 只在研讨轮被调，无取数任务。"""
|
||
gw = _RecGateway(_DISPATCH_ALL_THEN_CONSENSUS)
|
||
await _run(gw, gather_first=False)
|
||
assert not any(_GATHER_MARK in task for _, task in gw.seat_tasks)
|
||
# 研讨轮任务也不应带「已收集资料」块（没有取数）
|
||
assert not any(_GATHERED_BLOCK_MARK in task for _, task in gw.seat_tasks)
|
||
# 仅 3 次席位调用（一轮研讨派活）
|
||
assert len(gw.seat_tasks) == 3
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_gather_first_runs_parallel_gather_then_injects_material():
|
||
"""gather_first=True：研讨前对全席位并行跑「只取数」一轮，产出注入后续研讨轮任务。"""
|
||
gw = _RecGateway(_DISPATCH_ALL_THEN_CONSENSUS)
|
||
result = await _run(gw, gather_first=True)
|
||
assert result.status == "done"
|
||
|
||
gather_calls = [(aid, t) for aid, t in gw.seat_tasks if _GATHER_MARK in t]
|
||
discuss_calls = [(aid, t) for aid, t in gw.seat_tasks if _GATHER_MARK not in t]
|
||
|
||
# ① 取数阶段覆盖全部席位（每席位一次「只取数」任务）
|
||
assert {aid for aid, _ in gather_calls} == {"seat-a", "seat-b", "seat-c"}
|
||
# ② 取数任务在研讨任务之前（取数是前置阶段）
|
||
first_discuss_idx = gw.seat_tasks.index(discuss_calls[0])
|
||
last_gather_idx = max(gw.seat_tasks.index(c) for c in gather_calls)
|
||
assert last_gather_idx < first_discuss_idx
|
||
# ③ 研讨轮任务带上「研讨前各席位已并行收集的资料」块，且包含各席位取数产出
|
||
assert discuss_calls, "应有研讨轮任务"
|
||
for _, dtask in discuss_calls:
|
||
assert _GATHERED_BLOCK_MARK in dtask
|
||
assert "seat-a 的产出" in dtask # 注入了取数阶段的产物
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_gather_phase_is_parallel():
|
||
"""取数阶段本身并行（全席位并发跑，峰值并发 = 席位数）。"""
|
||
# 首轮 leader 直接共识（dispatched=[]），让研讨轮不抢占并发，单独度量取数阶段峰值。
|
||
gw = _RecGateway([[]])
|
||
await _run(gw, gather_first=True)
|
||
# 只有取数阶段跑了 run_seat（研讨轮因首轮共识没派活）
|
||
assert all(_GATHER_MARK in t for _, t in gw.seat_tasks)
|
||
assert len(gw.seat_tasks) == 3
|
||
assert gw.max_active == 3 # 取数阶段三席位真并行
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_recommend_eighth_round_forces_text_closeout_instead_of_stopping():
|
||
"""前 7 轮仍派活时，第 8 轮必须关闭工具并要求总控总结，不能直接停止。"""
|
||
gw = _RecGateway([[('seat-a', f'task-{i}')] for i in range(8)])
|
||
|
||
result = await _run(gw, gather_first=False)
|
||
|
||
assert result.status == "done"
|
||
assert result.consensus is True
|
||
assert len(gw.leader_requests) == 8
|
||
assert [finalize for _, finalize in gw.leader_requests] == [False] * 7 + [True]
|
||
assert "最终方案与结论正文" in gw.leader_requests[-1][0]
|
||
assert "不要调用任何工具" in gw.leader_requests[-1][0]
|
||
# 第 8 轮是总控收口，不再执行假网关返回的派活。
|
||
assert len(gw.seat_tasks) == 7
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_recommend_uses_configured_closeout_round():
|
||
"""动态配置第 3 轮收口时，只执行前两轮席位派活。"""
|
||
gw = _RecGateway([[('seat-a', f'task-{i}')] for i in range(4)])
|
||
|
||
result = await _run(gw, gather_first=False, max_cycles=3)
|
||
|
||
assert result.status == "done"
|
||
assert len(gw.leader_requests) == 3
|
||
assert [finalize for _, finalize in gw.leader_requests] == [False, False, True]
|
||
assert len(gw.seat_tasks) == 2
|