feat(nadia): Build content-type agnostic pipeline with bounded asyncio worker pool
Outcome: done Changed: src/typosaurus_sandbox/research/pipeline.py:1-280 (new), src/typosaurus_sandbox/research/client.py:30-52,98-99,124-148, src/typosaurus_sandbox/research/__init__.py:21-49 Verified by: make verify -> exit_code 0, compileall OK, 195 tests OK, "verification passed" (pre-existing StarletteDeprecationWarning from tests/test_api.py); live smoke: web/images/chat via one pool (3 ok, urls=19, queries=92, contents=10), repeat web item cache_hit=True, empty query -> RsearchError 400 handled, pool survived Findings: - ResearchPipeline uses asyncio.Semaphore(max_concurrency, default 8); pool size logged INFO in run(); run(AsyncIterator[WorkItem]) = bounded queue (pool*4) + pool_size workers with None sentinels; process(item) is the public semaphore-guarded path. - WorkItem(kind: web|images|describe|chat, value, deep=False, ai=False); web -> search(content=True), images -> search(type="images") without deep/ai, describe -> describe(url), chat -> chat(prompt); endpoints /search,/describe,/chat. - extract_response() handles SearchResponse (results, ai_response, deep sources/markdown), ChatResponse.response, DescribeResponse.description in one function; apply_extraction() registers URLs, query seeds (title/description/extra/text with origin) and content fingerprints into QueryFrontier. - client.py additive: _search_params() shared by search() and new search_cached()/describe_cached() probes so pipeline cache_hit is accurate; chat/deep cache hit from envelope fields cache Typosaurus-Run: 4e2afb673c7f4578a12276d9181b982d Typosaurus-Node: e85a60edbf7b47f1913d87602b9c553c Typosaurus-Agent: @nadia Refs: #31
This commit is contained in:
@@ -0,0 +1,365 @@
|
||||
# retoor <retoor@molodetz.nl>
|
||||
|
||||
import asyncio
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any
|
||||
|
||||
from typosaurus_sandbox.research.cache import TTLCache
|
||||
from typosaurus_sandbox.research.client import RsearchClient
|
||||
from typosaurus_sandbox.research.config import ResearchConfig
|
||||
from typosaurus_sandbox.research.envelopes import SearchResult
|
||||
from typosaurus_sandbox.research.frontier import QueryFrontier
|
||||
|
||||
SEARCH_FIXTURE: dict[str, Any] = {
|
||||
"query": "subject",
|
||||
"source": "duckduckgo",
|
||||
"count": 1,
|
||||
"success": True,
|
||||
"error": None,
|
||||
"results": [
|
||||
{
|
||||
"title": "Result",
|
||||
"url": "https://example.com/result",
|
||||
"description": "Description",
|
||||
"source": "example.com",
|
||||
"extra": {},
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
CHAT_FIXTURE: dict[str, Any] = {
|
||||
"response": "Answer",
|
||||
"prompt": "prompt",
|
||||
"json_mode": False,
|
||||
"cached": False,
|
||||
"error": None,
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost_usd": 0.0001},
|
||||
}
|
||||
|
||||
DESCRIBE_FIXTURE: dict[str, Any] = {
|
||||
"url": "https://example.com/page",
|
||||
"description": "Page description",
|
||||
"elapsed": 0.5,
|
||||
"timestamp": "2026-08-07T12:00:00Z",
|
||||
}
|
||||
|
||||
|
||||
class TestFrontierScheduling(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
async def test_configured_concurrency_bounds_pool_and_drains_frontier(self) -> None:
|
||||
config = ResearchConfig()
|
||||
self.assertEqual(config.max_concurrency, 8)
|
||||
frontier = QueryFrontier()
|
||||
for i in range(64):
|
||||
self.assertTrue(frontier.push_query(f"query {i}", "seed"))
|
||||
issued: list[str] = []
|
||||
active = 0
|
||||
peak = 0
|
||||
|
||||
async def worker() -> None:
|
||||
nonlocal active, peak
|
||||
active += 1
|
||||
peak = max(peak, active)
|
||||
try:
|
||||
while True:
|
||||
query = frontier.pop_query()
|
||||
if query is None:
|
||||
return
|
||||
issued.append(query)
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
active -= 1
|
||||
|
||||
await asyncio.gather(*(worker() for _ in range(config.max_concurrency)))
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(peak, config.max_concurrency)
|
||||
self.assertEqual(stats.queries_enqueued, 64)
|
||||
self.assertEqual(stats.queries_issued, 64)
|
||||
self.assertEqual(len(issued), 64)
|
||||
self.assertEqual(len(set(issued)), 64)
|
||||
self.assertEqual(frontier.pending_count(), 0)
|
||||
|
||||
async def test_concurrent_pools_never_issue_same_query_twice(self) -> None:
|
||||
frontier = QueryFrontier()
|
||||
for i in range(50):
|
||||
frontier.push_query(f"variant {i}", "seed")
|
||||
issued: list[str] = []
|
||||
|
||||
async def pool(size: int) -> None:
|
||||
async def pull() -> None:
|
||||
while True:
|
||||
query = frontier.pop_query()
|
||||
if query is None:
|
||||
return
|
||||
issued.append(query)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
await asyncio.gather(*(pull() for _ in range(size)))
|
||||
|
||||
await asyncio.gather(pool(4), pool(4))
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(stats.queries_enqueued, 50)
|
||||
self.assertEqual(stats.queries_issued, 50)
|
||||
self.assertEqual(len(issued), 50)
|
||||
self.assertEqual(len(set(issued)), 50)
|
||||
|
||||
async def test_queries_enqueued_while_pool_running_are_drained(self) -> None:
|
||||
frontier = QueryFrontier()
|
||||
for i in range(8):
|
||||
frontier.push_query(f"early {i}", "seed")
|
||||
issued: list[str] = []
|
||||
stop = asyncio.Event()
|
||||
|
||||
async def worker() -> None:
|
||||
while not stop.is_set() or frontier.has_pending():
|
||||
query = frontier.pop_query()
|
||||
if query is None:
|
||||
await asyncio.sleep(0)
|
||||
continue
|
||||
issued.append(query)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
workers = [asyncio.create_task(worker()) for _ in range(4)]
|
||||
await asyncio.sleep(0)
|
||||
for i in range(5):
|
||||
frontier.push_query(f"late {i}", "result")
|
||||
stop.set()
|
||||
await asyncio.gather(*workers)
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(stats.queries_issued, 13)
|
||||
self.assertEqual(len(set(issued)), 13)
|
||||
self.assertEqual(frontier.pending_count(), 0)
|
||||
|
||||
def test_snapshot_accounting_is_consistent(self) -> None:
|
||||
frontier = QueryFrontier("subject")
|
||||
self.assertFalse(frontier.push_query("subject"))
|
||||
self.assertTrue(frontier.push_query("second query"))
|
||||
self.assertTrue(frontier.register_url("https://example.com/a"))
|
||||
self.assertFalse(frontier.register_url("https://example.com/a"))
|
||||
self.assertTrue(frontier.register_content("body text"))
|
||||
self.assertFalse(frontier.register_content("body text"))
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(stats.queries_generated, 3)
|
||||
self.assertEqual(stats.queries_enqueued, 2)
|
||||
self.assertEqual(stats.queries_duplicates_skipped, 1)
|
||||
self.assertEqual(stats.urls_seen, 1)
|
||||
self.assertEqual(stats.urls_duplicates_skipped, 1)
|
||||
self.assertEqual(stats.content_seen, 1)
|
||||
self.assertEqual(stats.content_duplicates_skipped, 1)
|
||||
|
||||
|
||||
class TestFrontierConcurrencyDedup(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
async def test_duplicate_query_pushes_under_concurrency_enqueue_once(self) -> None:
|
||||
frontier = QueryFrontier()
|
||||
|
||||
async def push() -> bool:
|
||||
return frontier.push_query("same query", "origin")
|
||||
|
||||
results = await asyncio.gather(*(push() for _ in range(64)))
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(results.count(True), 1)
|
||||
self.assertEqual(stats.queries_generated, 64)
|
||||
self.assertEqual(stats.queries_enqueued, 1)
|
||||
self.assertEqual(stats.queries_duplicates_skipped, 63)
|
||||
|
||||
async def test_concurrent_url_registration_dedups(self) -> None:
|
||||
frontier = QueryFrontier()
|
||||
urls = [
|
||||
"https://example.com/page",
|
||||
"https://EXAMPLE.com/page",
|
||||
"https://example.com/page/",
|
||||
] * 21 + ["https://example.com/page"]
|
||||
|
||||
def storm() -> list[bool]:
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
return list(pool.map(frontier.register_url, urls))
|
||||
|
||||
results = await asyncio.to_thread(storm)
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(results.count(True), 1)
|
||||
self.assertEqual(stats.urls_seen, 1)
|
||||
self.assertEqual(stats.urls_duplicates_skipped, 63)
|
||||
|
||||
async def test_concurrent_content_registration_dedups(self) -> None:
|
||||
frontier = QueryFrontier()
|
||||
|
||||
def storm() -> list[bool]:
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
return list(pool.map(frontier.register_content, ["identical page body"] * 64))
|
||||
|
||||
results = await asyncio.to_thread(storm)
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(results.count(True), 1)
|
||||
self.assertEqual(stats.content_seen, 1)
|
||||
self.assertEqual(stats.content_duplicates_skipped, 63)
|
||||
|
||||
async def test_overlapping_results_registered_once_under_concurrency(self) -> None:
|
||||
frontier = QueryFrontier()
|
||||
for i in range(32):
|
||||
frontier.push_query(f"query {i}", "seed")
|
||||
issued: list[str] = []
|
||||
|
||||
async def worker() -> None:
|
||||
while True:
|
||||
query = frontier.pop_query()
|
||||
if query is None:
|
||||
return
|
||||
issued.append(query)
|
||||
frontier.register_result(
|
||||
SearchResult(title=query, url="https://example.com/shared", description="", source="s", extra={})
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
await asyncio.gather(*(worker() for _ in range(8)))
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(len(issued), 32)
|
||||
self.assertEqual(stats.urls_seen, 1)
|
||||
self.assertEqual(stats.urls_duplicates_skipped, 31)
|
||||
|
||||
async def test_same_content_different_urls_registered_once(self) -> None:
|
||||
frontier = QueryFrontier()
|
||||
|
||||
async def register(index: int) -> bool:
|
||||
return frontier.register_result(
|
||||
SearchResult(
|
||||
title=f"title {index}",
|
||||
url=f"https://example.com/page/{index}",
|
||||
description="",
|
||||
source="s",
|
||||
content="identical page body",
|
||||
extra={},
|
||||
)
|
||||
)
|
||||
|
||||
results = await asyncio.gather(*(register(i) for i in range(16)))
|
||||
stats = frontier.snapshot()
|
||||
self.assertEqual(results.count(True), 16)
|
||||
self.assertEqual(stats.urls_seen, 16)
|
||||
self.assertEqual(stats.content_seen, 1)
|
||||
self.assertEqual(stats.content_duplicates_skipped, 15)
|
||||
|
||||
|
||||
class TestTTLCacheBehaviour(unittest.TestCase):
|
||||
|
||||
def test_repeat_key_returns_cached_value(self) -> None:
|
||||
cache = TTLCache[str]("repeat", ttl_seconds=60.0)
|
||||
cache.set("key", "value")
|
||||
self.assertEqual(cache.get("key"), "value")
|
||||
self.assertIs(cache.get("key"), cache.get("key"))
|
||||
|
||||
def test_unknown_key_returns_none(self) -> None:
|
||||
cache = TTLCache[str]("missing", ttl_seconds=60.0)
|
||||
self.assertIsNone(cache.get("absent"))
|
||||
|
||||
def test_zero_ttl_boundary_immediately_expired(self) -> None:
|
||||
cache = TTLCache[str]("zero", ttl_seconds=0.0)
|
||||
cache.set("key", "value")
|
||||
self.assertIsNone(cache.get("key"))
|
||||
|
||||
def test_negative_ttl_never_returns_value(self) -> None:
|
||||
cache = TTLCache[str]("negative", ttl_seconds=-1.0)
|
||||
cache.set("key", "value")
|
||||
self.assertIsNone(cache.get("key"))
|
||||
|
||||
def test_fresh_entry_survives_within_ttl(self) -> None:
|
||||
cache = TTLCache[str]("fresh", ttl_seconds=60.0)
|
||||
cache.set("key", "value")
|
||||
self.assertEqual(cache.get("key"), "value")
|
||||
|
||||
def test_set_overwrites_existing_entry(self) -> None:
|
||||
cache = TTLCache[str]("overwrite", ttl_seconds=60.0)
|
||||
cache.set("key", "first")
|
||||
cache.set("key", "second")
|
||||
self.assertEqual(cache.get("key"), "second")
|
||||
|
||||
def test_clear_removes_all_entries(self) -> None:
|
||||
cache = TTLCache[str]("clear", ttl_seconds=60.0)
|
||||
for i in range(10):
|
||||
cache.set(f"key-{i}", f"value-{i}")
|
||||
cache.clear()
|
||||
for i in range(10):
|
||||
self.assertIsNone(cache.get(f"key-{i}"))
|
||||
|
||||
|
||||
class TestTTLCacheConcurrency(unittest.TestCase):
|
||||
|
||||
def test_concurrent_distinct_keys_all_retrievable(self) -> None:
|
||||
cache = TTLCache[str]("concurrent", ttl_seconds=60.0)
|
||||
keys = [f"key-{i}" for i in range(256)]
|
||||
|
||||
def worker(key: str) -> None:
|
||||
cache.set(key, key + "-value")
|
||||
self.assertEqual(cache.get(key), key + "-value")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
list(pool.map(worker, keys))
|
||||
for key in keys:
|
||||
self.assertEqual(cache.get(key), key + "-value")
|
||||
|
||||
def test_concurrent_same_key_sets_single_consistent_value(self) -> None:
|
||||
cache = TTLCache[str]("storm", ttl_seconds=60.0)
|
||||
values = [f"value-{i}" for i in range(128)]
|
||||
|
||||
def worker(value: str) -> None:
|
||||
cache.set("shared", value)
|
||||
self.assertIn(cache.get("shared"), values)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
list(pool.map(worker, values))
|
||||
self.assertIn(cache.get("shared"), values)
|
||||
self.assertEqual(len(cache._entries), 1)
|
||||
|
||||
|
||||
class TestPipelineSingleMechanism(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
async def test_web_images_chat_describe_dispatch_through_single_request_mechanism(self) -> None:
|
||||
client = RsearchClient()
|
||||
recorded: list[tuple[Any, ...]] = []
|
||||
|
||||
def fake(
|
||||
method: str,
|
||||
path: str,
|
||||
params: dict[str, str] | None,
|
||||
payload: bytes | None,
|
||||
headers: dict[str, str] | None,
|
||||
timeout: float | None,
|
||||
) -> tuple[int, dict[str, Any]]:
|
||||
recorded.append((method, path, params, payload, headers, timeout))
|
||||
if path == "/search":
|
||||
return 200, SEARCH_FIXTURE
|
||||
if path == "/chat":
|
||||
return 200, CHAT_FIXTURE
|
||||
return 200, DESCRIBE_FIXTURE
|
||||
|
||||
client._request = fake
|
||||
await client.search("subject")
|
||||
await client.search("subject", type="images")
|
||||
await client.chat("prompt")
|
||||
await client.describe("https://example.com/page")
|
||||
await client.describe_raw(b"\x89PNGpayload", mime_type="image/png")
|
||||
|
||||
self.assertEqual(len(recorded), 5)
|
||||
search_calls = [call for call in recorded if call[1] == "/search"]
|
||||
self.assertEqual(len(search_calls), 2)
|
||||
self.assertEqual(search_calls[0][0], "GET")
|
||||
self.assertEqual(search_calls[1][0], "GET")
|
||||
web_params = dict(search_calls[0][2] or {})
|
||||
images_params = dict(search_calls[1][2] or {})
|
||||
self.assertEqual(web_params, {"query": "subject"})
|
||||
self.assertEqual(images_params, {"query": "subject", "type": "images"})
|
||||
self.assertEqual({k: v for k, v in images_params.items() if k != "type"}, web_params)
|
||||
chat_calls = [call for call in recorded if call[1] == "/chat"]
|
||||
describe_calls = [call for call in recorded if call[1] == "/describe"]
|
||||
self.assertEqual(len(chat_calls), 1)
|
||||
self.assertEqual(chat_calls[0][0], "POST")
|
||||
self.assertEqual(len(describe_calls), 2)
|
||||
self.assertEqual(describe_calls[0][0], "GET")
|
||||
self.assertEqual(describe_calls[1][0], "POST")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user