236 lines
10 KiB
Python
236 lines
10 KiB
Python
"""RC-BE-009: TaskLedger waves, controlled member runner, repair attempts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from app.report_collaboration.execution.ledger import TaskLedger, TaskResult
|
|
from app.report_collaboration.execution.member_runner import ScriptedMemberKernel, WaveExecutor
|
|
from app.report_collaboration.execution.schema_gate import validate_artifact_draft
|
|
from deerflow.persistence.report_collaboration import MemoryReportCollaborationStore
|
|
|
|
|
|
def _evidence(text: str = "2024 年新能源乘用车零售份额上升") -> dict:
|
|
return {
|
|
"schema_version": 1,
|
|
"claims": [
|
|
{
|
|
"claim_id": "c1",
|
|
"text": text,
|
|
"kind": "fact",
|
|
"source_ids": ["src_1"],
|
|
"excerpt": "份额上升",
|
|
}
|
|
],
|
|
"sources": [{"source_id": "src_1", "title": "乘联会月报", "excerpt": "份额上升"}],
|
|
}
|
|
|
|
|
|
def _draft() -> dict:
|
|
return {
|
|
"schema_version": 1,
|
|
"section_id": "overview",
|
|
"title": "市场趋势",
|
|
"markdown": "份额上升,来源 c1。",
|
|
"evidence_ids": ["c1"],
|
|
"claim_citation_map": {"份额上升": ["c1"]},
|
|
}
|
|
|
|
|
|
def _plan(*, max_attempts: int = 3, extra_research: bool = False) -> dict:
|
|
research = {
|
|
"id": "research",
|
|
"label": "检索",
|
|
"role_key": "researcher",
|
|
"role_display_name": "研究员",
|
|
"mission": "收集证据",
|
|
"output_artifact_type": "EvidenceBundle",
|
|
"allowed_tools": ["web_search", "web_fetch", "knowledge_search"],
|
|
"acceptance_criteria": ["每条主张有来源"],
|
|
"depends_on": [],
|
|
"max_attempts": max_attempts,
|
|
}
|
|
nodes = [research]
|
|
edges = []
|
|
if extra_research:
|
|
nodes.append({**research, "id": "research-b", "label": "检索B", "angle": "价格"})
|
|
edges.append({"source": "research", "target": "write"})
|
|
edges.append({"source": "research-b", "target": "write"})
|
|
write_deps = ["research", "research-b"]
|
|
else:
|
|
edges.append({"source": "research", "target": "write"})
|
|
write_deps = ["research"]
|
|
nodes.append(
|
|
{
|
|
"id": "write",
|
|
"label": "写作",
|
|
"role_key": "writer",
|
|
"role_display_name": "撰写员",
|
|
"mission": "成稿",
|
|
"output_artifact_type": "ReportSectionDraft",
|
|
"allowed_tools": ["submit_artifact", "knowledge_search"],
|
|
"depends_on": write_deps,
|
|
"max_attempts": max_attempts,
|
|
}
|
|
)
|
|
return {
|
|
"id": "plan_poc3",
|
|
"proposal_group_id": "grp_poc",
|
|
"title": "三角色 POC",
|
|
"strategy": "focused_fast",
|
|
"summary": "研究后写作",
|
|
"rationale": "POC-2",
|
|
"recommended": True,
|
|
"estimated_duration_seconds": 600,
|
|
"estimated_cost_level": "low",
|
|
"requirement_revision": 0,
|
|
"roles": [
|
|
{"key": "coordinator", "display_name": "协调者", "responsibility": "建议"},
|
|
{"key": "researcher", "display_name": "研究员", "responsibility": "检索"},
|
|
{"key": "writer", "display_name": "撰写员", "responsibility": "成稿"},
|
|
],
|
|
"nodes": nodes,
|
|
"edges": edges,
|
|
"quality_gates": [{"id": "q1", "label": "引用覆盖"}],
|
|
"validation": {"ok": True, "errors": []},
|
|
"revision": 1,
|
|
"status": "proposed",
|
|
}
|
|
|
|
|
|
async def _boot(store: MemoryReportCollaborationStore, plan: dict | None = None) -> str:
|
|
session = await store.create_session(owner_id="u1", title="POC", idempotency_key="ik-s")
|
|
inserted = await store.insert_plan(session["id"], plan or _plan())
|
|
await store.select_plan(session["id"], inserted["id"], idempotency_key="ik-sel", expected_revision=None)
|
|
run = await store.create_run(session["id"], plan_id=inserted["id"], idempotency_key="ik-run", expected_revision=None)
|
|
return run["id"]
|
|
|
|
|
|
def test_schema_gate_rejects_empty_and_illegal_json() -> None:
|
|
assert validate_artifact_draft("EvidenceBundle", None).reason == "empty_output"
|
|
assert validate_artifact_draft("EvidenceBundle", {}).reason == "empty_output"
|
|
assert validate_artifact_draft("EvidenceBundle", "{not json").reason == "illegal_json"
|
|
assert validate_artifact_draft("EvidenceBundle", _evidence()).ok
|
|
assert not validate_artifact_draft("ReportSectionDraft", {"schema_version": 1, "section_id": "x", "title": "t", "markdown": "hi"}).ok
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_waves_researcher_before_writer_and_text_does_not_complete() -> None:
|
|
store = MemoryReportCollaborationStore()
|
|
run_id = await _boot(store)
|
|
kernel = ScriptedMemberKernel()
|
|
kernel.enqueue(
|
|
"research",
|
|
TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence(), reply_text="检索完成,证据如下。"),
|
|
)
|
|
kernel.enqueue(
|
|
"write",
|
|
TaskResult(kind="artifact_draft", artifact_type="ReportSectionDraft", artifact_draft=_draft(), reply_text="章节草稿。"),
|
|
)
|
|
executor = WaveExecutor(store, kernel)
|
|
status = await executor.run_until_idle(
|
|
run_id,
|
|
coordinator=lambda: TaskResult(kind="advice", suggested_next_node_ids=["write"], reply_text="请先写报告"),
|
|
)
|
|
assert executor.dispatched == ["research", "write"]
|
|
assert status == "completed"
|
|
nodes = await store.list_node_runs(run_id)
|
|
latest = {item["node_id"]: item for item in nodes if item["attempt"] == 1}
|
|
assert latest["research"]["status"] == "completed"
|
|
assert latest["write"]["status"] == "completed"
|
|
artifacts = await store.list_artifacts_full(run_id)
|
|
assert {item["artifact_type"] for item in artifacts} == {"EvidenceBundle", "ReportSectionDraft"}
|
|
assert all(item["validation_status"] == "validated" for item in artifacts)
|
|
events = await store.list_events(run_id)
|
|
types = [item["type"] for item in events]
|
|
assert "session.snapshot" in types
|
|
assert "message.created" in types
|
|
assert "message.delta" in types
|
|
advice = [item for item in events if item["type"] == "heartbeat" and (item.get("data") or {}).get("kind") == "coordinator_advice"]
|
|
assert advice
|
|
assert advice[0]["data"]["ignored"] is True
|
|
assert advice[0]["data"]["suggested_next_node_ids"] == ["write"]
|
|
agents = await store.list_agent_runs(run_id)
|
|
phases = {item["node_run_id"]: item["phase_id"] for item in agents}
|
|
assert phases[f"{run_id}-research-attempt1"] != phases[f"{run_id}-write-attempt1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_researchers_share_phase() -> None:
|
|
store = MemoryReportCollaborationStore()
|
|
run_id = await _boot(store, _plan(extra_research=True))
|
|
kernel = ScriptedMemberKernel()
|
|
kernel.enqueue("research", TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence("品牌")))
|
|
kernel.enqueue("research-b", TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence("价格")))
|
|
kernel.enqueue("write", TaskResult(kind="artifact_draft", artifact_type="ReportSectionDraft", artifact_draft=_draft()))
|
|
executor = WaveExecutor(store, kernel)
|
|
status = await executor.run_until_idle(run_id)
|
|
assert status == "completed"
|
|
assert set(executor.dispatched[:2]) == {"research", "research-b"}
|
|
assert executor.dispatched[2] == "write"
|
|
agents = await store.list_agent_runs(run_id)
|
|
research_phases = {item["phase_id"] for item in agents if item["role"] == "researcher"}
|
|
assert len(research_phases) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_and_illegal_json_repair_then_succeed() -> None:
|
|
store = MemoryReportCollaborationStore()
|
|
run_id = await _boot(store)
|
|
kernel = ScriptedMemberKernel()
|
|
kernel.enqueue(
|
|
"research",
|
|
TaskResult(kind="artifact_draft", artifact_draft=None, reply_text="先聊两句"),
|
|
TaskResult(kind="artifact_draft", artifact_draft="{broken"),
|
|
TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence()),
|
|
)
|
|
kernel.enqueue("write", TaskResult(kind="artifact_draft", artifact_type="ReportSectionDraft", artifact_draft=_draft()))
|
|
executor = WaveExecutor(store, kernel)
|
|
status = await executor.run_until_idle(run_id)
|
|
assert status == "completed"
|
|
nodes = await store.list_node_runs(run_id)
|
|
research = [item for item in nodes if item["node_id"] == "research"]
|
|
assert [item["attempt"] for item in research] == [1, 2, 3]
|
|
assert research[0]["status"] == "superseded"
|
|
assert research[1]["status"] == "superseded"
|
|
assert research[2]["status"] == "completed"
|
|
assert kernel.calls.count("research") == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_member_failure_does_not_complete_run_until_cap() -> None:
|
|
store = MemoryReportCollaborationStore()
|
|
run_id = await _boot(store, _plan(max_attempts=2))
|
|
kernel = ScriptedMemberKernel()
|
|
kernel.enqueue(
|
|
"research",
|
|
TaskResult(kind="failure", error="timeout"),
|
|
TaskResult(kind="failure", error="timeout"),
|
|
)
|
|
executor = WaveExecutor(store, kernel)
|
|
status = await executor.run_until_idle(run_id)
|
|
assert status == "failed"
|
|
assert "write" not in executor.dispatched
|
|
nodes = await store.list_node_runs(run_id)
|
|
research = [item for item in nodes if item["node_id"] == "research"]
|
|
assert research[-1]["status"] == "failed"
|
|
assert research[-1]["attempt"] == 2
|
|
write = [item for item in nodes if item["node_id"] == "write"]
|
|
assert write[0]["status"] == "planned"
|
|
run = await store.get_run(run_id)
|
|
assert run["status"] == "failed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_coordinator_cannot_dispatch_writer() -> None:
|
|
store = MemoryReportCollaborationStore()
|
|
run_id = await _boot(store)
|
|
ledger = TaskLedger(store)
|
|
advice = TaskResult(kind="advice", suggested_next_node_ids=["write"])
|
|
outcome = await ledger.record_coordinator_advice(run_id, advice)
|
|
assert outcome.ignored_coordinator_suggestions == ["write"]
|
|
wave = await ledger.ready_assignments(run_id)
|
|
assert [item.node_id for item in wave] == ["research"]
|
|
second = await ledger.ready_assignments(run_id)
|
|
assert second == []
|