deep_research/tests/test_v020_workers.py

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.tasks import TaskCard, validate_packet
from scripts.runtime.workers import PacketWorker, build_packet_user_prompt, run_packet_workers


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)


class SequenceClient:
    def __init__(self, responses: list[str]) -> None:
        self.responses = responses
        self.calls: list[dict[str, object]] = []

    def chat_complete(self, **kwargs) -> str:
        self.calls.append(kwargs)
        if self.responses:
            return self.responses.pop(0)
        return "not json"


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.get(str(kwargs.get("tag")), "not json")


def sample_card() -> TaskCard:
    return TaskCard(
        task_id="ch01-clinical",
        chapter_ids=["ch01"],
        topic_axis="clinical",
        questions=["围绕临床证据提炼结论。"],
        search_routes=["scholar", "general"],
        output_packet="phase2/packets/ch01-clinical.json",
    )


def second_card() -> TaskCard:
    return TaskCard(
        task_id="ch02-market",
        chapter_ids=["ch02"],
        topic_axis="market",
        questions=["围绕市场证据提炼结论。"],
        search_routes=["news", "general"],
        output_packet="phase2/packets/ch02-market.json",
    )


def valid_response() -> dict:
    return {
        "task_id": "ch01-clinical",
        "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"],
        "sources": [
            {"id": "src_001", "title": "来源1", "url": "https://example.com/1"},
            {"id": "src_002", "title": "来源2", "url": "https://example.com/2"},
        ],
        "source_quality_notes": ["src_001 Tier 1; src_002 Tier 2"],
        "open_questions": [],
        "raw_quotes_or_notes": ["Original English evidence note is allowed."],
    }


def test_build_packet_prompt_contains_card_and_chinese_policy() -> None:
    prompt = build_packet_user_prompt(sample_card())

    assert "ch01-clinical" in prompt
    assert "中文" in prompt
    assert "scholar" in prompt


def test_packet_worker_generates_valid_packet_with_fake_client(tmp_path: Path) -> None:
    runtime = resolve_runtime_profile(profile="medium")
    role = runtime.role_for_task("evidence_packet")
    fake = FakeClient(valid_response())
    worker = PacketWorker(role=role, client=fake)

    packet = worker.run(sample_card())

    validate_packet(packet)
    assert fake.calls[0]["model"] == role.model
    assert "章节证据分析师" in fake.calls[0]["system"]
    assert "search-strategy" in fake.calls[0]["system"]


def test_run_packet_workers_writes_packet_files(tmp_path: Path) -> None:
    project = tmp_path / "project"
    fake = FakeClient(valid_response())
    runtime = resolve_runtime_profile(profile="medium")

    count = run_packet_workers(
        project_root=project,
        cards=[sample_card()],
        runtime=runtime,
        client_factory=lambda _role: fake,
        workers=2,
    )

    packet_path = project / "phase2/packets/ch01-clinical.json"
    assert count == 1
    assert packet_path.exists()
    validate_packet(json.loads(packet_path.read_text(encoding="utf-8")))


def test_packet_worker_repairs_malformed_json_once() -> None:
    runtime = resolve_runtime_profile(profile="medium")
    role = runtime.role_for_task("evidence_packet")
    fake = SequenceClient(["这里是说明，不是 JSON", json.dumps(valid_response(), ensure_ascii=False)])
    worker = PacketWorker(role=role, client=fake)

    packet = worker.run(sample_card())

    validate_packet(packet)
    assert len(fake.calls) == 2
    assert "修复" in str(fake.calls[1]["user"])
    assert fake.calls[1]["temperature"] == 0
    assert fake.calls[1]["tag"] == "packet-repair:ch01-clinical"


def test_run_packet_workers_writes_error_file_without_aborting_batch(tmp_path: Path) -> None:
    project = tmp_path / "project"
    runtime = resolve_runtime_profile(profile="medium")
    valid = json.dumps(valid_response(), ensure_ascii=False)
    fake = TaggedClient(
        {
            "packet:ch01-clinical": valid,
            "packet:ch02-market": "not json",
            "packet-repair:ch02-market": "still not json",
        }
    )

    count = run_packet_workers(
        project_root=project,
        cards=[sample_card(), second_card()],
        runtime=runtime,
        client_factory=lambda _role: fake,
        workers=2,
    )

    assert count == 1
    assert (project / "phase2/packets/ch01-clinical.json").exists()
    error_path = project / "phase2/packet_errors/ch02-market.json"
    assert error_path.exists()
    error = json.loads(error_path.read_text(encoding="utf-8"))
    assert error["task_id"] == "ch02-market"
    assert error["status"] == "failed"
    assert "worker response does not contain a JSON object" in error["error"]