240 lines
9.7 KiB
Python
240 lines
9.7 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from scripts.runtime.assembly import (
|
|
ChapterAssemblyWorker,
|
|
build_chapter_briefs,
|
|
build_compressed_findings,
|
|
build_chapter_user_prompt,
|
|
run_chapter_assembly_workers,
|
|
validate_chapter_markdown_citations,
|
|
validate_chapter_brief,
|
|
validate_compressed_finding,
|
|
)
|
|
from scripts.runtime.roles import resolve_runtime_profile
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self, response: str) -> None:
|
|
self.response = response
|
|
self.calls: list[dict[str, object]] = []
|
|
|
|
def chat_complete(self, **kwargs) -> str:
|
|
self.calls.append(kwargs)
|
|
return self.response
|
|
|
|
|
|
class TaggedClient:
|
|
def __init__(self, responses_by_tag: dict[str, str]) -> None:
|
|
self.responses_by_tag = responses_by_tag
|
|
self.calls: list[dict[str, object]] = []
|
|
|
|
def chat_complete(self, **kwargs) -> str:
|
|
self.calls.append(kwargs)
|
|
return self.responses_by_tag[str(kwargs["tag"])]
|
|
|
|
|
|
def chapter_brief(chapter_id: str, source_ids: list[str]) -> dict:
|
|
return {
|
|
"chapter_id": chapter_id,
|
|
"chapter_title": "临床证据正在重塑需求判断",
|
|
"packet_ids": [f"{chapter_id}-clinical"],
|
|
"core_claims": [{"claim": "临床证据支持核心判断", "source_ids": source_ids[:1]}],
|
|
"evidence_items": [{"source_id": source_ids[0], "summary": "III 期数据支持主要终点"}],
|
|
"counter_evidence": [{"claim": "长期安全性仍需随访", "source_ids": source_ids[-1:]}],
|
|
"source_ids": source_ids,
|
|
"open_questions": [],
|
|
"assembly_notes": ["按金字塔结构组织。"],
|
|
}
|
|
|
|
|
|
def write_packet(path: Path, task_id: str, claim: str, source_id: str) -> None:
|
|
packet = {
|
|
"task_id": task_id,
|
|
"claims": [{"claim": claim, "source_ids": [source_id]}],
|
|
"evidence_items": [{"source_id": source_id, "summary": f"{claim} 的证据"}],
|
|
"counter_evidence": [{"claim": "仍需关注样本量和外推限制", "source_ids": ["src_counter"]}],
|
|
"source_ids": [source_id, "src_counter"],
|
|
"source_quality_notes": [f"{source_id} Tier 1"],
|
|
"open_questions": ["还需要补充中国市场数据"],
|
|
"raw_quotes_or_notes": ["English note can remain as source material."],
|
|
}
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(packet, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
def test_build_chapter_briefs_aggregates_packets_by_chapter(tmp_path: Path) -> None:
|
|
project = tmp_path / "project"
|
|
cards = [
|
|
{
|
|
"task_id": "ch01-clinical",
|
|
"chapter_ids": ["ch01"],
|
|
"topic_axis": "clinical",
|
|
"questions": ["q"],
|
|
"search_routes": ["scholar"],
|
|
"output_packet": "phase2/packets/ch01-clinical.json",
|
|
},
|
|
{
|
|
"task_id": "ch01-market",
|
|
"chapter_ids": ["ch01"],
|
|
"topic_axis": "market",
|
|
"questions": ["q"],
|
|
"search_routes": ["news"],
|
|
"output_packet": "phase2/packets/ch01-market.json",
|
|
},
|
|
]
|
|
(project / "phase2").mkdir(parents=True)
|
|
(project / "phase2" / "task_cards.json").write_text(json.dumps(cards, ensure_ascii=False), encoding="utf-8")
|
|
write_packet(project / "phase2/packets/ch01-clinical.json", "ch01-clinical", "临床证据支持核心判断", "src_001")
|
|
write_packet(project / "phase2/packets/ch01-market.json", "ch01-market", "市场数据支持需求增长", "src_002")
|
|
|
|
briefs = build_chapter_briefs(project)
|
|
|
|
assert len(briefs) == 1
|
|
brief = briefs[0]
|
|
validate_chapter_brief(brief)
|
|
assert brief["chapter_id"] == "ch01"
|
|
assert brief["packet_ids"] == ["ch01-clinical", "ch01-market"]
|
|
assert "src_001" in brief["source_ids"]
|
|
assert "src_002" in brief["source_ids"]
|
|
assert (project / "phase2/chapter_briefs/ch01.json").exists()
|
|
|
|
compressed = build_compressed_findings(project)
|
|
|
|
assert len(compressed) == 1
|
|
finding = compressed[0]
|
|
validate_compressed_finding(finding)
|
|
assert finding["chapter_id"] == "ch01"
|
|
assert "chapter_thesis" in finding
|
|
assert "evidence_landings" in finding
|
|
assert "src_001" in finding["source_ids"]
|
|
assert (project / "phase2/compressed_findings/ch01.json").exists()
|
|
|
|
|
|
def test_chapter_prompt_contains_brief_and_fragmentation_guard() -> None:
|
|
brief = {
|
|
"chapter_id": "ch01",
|
|
"chapter_title": "临床证据正在重塑需求判断",
|
|
"packet_ids": ["ch01-clinical"],
|
|
"core_claims": [{"claim": "临床证据支持核心判断", "source_ids": ["src_001"]}],
|
|
"evidence_items": [{"source_id": "src_001", "summary": "III 期数据支持主要终点"}],
|
|
"counter_evidence": [{"claim": "长期安全性仍需随访", "source_ids": ["src_002"]}],
|
|
"source_ids": ["src_001", "src_002"],
|
|
"open_questions": [],
|
|
"assembly_notes": ["避免重复 packet 原文；按金字塔结构组织。"],
|
|
}
|
|
|
|
prompt = build_chapter_user_prompt(brief)
|
|
|
|
assert "临床证据正在重塑需求判断" in prompt
|
|
assert "避免碎片化" in prompt
|
|
assert "compressed finding" in prompt
|
|
assert "只输出 Markdown" in prompt
|
|
|
|
|
|
def test_chapter_assembly_worker_writes_markdown(tmp_path: Path) -> None:
|
|
runtime = resolve_runtime_profile(profile="medium")
|
|
role = runtime.role_for_task("chapter_assembly")
|
|
fake = FakeClient("# 第1章 临床证据正在重塑需求判断\n\n结论先行。[src_001]\n")
|
|
worker = ChapterAssemblyWorker(role=role, client=fake)
|
|
brief = {
|
|
"chapter_id": "ch01",
|
|
"chapter_title": "临床证据正在重塑需求判断",
|
|
"packet_ids": ["ch01-clinical"],
|
|
"core_claims": [{"claim": "临床证据支持核心判断", "source_ids": ["src_001"]}],
|
|
"evidence_items": [{"source_id": "src_001", "summary": "III 期数据支持主要终点"}],
|
|
"counter_evidence": [{"claim": "长期安全性仍需随访", "source_ids": ["src_002"]}],
|
|
"source_ids": ["src_001", "src_002"],
|
|
"open_questions": [],
|
|
"assembly_notes": ["按金字塔结构组织。"],
|
|
}
|
|
|
|
output = worker.write_chapter(project_root=tmp_path, brief=brief)
|
|
|
|
assert output == tmp_path / "phase2/drafts/ch01.md"
|
|
assert "结论先行" in output.read_text(encoding="utf-8")
|
|
assert fake.calls[0]["model"] == role.model
|
|
|
|
|
|
def test_validate_chapter_markdown_rejects_unknown_source_ids() -> None:
|
|
brief = {
|
|
"chapter_id": "ch01",
|
|
"chapter_title": "临床证据正在重塑需求判断",
|
|
"packet_ids": ["ch01-clinical"],
|
|
"core_claims": [{"claim": "临床证据支持核心判断", "source_ids": ["src_001"]}],
|
|
"evidence_items": [{"source_id": "src_001", "summary": "III 期数据支持主要终点"}],
|
|
"counter_evidence": [{"claim": "长期安全性仍需随访", "source_ids": ["src_002"]}],
|
|
"source_ids": ["src_001", "src_002"],
|
|
"open_questions": [],
|
|
"assembly_notes": ["按金字塔结构组织。"],
|
|
}
|
|
|
|
try:
|
|
validate_chapter_markdown_citations("结论引用了不存在的来源。[src_fake]", brief)
|
|
except ValueError as exc:
|
|
assert "unknown citation ids" in str(exc)
|
|
assert "src_fake" in str(exc)
|
|
else:
|
|
raise AssertionError("unknown source id should fail validation")
|
|
|
|
|
|
def test_chapter_assembly_worker_refuses_to_write_unknown_citations(tmp_path: Path) -> None:
|
|
runtime = resolve_runtime_profile(profile="medium")
|
|
role = runtime.role_for_task("chapter_assembly")
|
|
fake = FakeClient("# 第1章 临床证据正在重塑需求判断\n\n结论先行。[src_fake]\n")
|
|
worker = ChapterAssemblyWorker(role=role, client=fake)
|
|
brief = {
|
|
"chapter_id": "ch01",
|
|
"chapter_title": "临床证据正在重塑需求判断",
|
|
"packet_ids": ["ch01-clinical"],
|
|
"core_claims": [{"claim": "临床证据支持核心判断", "source_ids": ["src_001"]}],
|
|
"evidence_items": [{"source_id": "src_001", "summary": "III 期数据支持主要终点"}],
|
|
"counter_evidence": [{"claim": "长期安全性仍需随访", "source_ids": ["src_002"]}],
|
|
"source_ids": ["src_001", "src_002"],
|
|
"open_questions": [],
|
|
"assembly_notes": ["按金字塔结构组织。"],
|
|
}
|
|
|
|
try:
|
|
worker.write_chapter(project_root=tmp_path, brief=brief)
|
|
except ValueError as exc:
|
|
assert "unknown citation ids" in str(exc)
|
|
else:
|
|
raise AssertionError("chapter with unknown citation should not be written")
|
|
assert not (tmp_path / "phase2/drafts/ch01.md").exists()
|
|
|
|
|
|
def test_run_chapter_assembly_workers_records_errors_without_aborting_batch(tmp_path: Path) -> None:
|
|
runtime = resolve_runtime_profile(profile="medium")
|
|
fake = TaggedClient(
|
|
{
|
|
"chapter:ch01": "# 第1章 临床证据正在重塑需求判断\n\n结论先行。[src_001]\n",
|
|
"chapter:ch02": "# 第2章 临床证据存在不确定性\n\n错误引用。[src_fake]\n",
|
|
}
|
|
)
|
|
|
|
count = run_chapter_assembly_workers(
|
|
project_root=tmp_path,
|
|
briefs=[chapter_brief("ch01", ["src_001", "src_002"]), chapter_brief("ch02", ["src_003", "src_004"])],
|
|
runtime=runtime,
|
|
client_factory=lambda _role: fake,
|
|
workers=2,
|
|
)
|
|
|
|
assert count == 1
|
|
assert (tmp_path / "phase2/drafts/ch01.md").exists()
|
|
assert not (tmp_path / "phase2/drafts/ch02.md").exists()
|
|
error_path = tmp_path / "phase2/chapter_errors/ch02.json"
|
|
assert error_path.exists()
|
|
error = json.loads(error_path.read_text(encoding="utf-8"))
|
|
assert error["chapter_id"] == "ch02"
|
|
assert error["status"] == "failed"
|
|
assert "src_fake" in error["error"]
|