Files
deep_research/tests/test_search_grounded_packets.py
T

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]