devplacepy/tests/unit/services/jobs/deepsearch/orchestrate.py

# retoor <retoor@molodetz.nl>

import json

from tests.conftest import run_async

from devplacepy.services.deepsearch.store import Chunk
from devplacepy.services.jobs.deepsearch import orchestrate as orchestrate_module
from devplacepy.services.jobs.deepsearch.crawl import CrawledPage
from devplacepy.services.jobs.deepsearch.orchestrate import (
    Orchestration,
    _cosine,
    _mmr_select,
    orchestrate,
    source_diversity,
)


def _page(url, text="content " * 30):
    return CrawledPage(url=url, title="T", text=text, source="httpx", status=200)


def test_source_diversity_empty_is_zero():
    assert source_diversity([]) == 0.0


def test_source_diversity_increases_with_distinct_domains():
    one = source_diversity([_page("https://a.example/1"), _page("https://a.example/2")])
    many = source_diversity([_page("https://a.example/1"), _page("https://b.example/2")])
    assert 0.0 < one <= 1.0
    assert many > one


def test_source_diversity_capped_at_one():
    pages = [_page(f"https://d{i}.example") for i in range(8)]
    assert source_diversity(pages) <= 1.0


def test_orchestrate_with_no_pages_is_heuristic():
    result = run_async(orchestrate("q", [], "k", lambda frame: None))
    assert isinstance(result, Orchestration)
    assert result.source_diversity == 0.0
    assert result.synthesis == "heuristic"
    assert not result.findings


def test_orchestrate_grounded_run_emits_agent_frames(monkeypatch):
    frames = []

    report_payload = "## Answer\nA grounded answer [1]."
    findings_payload = json.dumps(
        {
            "findings": [
                {"title": "Finding", "detail": "Detail", "confidence": 0.7, "citations": [1]}
            ]
        }
    )
    link_payload = json.dumps({"confidence": 0.8})
    replies = iter([report_payload, findings_payload, link_payload])

    async def fake_request_completion(messages, api_key, **kwargs):
        return ({"choices": [{"message": {"content": next(replies)}}]}, {}, 5)

    monkeypatch.setattr(orchestrate_module, "request_completion", fake_request_completion)

    pages = [_page("https://a.example"), _page("https://b.example")]
    result = run_async(orchestrate("question", pages, "k", frames.append))

    assert result.summary == "## Answer\nA grounded answer [1]."
    assert result.findings and result.findings[0]["citations"] == [1]
    assert result.synthesis == "agents"
    assert 0.0 < result.confidence <= 1.0
    assert result.source_diversity > 0.0
    assert result.score > 0
    agents = {f["agent"] for f in frames if f.get("type") == "agent"}
    assert agents == {"summarizer", "extractor", "linker"}


def test_orchestrate_falls_back_to_heuristic_on_failure(monkeypatch):
    async def boom(messages, api_key, **kwargs):
        raise RuntimeError("upstream down")

    monkeypatch.setattr(orchestrate_module, "request_completion", boom)
    frames = []
    pages = [_page("https://a.example"), _page("https://b.example")]
    result = run_async(orchestrate("question", pages, "k", frames.append))
    assert result.findings
    assert result.synthesis == "heuristic"
    assert any(f.get("status") == "failed" for f in frames if f.get("type") == "agent")
    assert result.source_diversity > 0.0


def test_orchestrate_survives_linker_failure(monkeypatch):
    replies = iter(
        [
            "A full report [1].",
            json.dumps(
                {
                    "findings": [
                        {"title": "F", "detail": "D", "confidence": 0.5, "citations": [1]}
                    ]
                }
            ),
        ]
    )

    async def flaky(messages, api_key, **kwargs):
        try:
            content = next(replies)
        except StopIteration:
            raise RuntimeError("upstream down")
        return ({"choices": [{"message": {"content": content}}]}, {}, 5)

    monkeypatch.setattr(orchestrate_module, "request_completion", flaky)
    pages = [_page("https://a.example"), _page("https://b.example")]
    result = run_async(orchestrate("question", pages, "k", lambda frame: None))
    assert result.synthesis == "agents"
    assert result.summary == "A full report [1]."
    assert result.findings
    assert result.confidence >= 0.35


def test_numbered_source_digest_keeps_every_source_number_under_cap():
    from devplacepy.services.jobs.deepsearch.orchestrate import _numbered_source_digest
    import re

    pages = [_page(f"https://s{i}.example", text="word " * 500) for i in range(1, 13)]
    digest = _numbered_source_digest(pages, per_source=600, cap=4000)
    present = {int(n) for n in re.findall(r"\[(\d+)\]", digest)}
    assert present == set(range(1, 13))


def test_linker_receives_full_source_list(monkeypatch):
    seen = {}
    replies = iter(
        [
            "Report body [1].",
            json.dumps(
                {"findings": [{"title": "F", "detail": "D", "confidence": 1.0, "citations": [2]}]}
            ),
            json.dumps({"confidence": 0.9}),
        ]
    )

    async def capture(messages, api_key, **kwargs):
        content = messages[-1]["content"]
        if content.startswith("FINDINGS:") and "SOURCES:" in content:
            seen["linker"] = content
        return ({"choices": [{"message": {"content": next(replies)}}]}, {}, 5)

    monkeypatch.setattr(orchestrate_module, "request_completion", capture)
    pages = [_page(f"https://s{i}.example") for i in range(1, 13)]
    result = run_async(orchestrate("q", pages, "k", lambda f: None))
    assert result.confidence >= 0.35
    for n in range(1, 13):
        assert f"[{n}]" in seen["linker"]


def test_cosine_identical_vectors_is_one():
    assert round(_cosine([1.0, 0.0], [1.0, 0.0]), 6) == 1.0


def test_cosine_orthogonal_vectors_is_zero():
    assert _cosine([1.0, 0.0], [0.0, 1.0]) == 0.0


def test_cosine_mismatched_or_empty_is_zero():
    assert _cosine([], [1.0]) == 0.0
    assert _cosine([1.0], [1.0, 0.0]) == 0.0


def test_mmr_select_prefers_diverse_over_redundant():
    query_vector = [1.0, 0.0, 0.0]
    most_relevant = Chunk(uid="a", text="a", url="https://a", title="A", embedding=[0.9, 0.436, 0.0])
    near_duplicate = Chunk(uid="b", text="b", url="https://a2", title="B", embedding=[0.85, 0.527, 0.0])
    diverse = Chunk(uid="c", text="c", url="https://c", title="C", embedding=[0.85, 0.0, 0.527])
    selected = _mmr_select([most_relevant, near_duplicate, diverse], query_vector, limit=2)
    assert {chunk.uid for chunk in selected} == {"a", "c"}


def test_mmr_select_falls_back_when_embeddings_missing():
    chunks = [Chunk(uid="a", text="a", url="https://a", title="A")]
    assert _mmr_select(chunks, [1.0, 0.0], limit=1) == chunks


def test_mmr_select_empty_input():
    assert _mmr_select([], [1.0, 0.0], limit=3) == []


def test_orchestrate_grounded_run_includes_follow_up_questions(monkeypatch):
    replies = iter(
        [
            "## Answer\nA grounded answer [1].",
            json.dumps(
                {
                    "findings": [
                        {"title": "Finding", "detail": "Detail", "confidence": 0.7, "citations": [1]}
                    ]
                }
            ),
            json.dumps({"confidence": 0.8}),
            json.dumps({"questions": ["What about X?", "How does Y compare?"]}),
        ]
    )

    async def fake_request_completion(messages, api_key, **kwargs):
        return ({"choices": [{"message": {"content": next(replies)}}]}, {}, 5)

    monkeypatch.setattr(orchestrate_module, "request_completion", fake_request_completion)
    pages = [_page("https://a.example"), _page("https://b.example")]
    result = run_async(orchestrate("question", pages, "k", lambda frame: None))
    assert result.follow_up_questions == ["What about X?", "How does Y compare?"]


def test_orchestrate_heuristic_path_has_no_follow_up_questions():
    result = run_async(orchestrate("q", [], "k", lambda frame: None))
    assert result.follow_up_questions == []


def test_parse_json_tolerates_fences_and_trailing_garbage():
    from devplacepy.services.jobs.deepsearch.orchestrate import _parse_json

    assert _parse_json('```json\n{"gaps": ["g"]}\n```') == {"gaps": ["g"]}
    assert _parse_json('prose {"confidence": 0.7} more prose') == {"confidence": 0.7}
    assert _parse_json('{"gaps": ["a"]} {"broken": ') == {"gaps": ["a"]}
    assert _parse_json("no json here") == {}