116 lines
4.2 KiB
Python
116 lines
4.2 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.roles import resolve_runtime_profile
|
|
from scripts.runtime.sources import append_packet_sources, rebuild_sources_from_packets
|
|
from scripts.runtime.tasks import TaskCard
|
|
from scripts.runtime.workers import PacketWorker, build_search_context
|
|
|
|
|
|
class FakeSearchProvider:
|
|
def search(self, *, query: str, route: str, num_results: int):
|
|
return [
|
|
{
|
|
"title": f"{route} result for {query}",
|
|
"url": f"https://example.com/{route}",
|
|
"snippet": "候选证据摘要",
|
|
"route": route,
|
|
}
|
|
][:num_results]
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self, response: dict) -> None:
|
|
self.response = response
|
|
self.calls: list[dict[str, object]] = []
|
|
|
|
def chat_complete(self, **kwargs) -> str:
|
|
self.calls.append(kwargs)
|
|
return json.dumps(self.response, ensure_ascii=False)
|
|
|
|
|
|
def sample_card() -> TaskCard:
|
|
return TaskCard(
|
|
task_id="ch01-literature",
|
|
chapter_ids=["ch01"],
|
|
topic_axis="literature",
|
|
questions=["围绕临床证据提炼结论。"],
|
|
search_routes=["scholar", "general"],
|
|
output_packet="phase2/packets/ch01-literature.json",
|
|
)
|
|
|
|
|
|
def test_build_search_context_assigns_stable_source_ids() -> None:
|
|
context = build_search_context(sample_card(), FakeSearchProvider(), num_results_per_route=1)
|
|
|
|
assert [source["id"] for source in context["candidate_sources"]] == [
|
|
"src_ch01_literature_001",
|
|
"src_ch01_literature_002",
|
|
]
|
|
assert context["routes_used"] == ["scholar", "general"]
|
|
|
|
|
|
def test_packet_worker_includes_search_context_in_prompt() -> None:
|
|
context = build_search_context(sample_card(), FakeSearchProvider(), num_results_per_route=1)
|
|
response = {
|
|
"task_id": "ch01-literature",
|
|
"claims": [{"claim": "候选证据支持判断", "source_ids": ["src_ch01_literature_001"]}],
|
|
"evidence_items": [{"source_id": "src_ch01_literature_001", "summary": "摘要"}],
|
|
"counter_evidence": [{"claim": "仍需更多数据", "source_ids": ["src_ch01_literature_002"]}],
|
|
"source_ids": ["src_ch01_literature_001", "src_ch01_literature_002"],
|
|
"sources": context["candidate_sources"],
|
|
"source_quality_notes": ["候选来源需要后续评级"],
|
|
"open_questions": [],
|
|
"raw_quotes_or_notes": [],
|
|
}
|
|
role = resolve_runtime_profile(profile="medium").role_for_task("evidence_packet")
|
|
fake = FakeClient(response)
|
|
|
|
packet = PacketWorker(role=role, client=fake, search_provider=FakeSearchProvider()).run(sample_card())
|
|
|
|
assert packet["sources"][0]["url"] == "https://example.com/scholar"
|
|
assert "candidate_sources" in fake.calls[0]["user"]
|
|
|
|
|
|
def test_append_packet_sources_dedupes_by_url(tmp_path: Path) -> None:
|
|
packet = {
|
|
"sources": [
|
|
{"id": "src_a", "title": "A", "url": "https://example.com/a", "tier": "Tier 2", "score": 7},
|
|
{"id": "src_b", "title": "B", "url": "https://example.com/a", "tier": "Tier 2", "score": 7},
|
|
]
|
|
}
|
|
|
|
written = append_packet_sources(tmp_path / "sources.jsonl", packet)
|
|
|
|
assert written == 1
|
|
assert len((tmp_path / "sources.jsonl").read_text(encoding="utf-8").splitlines()) == 1
|
|
|
|
|
|
def test_rebuild_sources_from_packets_dedupes_manual_packets(tmp_path: Path) -> None:
|
|
project = tmp_path / "project"
|
|
packets = project / "phase2" / "packets"
|
|
packets.mkdir(parents=True)
|
|
packet = {
|
|
"sources": [
|
|
{"id": "src_001", "title": "A", "url": "https://example.com/a"},
|
|
{"id": "src_002", "title": "A duplicate", "url": "https://example.com/a"},
|
|
{"id": "src_003", "title": "Local", "url": "phase0/extracted/local.md"},
|
|
]
|
|
}
|
|
(packets / "ch01-a.json").write_text(json.dumps(packet, ensure_ascii=False), encoding="utf-8")
|
|
|
|
count = rebuild_sources_from_packets(project)
|
|
|
|
lines = (project / "phase2" / "sources.jsonl").read_text(encoding="utf-8").splitlines()
|
|
assert count == 2
|
|
assert len(lines) == 2
|
|
assert "src_001" in lines[0]
|
|
assert "src_003" in lines[1]
|