# retoor 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()