typosaurus-sandbox/tests/test_research_pipeline.py

# retoor <retoor@molodetz.nl>

import unittest
from typing import Any, AsyncIterator
from unittest.mock import patch

from typosaurus_sandbox.research.client import RsearchError
from typosaurus_sandbox.research.config import ResearchConfig
from typosaurus_sandbox.research.envelopes import ChatResponse, DeepReport, DescribeResponse, SearchResponse, SearchResult
from typosaurus_sandbox.research.frontier import QueryFrontier
from typosaurus_sandbox.research.pipeline import (
    RETRY_MAX_ATTEMPTS,
    Extraction,
    ResearchPipeline,
    WorkItem,
    apply_extraction,
    extract_response,
)


async def _no_sleep(delay: float) -> None:
    return None


class StubResearchClient:
    def __init__(self) -> None:
        self.config = ResearchConfig(max_concurrency=4, default_count=5)
        self.cache_hits: dict[tuple[str, str], Any] = {}
        self.error_on: set[tuple[str, str]] = set()
        self.explode_on: set[tuple[str, str]] = set()
        self.failures_remaining: dict[tuple[str, str], int] = {}
        self.chat_response = ChatResponse(response="chat answer")
        self.describe_response = DescribeResponse(description="described page")
        self.calls: list[tuple[str, str]] = []

    def _maybe_fail(self, key: tuple[str, str], error: RsearchError) -> None:
        if key in self.explode_on:
            raise ValueError("unexpected boom")
        if key in self.error_on:
            raise error
        remaining = self.failures_remaining.get(key, 0)
        if remaining > 0:
            self.failures_remaining[key] = remaining - 1
            raise error

    def search_cached(
        self,
        query: str,
        *,
        source: str | None = None,
        count: int | None = None,
        content: bool = False,
        type: str | None = None,
        deep: bool = False,
        ai: bool = False,
        cache: bool = True,
    ) -> Any:
        if type == "images":
            return self.cache_hits.get(("images", query))
        return self.cache_hits.get(("web", query))

    def describe_cached(self, url: str) -> Any:
        return self.cache_hits.get(("describe", url))

    async def search(
        self,
        query: str,
        *,
        source: str | None = None,
        count: int | None = None,
        content: bool = False,
        type: str | None = None,
        deep: bool = False,
        ai: bool = False,
        cache: bool = True,
    ) -> SearchResponse:
        self.calls.append(("search", query))
        self._maybe_fail(("search", query), RsearchError("search failed", 503))
        return SearchResponse(
            query=query,
            source="duckduckgo",
            count=1,
            success=True,
            results=[SearchResult(title=query, url=f"https://example.com/{query}", description="details", content="body")],
        )

    async def chat(self, prompt: str, *, system: str | None = None, json_mode: bool = False, cache: bool = True) -> ChatResponse:
        self.calls.append(("chat", prompt))
        if ("chat", prompt) in self.error_on:
            raise RsearchError("chat failed", 400)
        return self.chat_response

    async def describe(self, url: str) -> DescribeResponse:
        self.calls.append(("describe", url))
        if ("describe", url) in self.error_on:
            raise RsearchError("describe failed", 500)
        return self.describe_response


class TestExtractResponse(unittest.TestCase):

    def test_web_results_extract_urls_seeds_and_content(self) -> None:
        item = WorkItem("web", "query")
        response = SearchResponse(
            query="query",
            source="duckduckgo",
            count=2,
            success=True,
            results=[
                SearchResult(title="First", url="https://a.example/1", description="First details", content="first body"),
                SearchResult(title="Second", url="https://b.example/2", description="Second details", content=None),
            ],
        )
        extraction = extract_response(item, response)
        self.assertEqual(extraction.urls, ("https://a.example/1", "https://b.example/2"))
        self.assertEqual(
            extraction.query_seeds,
            (
                ("First", "title"),
                ("First details", "description"),
                ("Second", "title"),
                ("Second details", "description"),
            ),
        )
        self.assertEqual(extraction.content_texts, ("first body",))

    def test_ai_response_adds_content_seed_and_urls(self) -> None:
        item = WorkItem("web", "query")
        response = SearchResponse(
            query="query",
            source="ai",
            count=0,
            success=True,
            results=[],
            ai_response="Overview at https://docs.example.org/x and https://blog.example.org/y.",
        )
        extraction = extract_response(item, response)
        self.assertEqual(extraction.urls, ("https://docs.example.org/x", "https://blog.example.org/y"))
        self.assertEqual(
            extraction.query_seeds,
            (("Overview at https://docs.example.org/x and https://blog.example.org/y.", "ai_response"),),
        )
        self.assertEqual(
            extraction.content_texts,
            ("Overview at https://docs.example.org/x and https://blog.example.org/y.",),
        )

    def test_deep_report_sources_and_markdown_extracted(self) -> None:
        item = WorkItem("web", "query")
        response = SearchResponse(
            query="query",
            source="google",
            count=1,
            success=True,
            results=[],
            deep=DeepReport(
                query="query",
                markdown="# Deep\n\nSee https://deep.example.org/report for details.",
                sources=[SearchResult(title="Deep source", url="https://deep.example.org/source", description="Deep details")],
            ),
        )
        extraction = extract_response(item, response)
        self.assertEqual(extraction.urls, ("https://deep.example.org/source", "https://deep.example.org/report"))
        self.assertEqual(
            extraction.query_seeds,
            (("Deep source", "title"), ("Deep details", "description")),
        )
        self.assertEqual(extraction.content_texts, ("# Deep\n\nSee https://deep.example.org/report for details.",))

    def test_chat_response_extracts_content_seed_and_urls(self) -> None:
        item = WorkItem("chat", "prompt")
        response = ChatResponse(response="Answer at https://chat.example.org/a.")
        extraction = extract_response(item, response)
        self.assertEqual(extraction.urls, ("https://chat.example.org/a",))
        self.assertEqual(extraction.query_seeds, (("Answer at https://chat.example.org/a.", "chat"),))
        self.assertEqual(extraction.content_texts, ("Answer at https://chat.example.org/a.",))

    def test_describe_response_extracts_description_seed_and_urls(self) -> None:
        item = WorkItem("describe", "https://page.example.org/x")
        response = DescribeResponse(description="Image shows a cat. More at https://gallery.example.org/cat.")
        extraction = extract_response(item, response)
        self.assertEqual(extraction.urls, ("https://gallery.example.org/cat",))
        self.assertEqual(
            extraction.query_seeds,
            (("Image shows a cat. More at https://gallery.example.org/cat.", "describe"),),
        )
        self.assertEqual(extraction.content_texts, ("Image shows a cat. More at https://gallery.example.org/cat.",))


class TestApplyExtraction(unittest.IsolatedAsyncioTestCase):

    async def test_registers_each_kind_and_returns_counts(self) -> None:
        frontier = QueryFrontier()
        extraction = Extraction(
            urls=("https://x.example/1", "https://y.example/2"),
            query_seeds=(("variant one", "title"), ("variant two", "description")),
            content_texts=("body one", "body two"),
        )
        self.assertEqual(apply_extraction(frontier, extraction), (2, 2, 2))
        self.assertEqual(apply_extraction(frontier, extraction), (0, 0, 0))
        stats = frontier.snapshot()
        self.assertEqual(stats.urls_seen, 2)
        self.assertEqual(stats.queries_enqueued, 2)
        self.assertEqual(stats.content_seen, 2)


class TestPipelineCacheProbe(unittest.IsolatedAsyncioTestCase):

    async def test_web_probe_reflects_search_cache(self) -> None:
        client = StubResearchClient()
        frontier = QueryFrontier()
        pipeline = ResearchPipeline(client, frontier)
        self.assertFalse(pipeline._probe_cache(WorkItem("web", "alpha")))
        client.cache_hits[("web", "alpha")] = SearchResponse(query="alpha", success=True)
        self.assertTrue(pipeline._probe_cache(WorkItem("web", "alpha")))

    async def test_images_probe_uses_images_search_cache(self) -> None:
        client = StubResearchClient()
        frontier = QueryFrontier()
        pipeline = ResearchPipeline(client, frontier)
        self.assertFalse(pipeline._probe_cache(WorkItem("images", "alpha")))
        client.cache_hits[("images", "alpha")] = SearchResponse(query="alpha", success=True)
        self.assertTrue(pipeline._probe_cache(WorkItem("images", "alpha")))

    async def test_describe_probe_reflects_describe_cache(self) -> None:
        client = StubResearchClient()
        frontier = QueryFrontier()
        pipeline = ResearchPipeline(client, frontier)
        self.assertFalse(pipeline._probe_cache(WorkItem("describe", "https://z.example/page")))
        client.cache_hits[("describe", "https://z.example/page")] = DescribeResponse(description="cached")
        self.assertTrue(pipeline._probe_cache(WorkItem("describe", "https://z.example/page")))

    async def test_chat_never_probes_cache(self) -> None:
        client = StubResearchClient()
        frontier = QueryFrontier()
        pipeline = ResearchPipeline(client, frontier)
        self.assertFalse(pipeline._probe_cache(WorkItem("chat", "question")))


class TestPipelineOutcomes(unittest.IsolatedAsyncioTestCase):

    async def test_chat_cached_response_marks_outcome_cache_hit(self) -> None:
        client = StubResearchClient()
        client.chat_response = ChatResponse(response="cached answer", cached=True)
        frontier = QueryFrontier()
        pipeline = ResearchPipeline(client, frontier)
        outcome = await pipeline.process(WorkItem("chat", "question"))
        self.assertTrue(outcome.success)
        self.assertTrue(outcome.cache_hit)
        self.assertEqual(outcome.endpoint, "/chat")

    async def test_search_error_produces_failure_outcome_and_pool_survives(self) -> None:
        client = StubResearchClient()
        client.error_on.add(("search", "bad"))
        frontier = QueryFrontier()

        async def items() -> AsyncIterator[WorkItem]:
            yield WorkItem("web", "bad")
            yield WorkItem("web", "good")

        pipeline = ResearchPipeline(client, frontier)
        with patch("typosaurus_sandbox.research.pipeline.asyncio.sleep", side_effect=_no_sleep):
            report = await pipeline.run(items())
        self.assertEqual(report.requests_succeeded, 1)
        self.assertEqual(report.requests_failed, 1)
        self.assertEqual(len(report.outcomes), 2)
        failed = next(outcome for outcome in report.outcomes if not outcome.success)
        self.assertEqual(failed.endpoint, "/search")
        self.assertEqual(failed.status_code, 503)
        self.assertEqual(failed.error, "search failed")

    async def test_run_drains_all_four_content_types(self) -> None:
        client = StubResearchClient()
        frontier = QueryFrontier()

        async def items() -> AsyncIterator[WorkItem]:
            yield WorkItem("web", "alpha")
            yield WorkItem("images", "alpha")
            yield WorkItem("describe", "https://z.example/page")
            yield WorkItem("chat", "question")

        pipeline = ResearchPipeline(client, frontier)
        report = await pipeline.run(items())
        self.assertEqual(report.requests_succeeded, 4)
        self.assertEqual(report.requests_failed, 0)
        self.assertEqual(len(report.outcomes), 4)
        self.assertTrue(all(outcome.success for outcome in report.outcomes))
        endpoints = {outcome.item.kind: outcome.endpoint for outcome in report.outcomes}
        self.assertEqual(endpoints, {"web": "/search", "images": "/search", "describe": "/describe", "chat": "/chat"})


class TestPipelineRetryAndFailureAccounting(unittest.IsolatedAsyncioTestCase):

    async def test_transient_failure_retried_with_backoff_then_succeeds(self) -> None:
        client = StubResearchClient()
        client.failures_remaining[("search", "flaky")] = 2
        frontier = QueryFrontier()
        pipeline = ResearchPipeline(client, frontier)
        delays: list[float] = []

        async def fake_sleep(delay: float) -> None:
            delays.append(delay)

        with patch("typosaurus_sandbox.research.pipeline.asyncio.sleep", side_effect=fake_sleep):
            outcome = await pipeline.process(WorkItem("web", "flaky"))
        self.assertTrue(outcome.success)
        self.assertEqual(client.calls.count(("search", "flaky")), 3)
        self.assertEqual(delays, [0.5, 1.0])

    async def test_transient_failure_exhausts_retries_and_reports_failure(self) -> None:
        client = StubResearchClient()
        client.failures_remaining[("search", "persistent")] = 100
        frontier = QueryFrontier()
        pipeline = ResearchPipeline(client, frontier)
        with patch("typosaurus_sandbox.research.pipeline.asyncio.sleep", side_effect=_no_sleep):
            outcome = await pipeline.process(WorkItem("web", "persistent"))
        self.assertFalse(outcome.success)
        self.assertEqual(outcome.status_code, 503)
        self.assertEqual(outcome.error, "search failed")
        self.assertEqual(client.calls.count(("search", "persistent")), RETRY_MAX_ATTEMPTS)

    async def test_client_error_is_not_retried(self) -> None:
        client = StubResearchClient()
        client.error_on.add(("chat", "bad request"))
        frontier = QueryFrontier()
        pipeline = ResearchPipeline(client, frontier)
        with patch("typosaurus_sandbox.research.pipeline.asyncio.sleep", side_effect=_no_sleep):
            outcome = await pipeline.process(WorkItem("chat", "bad request"))
        self.assertFalse(outcome.success)
        self.assertEqual(outcome.status_code, 400)
        self.assertEqual(client.calls.count(("chat", "bad request")), 1)

    async def test_unexpected_exception_recorded_as_failure_outcome(self) -> None:
        client = StubResearchClient()
        client.explode_on.add(("search", "boom"))
        frontier = QueryFrontier()

        async def items() -> AsyncIterator[WorkItem]:
            yield WorkItem("web", "boom")
            yield WorkItem("web", "fine")

        pipeline = ResearchPipeline(client, frontier)
        report = await pipeline.run(items())
        self.assertEqual(report.requests_succeeded, 1)
        self.assertEqual(report.requests_failed, 1)
        self.assertEqual(len(report.outcomes), 2)
        failed = next(outcome for outcome in report.outcomes if not outcome.success)
        self.assertEqual(failed.endpoint, "/search")
        self.assertIsNone(failed.status_code)
        self.assertIn("boom", failed.error or "")
        self.assertEqual(failed.item.value, "boom")


if __name__ == "__main__":
    unittest.main()