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

# retoor <retoor@molodetz.nl>

import json
import tempfile
from pathlib import Path

from tests.conftest import run_async

from devplacepy.services.jobs.deepsearch import crawl as crawl_module
from devplacepy.services.jobs.deepsearch import worker as worker_module
from devplacepy.services.jobs.deepsearch.crawl import CrawledPage, CrawlOutcome
from devplacepy.services.deepsearch.embeddings import local_embed
from devplacepy.services.deepsearch.store import VectorStore


def _patch_pipeline(monkeypatch, pages):
    async def fake_plan(query, api_key, emit=lambda frame: None, model=None):
        return [query, f"{query} overview"]

    async def fake_search(queries, emit=lambda frame: None):
        return [{"url": page.url, "title": page.title, "description": ""} for page in pages]

    async def fake_crawl(
        candidates, max_pages, emit, is_cached, should_stop, query="", depth=1, seen_hashes=None
    ):
        outcome = CrawlOutcome(seen_hashes=seen_hashes if seen_hashes is not None else set())
        for page in pages[:max_pages]:
            emit({"type": "page_loaded", "url": page.url, "done": 1, "total": len(pages)})
            outcome.pages.append(page)
        return outcome

    async def fake_plan_followups(query, covered_titles, api_key, emit=lambda frame: None, model=None):
        return []

    def fake_embed(texts, api_key):
        return local_embed(texts)

    async def fake_embed_async(texts, api_key, **kwargs):
        return local_embed(texts)

    async def fake_orchestrate(question, crawled, api_key, emit, store=None, queries=None, model=None):
        from devplacepy.services.jobs.deepsearch.orchestrate import Orchestration

        return Orchestration(
            summary="A summary.",
            findings=[{"title": "F", "detail": "D", "confidence": 0.5, "citations": [1]}],
            confidence=0.6,
            source_diversity=0.5,
            score=42,
        )

    monkeypatch.setattr(worker_module, "plan_queries", fake_plan)
    monkeypatch.setattr(worker_module, "search_queries", fake_search)
    monkeypatch.setattr(worker_module, "crawl", fake_crawl)
    monkeypatch.setattr(worker_module, "plan_followup_queries", fake_plan_followups)
    monkeypatch.setattr(worker_module, "embed_texts", fake_embed_async)
    monkeypatch.setattr(worker_module, "orchestrate", fake_orchestrate)


def test_worker_run_produces_report(monkeypatch):
    pages = [
        CrawledPage(
            url="https://example.com/a",
            title="Page A",
            text="The transistor was invented at Bell Labs. " * 20,
            source="httpx",
            status=200,
        ),
        CrawledPage(
            url="https://other.example/b",
            title="Page B",
            text="Semiconductors are made from silicon. " * 20,
            source="playwright",
            status=200,
        ),
    ]
    _patch_pipeline(monkeypatch, pages)
    with tempfile.TemporaryDirectory() as tmp:
        output_dir = Path(tmp)
        payload = {
            "query": "history of the transistor",
            "max_pages": 5,
            "depth": 2,
            "api_key": "k",
            "collection": "ds_worker_test_one",
            "cached_hashes": [],
        }
        report = run_async(worker_module._run(payload, output_dir))
        assert report["query"] == "history of the transistor"
        assert report["page_count"] == 2
        assert report["chunk_count"] > 0
        assert report["summary"] == "A summary."
        assert report["findings"]
        assert (output_dir / "report.json").is_file()
        cache = json.loads((output_dir / "url_cache.json").read_text())
        assert len(cache) == 2
        VectorStore("ds_worker_test_one").drop()


def test_index_chunks_emits_batch_progress(monkeypatch):
    async def fake_embed_async(texts, api_key, **kwargs):
        return local_embed(texts)

    monkeypatch.setattr(worker_module, "embed_texts", fake_embed_async)
    monkeypatch.setattr(worker_module, "EMBED_BATCH", 1)

    frames = []
    pages = [
        CrawledPage(
            url="https://example.com/a",
            title="Page A",
            text="The transistor was invented at Bell Labs. " * 30,
            source="httpx",
            status=200,
        )
    ]
    store = VectorStore("ds_index_chunks_test")
    try:
        chunk_count, backend = run_async(
            worker_module._index_chunks(store, pages, "k", frames.append)
        )
        assert chunk_count > 0
        assert backend in ("gateway", "local")
        batch_frames = [f for f in frames if f.get("type") == "embed_batch"]
        assert batch_frames
        assert all(f["total"] == chunk_count for f in batch_frames)
        assert max(f["done"] for f in batch_frames) == chunk_count
        done_frames = [f for f in frames if f.get("type") == "embed_done"]
        assert done_frames and done_frames[-1]["chunk_count"] == chunk_count
    finally:
        store.drop()


def test_index_chunks_empty_pages_emits_done(monkeypatch):
    frames = []
    store = VectorStore("ds_index_chunks_empty_test")
    try:
        chunk_count, backend = run_async(
            worker_module._index_chunks(store, [], "k", frames.append)
        )
        assert chunk_count == 0
        assert backend == "empty"
        assert any(
            f.get("type") == "embed_done" and f.get("chunk_count") == 0 for f in frames
        )
    finally:
        store.drop()


def test_worker_run_performs_refinement_round_when_budget_remains(monkeypatch):
    initial_page = CrawledPage(
        url="https://example.com/a",
        title="Page A",
        text="The transistor was invented at Bell Labs. " * 20,
        source="httpx",
        status=200,
    )
    refined_page = CrawledPage(
        url="https://other.example/b",
        title="Page B",
        text="Semiconductors are made from silicon. " * 20,
        source="httpx",
        status=200,
    )
    _patch_pipeline(monkeypatch, [initial_page])

    followup_calls = []

    async def fake_plan_followups_once(query, covered_titles, api_key, emit=lambda frame: None, model=None):
        if followup_calls:
            return []
        followup_calls.append(covered_titles)
        return ["a more specific angle"]

    async def fake_search_followup(queries, emit=lambda frame: None):
        return [{"url": refined_page.url, "title": refined_page.title, "description": ""}]

    async def fake_crawl_refinement(
        candidates, max_pages, emit, is_cached, should_stop, query="", depth=1, seen_hashes=None
    ):
        outcome = CrawlOutcome(seen_hashes=seen_hashes if seen_hashes is not None else set())
        outcome.pages.append(refined_page)
        return outcome

    monkeypatch.setattr(worker_module, "plan_followup_queries", fake_plan_followups_once)

    real_search_queries = worker_module.search_queries
    real_crawl = worker_module.crawl
    call_count = {"search": 0, "crawl": 0}

    async def routed_search(queries, emit=lambda frame: None):
        call_count["search"] += 1
        if call_count["search"] == 1:
            return await real_search_queries(queries, emit)
        return await fake_search_followup(queries, emit)

    async def routed_crawl(*args, **kwargs):
        call_count["crawl"] += 1
        if call_count["crawl"] == 1:
            return await real_crawl(*args, **kwargs)
        return await fake_crawl_refinement(*args, **kwargs)

    monkeypatch.setattr(worker_module, "search_queries", routed_search)
    monkeypatch.setattr(worker_module, "crawl", routed_crawl)

    with tempfile.TemporaryDirectory() as tmp:
        output_dir = Path(tmp)
        payload = {
            "query": "history of the transistor",
            "max_pages": 5,
            "depth": 2,
            "api_key": "k",
            "collection": "ds_worker_refine_test",
            "cached_hashes": [],
        }
        report = run_async(worker_module._run(payload, output_dir))
        assert report["page_count"] == 2
        assert followup_calls and followup_calls[0] == ["Page A"]
        VectorStore("ds_worker_refine_test").drop()


def test_worker_control_cancel_stops(monkeypatch):
    pages = [
        CrawledPage(url="https://x.example", title="X", text="content " * 40, source="httpx", status=200)
    ]
    _patch_pipeline(monkeypatch, pages)
    with tempfile.TemporaryDirectory() as tmp:
        output_dir = Path(tmp)
        (output_dir / "control.json").write_text(json.dumps({"state": "cancelled"}))
        should_stop = worker_module._make_stop(output_dir)
        assert run_async(should_stop()) is True