# retoor <retoor@molodetz.nl>
import unittest
from typing import Any, AsyncIterator
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 (
Extraction,
ResearchPipeline,
WorkItem,
apply_extraction,
extract_response,
)
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.chat_response = ChatResponse(response="chat answer")
self.describe_response = DescribeResponse(description="described page")
self.calls: list[tuple[str, str]] = []
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))
if ("search", query) in self.error_on:
raise 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)
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"})
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink