WIP: feat: Most efficient deep research system ever made #32

Draft
typosaurus wants to merge 13 commits from typosaurus/31-most-efficient-deep-research-system-ever-made into main
pull from: typosaurus/31-most-efficient-deep-research-system-ever-made
merge into: typosaurus:main
Conversation 0 Commits 13 Files Changed 17 +291 -33
4 changed files with 291 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 472a59e122 - Show all commits

16
src/typosaurus_sandbox/research/engine.py
View File

@ -108,7 +108,7 @@ class ResearchEngine:
query = self._frontier.pop_query()
if query is None:
break
yield WorkItem("web", query)
yield WorkItem("web", query, deep=True, ai=True)
yield WorkItem("images", query)
yield WorkItem("chat", query)
for url in urls_to_describe:
@ -156,7 +156,11 @@ class ResearchEngine:
summary.new_urls = round_end.urls_seen - round_start.urls_seen
summary.new_queries = round_end.queries_enqueued - round_start.queries_enqueued
summary.new_contents = round_end.content_seen - round_start.content_seen
summary.closed = summary.new_urls == 0 and summary.new_queries == 0
summary.closed = (
summary.new_urls == 0
and summary.new_queries == 0
and summary.requests_failed == 0
)
report.rounds.append(summary)
logger.info(
"round %d finished new_urls=%d new_queries=%d new_contents=%d closed=%s",
@ -169,16 +173,17 @@ class ResearchEngine:
if summary.closed:
break
report.total_rounds = round_number
report.closed = True
self._finalize(report)
report.closed = report.requests_failed == 0
logger.info(
"research complete subject=%r rounds=%d queries_issued=%d urls_collected=%d contents_seen=%d cache_hits=%d",
"research complete subject=%r rounds=%d queries_issued=%d urls_collected=%d contents_seen=%d cache_hits=%d closed=%s",
report.subject,
report.total_rounds,
report.queries_issued,
report.urls_collected,
report.contents_seen,
report.cache_hits,
report.closed,
)
return report
@ -203,3 +208,6 @@ class ResearchEngine:
report.cache_misses = total_items - report.cache_hits

84
src/typosaurus_sandbox/research/pipeline.py
View File

@ -17,6 +17,11 @@ ContentKind = Literal["web", "images", "describe", "chat"]
URL_PATTERN = re.compile(r"https?://[^\s<>\"']+")
RETRY_MAX_ATTEMPTS = 3
RETRY_BACKOFF_BASE_SECONDS = 0.5
RETRY_BACKOFF_MAX_SECONDS = 8.0
TRANSIENT_STATUS_MIN = 500
@dataclass(frozen=True)
class WorkItem:
@ -182,27 +187,58 @@ class ResearchPipeline:
async def _handle(self, item: WorkItem) -> WorkOutcome:
endpoint = self._endpoint(item)
cache_hit = self._probe_cache(item)
try:
response = await self._fetch(item)
except RsearchError as exc:
outcome = WorkOutcome(
item=item,
endpoint=endpoint,
success=False,
cache_hit=cache_hit,
status_code=exc.status_code,
error=str(exc),
)
response: SearchResponse | ChatResponse | DescribeResponse | None = None
failure: RsearchError | None = None
for attempt in range(1, RETRY_MAX_ATTEMPTS + 1):
try:
response = await self._fetch(item)
failure = None
break
except RsearchError as exc:
failure = exc
if exc.status_code is None or exc.status_code < TRANSIENT_STATUS_MIN:
break
if attempt == RETRY_MAX_ATTEMPTS:
break
delay = min(RETRY_BACKOFF_BASE_SECONDS * (2 ** (attempt - 1)), RETRY_BACKOFF_MAX_SECONDS)
logger.warning(
"transient request failure endpoint=%s kind=%s target=%r status=%s attempt=%d/%d retry_in=%.1fs",
endpoint,
item.kind,
item.value,
exc.status_code,
attempt,
RETRY_MAX_ATTEMPTS,
delay,
)
await asyncio.sleep(delay)
if failure is not None:
logger.error(
"request failed endpoint=%s kind=%s target=%r status=%s cache_hit=%s error=%s",
endpoint,
item.kind,
item.value,
exc.status_code,
failure.status_code,
cache_hit,
exc,
failure,
)
return WorkOutcome(
item=item,
endpoint=endpoint,
success=False,
cache_hit=cache_hit,
status_code=failure.status_code,
error=str(failure),
)
if response is None:
return WorkOutcome(
item=item,
endpoint=endpoint,
success=False,
cache_hit=cache_hit,
status_code=None,
error="no response",
)
return outcome
if isinstance(response, ChatResponse) and response.cached:
cache_hit = True
if isinstance(response, SearchResponse) and response.deep is not None and response.deep.cache_hit:
@ -260,8 +296,20 @@ class ResearchPipeline:
try:
outcome = await self.process(item)
except Exception as exc:
logger.error("pool worker error kind=%s target=%r error=%s", item.kind, item.value, exc)
continue
logger.error(
"pool worker unexpected error kind=%s target=%r error=%s",
item.kind,
item.value,
exc,
)
outcome = WorkOutcome(
item=item,
endpoint=self._endpoint(item),
success=False,
cache_hit=False,
status_code=None,
error=f"unexpected error: {exc}",
)
outcomes.append(outcome)
producer_task = asyncio.create_task(produce())
@ -295,3 +343,7 @@ class ResearchPipeline:
report.contents_seen += outcome.contents_seen
return report

133
tests/test_research_engine.py
View File

@ -2,12 +2,18 @@
import unittest
from collections.abc import Callable
from unittest.mock import patch
from typosaurus_sandbox.research.client import RsearchError
from typosaurus_sandbox.research.config import ResearchConfig
from typosaurus_sandbox.research.engine import ResearchEngine
from typosaurus_sandbox.research.envelopes import ChatResponse, DescribeResponse, SearchResponse, SearchResult
async def _no_sleep(delay: float) -> None:
return None
class FakeResearchClient:
def __init__(
self,
@ -22,7 +28,7 @@ class FakeResearchClient:
self._web_result_factory = web_result_factory
self._chat_text = chat_text
self._describe_text = describe_text
self.calls: list[tuple[str, str, str | None]] = []
self.calls: list[tuple[str, str, str | None, bool, bool]] = []
def search_cached(
self,
@ -53,18 +59,18 @@ class FakeResearchClient:
ai: bool = False,
cache: bool = True,
) -> SearchResponse:
self.calls.append(("search", query, type))
self.calls.append(("search", query, type, deep, ai))
if type == "images":
return SearchResponse(query=query, source="wikimedia", count=0, success=True, results=[])
results = self._web_result_factory(query) if self._web_result_factory is not None else list(self._web_results)
return SearchResponse(query=query, source="duckduckgo", count=len(results), success=True, results=results)
async def chat(self, prompt: str, *, system: str | None = None, json_mode: bool = False, cache: bool = True) -> ChatResponse:
self.calls.append(("chat", prompt, None))
self.calls.append(("chat", prompt, None, False, False))
return ChatResponse(response=self._chat_text, prompt=prompt)
async def describe(self, url: str) -> DescribeResponse:
self.calls.append(("describe", url, None))
self.calls.append(("describe", url, None, False, False))
return DescribeResponse(description=self._describe_text, url=url)
@ -92,7 +98,7 @@ class TestEngineClosureDetection(unittest.IsolatedAsyncioTestCase):
self.assertEqual(report.contents_seen, 0)
self.assertEqual(report.content_types, {"web": 1, "images": 1, "chat": 1})
self.assertEqual(len(client.calls), 3)
self.assertEqual({kind for kind, _, _ in client.calls}, {"search", "chat"})
self.assertEqual({call[0] for call in client.calls}, {"search", "chat"})
async def test_run_discovery_rounds_then_closes(self) -> None:
client = FakeResearchClient(
@ -134,11 +140,11 @@ class TestEngineClosureDetection(unittest.IsolatedAsyncioTestCase):
self.assertEqual(report.cache_misses, 10)
self.assertEqual(report.content_types, {"web": 3, "images": 3, "chat": 3, "describe": 1})
self.assertEqual(len(client.calls), 10)
self.assertEqual(sum(1 for kind, _, type_value in client.calls if kind == "search" and type_value is None), 3)
self.assertEqual(sum(1 for kind, _, type_value in client.calls if kind == "search" and type_value == "images"), 3)
self.assertEqual(sum(1 for kind, _, _ in client.calls if kind == "chat"), 3)
self.assertEqual(sum(1 for kind, value, _ in client.calls if kind == "describe"), 1)
self.assertIn(("describe", "https://example.com/alpha", None), client.calls)
self.assertEqual(sum(1 for kind, _, type_value, _, _ in client.calls if kind == "search" and type_value is None), 3)
self.assertEqual(sum(1 for kind, _, type_value, _, _ in client.calls if kind == "search" and type_value == "images"), 3)
self.assertEqual(sum(1 for kind, _, _, _, _ in client.calls if kind == "chat"), 3)
self.assertEqual(sum(1 for kind, _, _, _, _ in client.calls if kind == "describe"), 1)
self.assertIn(("describe", "https://example.com/alpha", None, False, False), client.calls)
async def test_new_content_alone_does_not_prevent_closure(self) -> None:
def factory(query: str) -> list[SearchResult]:
@ -180,6 +186,107 @@ class TestEngineClosureDetection(unittest.IsolatedAsyncioTestCase):
self.assertTrue(report_dict["closed"])
class FailingWebClient(FakeResearchClient):
def __init__(
self,
*,
web_results: list[SearchResult] | None = None,
web_result_factory: Callable[[str], list[SearchResult]] | None = None,
chat_text: str = "",
describe_text: str = "",
failures_before_success: int = 0,
) -> None:
super().__init__(
web_results=web_results,
web_result_factory=web_result_factory,
chat_text=chat_text,
describe_text=describe_text,
)
self._web_failures_left = failures_before_success
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, type, deep, ai))
if type == "images":
return SearchResponse(query=query, source="wikimedia", count=0, success=True, results=[])
if self._web_failures_left > 0:
self._web_failures_left -= 1
raise RsearchError("search failed", 503)
results = self._web_result_factory(query) if self._web_result_factory is not None else list(self._web_results)
return SearchResponse(query=query, source="duckduckgo", count=len(results), success=True, results=results)
class TestEngineDeepAiWiring(unittest.IsolatedAsyncioTestCase):
async def test_web_search_work_items_issue_deep_and_ai_for_seed_and_subtopics(self) -> None:
def factory(query: str) -> list[SearchResult]:
return [
SearchResult(
title="subtopic alpha",
url="https://example.com/subtopic",
description="subtopic details",
content="subtopic body",
)
]
client = FakeResearchClient(web_result_factory=factory, chat_text="", describe_text="")
engine = ResearchEngine(client=client)
report = await engine.run("seed topic")
self.assertTrue(report.closed)
web_calls = [call for call in client.calls if call[0] == "search" and call[2] is None]
self.assertEqual(len(web_calls), 3)
self.assertEqual({call[1] for call in web_calls}, {"seed topic", "subtopic alpha", "subtopic details"})
self.assertTrue(all(call[3] and call[4] for call in web_calls))
class TestEngineClosureOnFailures(unittest.IsolatedAsyncioTestCase):
async def test_round_and_report_not_closed_when_request_failed(self) -> None:
client = FailingWebClient(failures_before_success=100)
engine = ResearchEngine(client=client)
with patch("typosaurus_sandbox.research.pipeline.asyncio.sleep", side_effect=_no_sleep):
report = await engine.run("subject")
self.assertFalse(report.closed)
self.assertEqual(report.requests_failed, 1)
self.assertEqual(report.requests_succeeded, 2)
self.assertEqual(report.total_rounds, 1)
first = report.rounds[0]
self.assertEqual(first.requests_failed, 1)
self.assertEqual(first.requests_succeeded, 2)
self.assertEqual(first.new_urls, 0)
self.assertEqual(first.new_queries, 0)
self.assertFalse(first.closed)
async def test_later_closed_round_keeps_report_unclosed_after_earlier_failure(self) -> None:
client = FailingWebClient(
chat_text="Reference at https://chat.example.org/note",
describe_text="",
failures_before_success=3,
)
engine = ResearchEngine(client=client)
with patch("typosaurus_sandbox.research.pipeline.asyncio.sleep", side_effect=_no_sleep):
report = await engine.run("subject")
self.assertEqual(report.total_rounds, 2)
first = report.rounds[0]
self.assertEqual(first.requests_failed, 1)
self.assertFalse(first.closed)
second = report.rounds[1]
self.assertEqual(second.requests_failed, 0)
self.assertTrue(second.closed)
self.assertEqual(report.requests_failed, 1)
self.assertFalse(report.closed)
class TestEngineInputValidation(unittest.IsolatedAsyncioTestCase):
async def test_empty_subject_raises_without_requests(self) -> None:
@ -196,3 +303,9 @@ if __name__ == "__main__":

91
tests/test_research_pipeline.py
View File

@ -2,12 +2,14 @@
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,
@ -16,15 +18,31 @@ from typosaurus_sandbox.research.pipeline import (
)
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,
@ -57,8 +75,7 @@ class StubResearchClient:
cache: bool = True,
) -> SearchResponse:
self.calls.append(("search", query))
if ("search", query) in self.error_on:
raise RsearchError("search failed", 503)
self._maybe_fail(("search", query), RsearchError("search failed", 503))
return SearchResponse(
query=query,
source="duckduckgo",
@ -242,7 +259,8 @@ class TestPipelineOutcomes(unittest.IsolatedAsyncioTestCase):
yield WorkItem("web", "good")
pipeline = ResearchPipeline(client, frontier)
report = await pipeline.run(items())
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)
@ -271,6 +289,73 @@ class TestPipelineOutcomes(unittest.IsolatedAsyncioTestCase):
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()