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:
parent
7977217013
commit
c0a59b4138
@ -19,14 +19,28 @@ from typosaurus_sandbox.research.frontier import (
|
|||||||
normalize_url,
|
normalize_url,
|
||||||
query_variants_from_result,
|
query_variants_from_result,
|
||||||
)
|
)
|
||||||
|
from typosaurus_sandbox.research.pipeline import (
|
||||||
|
ContentKind,
|
||||||
|
Extraction,
|
||||||
|
PipelineReport,
|
||||||
|
ResearchPipeline,
|
||||||
|
WorkItem,
|
||||||
|
WorkOutcome,
|
||||||
|
apply_extraction,
|
||||||
|
extract_response,
|
||||||
|
)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"ChatResponse",
|
"ChatResponse",
|
||||||
"ChatUsage",
|
"ChatUsage",
|
||||||
|
"ContentKind",
|
||||||
"DedupStats",
|
"DedupStats",
|
||||||
"DeepReport",
|
"DeepReport",
|
||||||
"DescribeResponse",
|
"DescribeResponse",
|
||||||
|
"Extraction",
|
||||||
|
"PipelineReport",
|
||||||
"QueryFrontier",
|
"QueryFrontier",
|
||||||
|
"ResearchPipeline",
|
||||||
"RsearchClient",
|
"RsearchClient",
|
||||||
"RsearchError",
|
"RsearchError",
|
||||||
"ResearchConfig",
|
"ResearchConfig",
|
||||||
@ -34,6 +48,10 @@ __all__ = [
|
|||||||
"SearchResponse",
|
"SearchResponse",
|
||||||
"SearchResult",
|
"SearchResult",
|
||||||
"TTLCache",
|
"TTLCache",
|
||||||
|
"WorkItem",
|
||||||
|
"WorkOutcome",
|
||||||
|
"apply_extraction",
|
||||||
|
"extract_response",
|
||||||
"fingerprint_text",
|
"fingerprint_text",
|
||||||
"normalize_url",
|
"normalize_url",
|
||||||
"query_variants_from_result",
|
"query_variants_from_result",
|
||||||
@ -41,3 +59,4 @@ __all__ = [
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@ -40,6 +40,35 @@ def _content_hash(image_bytes: bytes) -> str:
|
|||||||
return hashlib.sha256(image_bytes).hexdigest()
|
return hashlib.sha256(image_bytes).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
def _search_params(
|
||||||
|
query: str,
|
||||||
|
*,
|
||||||
|
source: str | None,
|
||||||
|
count: int | None,
|
||||||
|
content: bool,
|
||||||
|
type: str | None,
|
||||||
|
deep: bool,
|
||||||
|
ai: bool,
|
||||||
|
cache: bool,
|
||||||
|
) -> dict[str, str]:
|
||||||
|
params: dict[str, str] = {"query": query}
|
||||||
|
if source is not None:
|
||||||
|
params["source"] = source
|
||||||
|
if count is not None:
|
||||||
|
params["count"] = str(count)
|
||||||
|
if content:
|
||||||
|
params["content"] = "true"
|
||||||
|
if type is not None:
|
||||||
|
params["type"] = type
|
||||||
|
if deep:
|
||||||
|
params["deep"] = "true"
|
||||||
|
if ai:
|
||||||
|
params["ai"] = "true"
|
||||||
|
if not cache:
|
||||||
|
params["cache"] = "false"
|
||||||
|
return params
|
||||||
|
|
||||||
|
|
||||||
class RsearchClient:
|
class RsearchClient:
|
||||||
def __init__(self, config: ResearchConfig | None = None) -> None:
|
def __init__(self, config: ResearchConfig | None = None) -> None:
|
||||||
self._config = config if config is not None else ResearchConfig()
|
self._config = config if config is not None else ResearchConfig()
|
||||||
@ -66,21 +95,7 @@ class RsearchClient:
|
|||||||
ai: bool = False,
|
ai: bool = False,
|
||||||
cache: bool = True,
|
cache: bool = True,
|
||||||
) -> SearchResponse:
|
) -> SearchResponse:
|
||||||
params: dict[str, str] = {"query": query}
|
params = _search_params(query, source=source, count=count, content=content, type=type, deep=deep, ai=ai, cache=cache)
|
||||||
if source is not None:
|
|
||||||
params["source"] = source
|
|
||||||
if count is not None:
|
|
||||||
params["count"] = str(count)
|
|
||||||
if content:
|
|
||||||
params["content"] = "true"
|
|
||||||
if type is not None:
|
|
||||||
params["type"] = type
|
|
||||||
if deep:
|
|
||||||
params["deep"] = "true"
|
|
||||||
if ai:
|
|
||||||
params["ai"] = "true"
|
|
||||||
if not cache:
|
|
||||||
params["cache"] = "false"
|
|
||||||
key = urllib.parse.urlencode(sorted(params.items()))
|
key = urllib.parse.urlencode(sorted(params.items()))
|
||||||
if cache:
|
if cache:
|
||||||
cached_response = self._search_cache.get(key)
|
cached_response = self._search_cache.get(key)
|
||||||
@ -106,6 +121,27 @@ class RsearchClient:
|
|||||||
)
|
)
|
||||||
return response
|
return response
|
||||||
|
|
||||||
|
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,
|
||||||
|
) -> SearchResponse | None:
|
||||||
|
if not cache:
|
||||||
|
return None
|
||||||
|
params = _search_params(query, source=source, count=count, content=content, type=type, deep=deep, ai=ai, cache=cache)
|
||||||
|
key = urllib.parse.urlencode(sorted(params.items()))
|
||||||
|
return self._search_cache.get(key)
|
||||||
|
|
||||||
|
def describe_cached(self, url: str) -> DescribeResponse | None:
|
||||||
|
return self._describe_cache.get(f"url:{url}")
|
||||||
|
|
||||||
async def chat(
|
async def chat(
|
||||||
self,
|
self,
|
||||||
prompt: str,
|
prompt: str,
|
||||||
@ -210,3 +246,6 @@ class RsearchClient:
|
|||||||
raise RsearchError(self._error_message(data), status)
|
raise RsearchError(self._error_message(data), status)
|
||||||
return status, data
|
return status, data
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
297
src/typosaurus_sandbox/research/pipeline.py
Normal file
297
src/typosaurus_sandbox/research/pipeline.py
Normal file
@ -0,0 +1,297 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import re
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import AsyncIterator, Literal
|
||||||
|
|
||||||
|
from typosaurus_sandbox.research.client import RsearchClient, RsearchError
|
||||||
|
from typosaurus_sandbox.research.config import ResearchConfig
|
||||||
|
from typosaurus_sandbox.research.envelopes import ChatResponse, DescribeResponse, SearchResponse
|
||||||
|
from typosaurus_sandbox.research.frontier import QueryFrontier, query_variants_from_result
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
ContentKind = Literal["web", "images", "describe", "chat"]
|
||||||
|
|
||||||
|
URL_PATTERN = re.compile(r"https?://[^\s<>\"']+")
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class WorkItem:
|
||||||
|
kind: ContentKind
|
||||||
|
value: str
|
||||||
|
deep: bool = False
|
||||||
|
ai: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class Extraction:
|
||||||
|
urls: tuple[str, ...] = ()
|
||||||
|
query_seeds: tuple[tuple[str, str], ...] = ()
|
||||||
|
content_texts: tuple[str, ...] = ()
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class WorkOutcome:
|
||||||
|
item: WorkItem
|
||||||
|
endpoint: str
|
||||||
|
success: bool
|
||||||
|
cache_hit: bool
|
||||||
|
status_code: int | None = None
|
||||||
|
error: str | None = None
|
||||||
|
urls_found: int = 0
|
||||||
|
queries_seeded: int = 0
|
||||||
|
contents_seen: int = 0
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class PipelineReport:
|
||||||
|
outcomes: list[WorkOutcome] = field(default_factory=list)
|
||||||
|
requests_succeeded: int = 0
|
||||||
|
requests_failed: int = 0
|
||||||
|
urls_found: int = 0
|
||||||
|
queries_seeded: int = 0
|
||||||
|
contents_seen: int = 0
|
||||||
|
|
||||||
|
|
||||||
|
def _urls_from_text(text: str) -> list[str]:
|
||||||
|
cleaned: list[str] = []
|
||||||
|
for match in URL_PATTERN.findall(text):
|
||||||
|
cleaned.append(match.rstrip(".,;:!?)]}\"'"))
|
||||||
|
return cleaned
|
||||||
|
|
||||||
|
|
||||||
|
def extract_response(
|
||||||
|
item: WorkItem,
|
||||||
|
response: SearchResponse | ChatResponse | DescribeResponse,
|
||||||
|
) -> Extraction:
|
||||||
|
urls: list[str] = []
|
||||||
|
query_seeds: list[tuple[str, str]] = []
|
||||||
|
content_texts: list[str] = []
|
||||||
|
if isinstance(response, SearchResponse):
|
||||||
|
for result in response.results:
|
||||||
|
if result.url:
|
||||||
|
urls.append(result.url)
|
||||||
|
query_seeds.extend(query_variants_from_result(result))
|
||||||
|
if result.content:
|
||||||
|
content_texts.append(result.content)
|
||||||
|
if response.ai_response:
|
||||||
|
content_texts.append(response.ai_response)
|
||||||
|
query_seeds.append((response.ai_response, "ai_response"))
|
||||||
|
urls.extend(_urls_from_text(response.ai_response))
|
||||||
|
if response.deep is not None:
|
||||||
|
for source in response.deep.sources:
|
||||||
|
if source.url:
|
||||||
|
urls.append(source.url)
|
||||||
|
query_seeds.extend(query_variants_from_result(source))
|
||||||
|
if response.deep.markdown:
|
||||||
|
content_texts.append(response.deep.markdown)
|
||||||
|
urls.extend(_urls_from_text(response.deep.markdown))
|
||||||
|
elif isinstance(response, ChatResponse):
|
||||||
|
if response.response:
|
||||||
|
content_texts.append(response.response)
|
||||||
|
query_seeds.append((response.response, "chat"))
|
||||||
|
urls.extend(_urls_from_text(response.response))
|
||||||
|
elif isinstance(response, DescribeResponse):
|
||||||
|
if response.description:
|
||||||
|
content_texts.append(response.description)
|
||||||
|
query_seeds.append((response.description, "describe"))
|
||||||
|
urls.extend(_urls_from_text(response.description))
|
||||||
|
return Extraction(
|
||||||
|
urls=tuple(dict.fromkeys(urls)),
|
||||||
|
query_seeds=tuple(query_seeds),
|
||||||
|
content_texts=tuple(content_texts),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def apply_extraction(frontier: QueryFrontier, extraction: Extraction) -> tuple[int, int, int]:
|
||||||
|
new_urls = 0
|
||||||
|
new_queries = 0
|
||||||
|
new_contents = 0
|
||||||
|
for url in extraction.urls:
|
||||||
|
if frontier.register_url(url):
|
||||||
|
new_urls += 1
|
||||||
|
for text, origin in extraction.query_seeds:
|
||||||
|
if frontier.push_query(text, origin):
|
||||||
|
new_queries += 1
|
||||||
|
for text in extraction.content_texts:
|
||||||
|
if frontier.register_content(text):
|
||||||
|
new_contents += 1
|
||||||
|
return new_urls, new_queries, new_contents
|
||||||
|
|
||||||
|
|
||||||
|
class ResearchPipeline:
|
||||||
|
def __init__(self, client: RsearchClient, frontier: QueryFrontier, config: ResearchConfig | None = None) -> None:
|
||||||
|
self._client = client
|
||||||
|
self._frontier = frontier
|
||||||
|
self._config = config if config is not None else client.config
|
||||||
|
self._pool_size = max(1, self._config.max_concurrency)
|
||||||
|
self._semaphore = asyncio.Semaphore(self._pool_size)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pool_size(self) -> int:
|
||||||
|
return self._pool_size
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _endpoint(item: WorkItem) -> str:
|
||||||
|
if item.kind in ("web", "images"):
|
||||||
|
return "/search"
|
||||||
|
if item.kind == "describe":
|
||||||
|
return "/describe"
|
||||||
|
return "/chat"
|
||||||
|
|
||||||
|
def _probe_cache(self, item: WorkItem) -> bool:
|
||||||
|
if item.kind == "web":
|
||||||
|
return (
|
||||||
|
self._client.search_cached(
|
||||||
|
item.value,
|
||||||
|
content=True,
|
||||||
|
count=self._config.default_count,
|
||||||
|
deep=item.deep,
|
||||||
|
ai=item.ai,
|
||||||
|
)
|
||||||
|
is not None
|
||||||
|
)
|
||||||
|
if item.kind == "images":
|
||||||
|
return self._client.search_cached(item.value, type="images", count=self._config.default_count) is not None
|
||||||
|
if item.kind == "describe":
|
||||||
|
return self._client.describe_cached(item.value) is not None
|
||||||
|
return False
|
||||||
|
|
||||||
|
async def _fetch(self, item: WorkItem) -> SearchResponse | ChatResponse | DescribeResponse:
|
||||||
|
if item.kind == "web":
|
||||||
|
return await self._client.search(
|
||||||
|
item.value,
|
||||||
|
content=True,
|
||||||
|
count=self._config.default_count,
|
||||||
|
deep=item.deep,
|
||||||
|
ai=item.ai,
|
||||||
|
)
|
||||||
|
if item.kind == "images":
|
||||||
|
return await self._client.search(item.value, type="images", count=self._config.default_count)
|
||||||
|
if item.kind == "describe":
|
||||||
|
return await self._client.describe(item.value)
|
||||||
|
return await self._client.chat(item.value)
|
||||||
|
|
||||||
|
async def process(self, item: WorkItem) -> WorkOutcome:
|
||||||
|
async with self._semaphore:
|
||||||
|
return await self._handle(item)
|
||||||
|
|
||||||
|
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),
|
||||||
|
)
|
||||||
|
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,
|
||||||
|
cache_hit,
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
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:
|
||||||
|
cache_hit = True
|
||||||
|
extraction = extract_response(item, response)
|
||||||
|
urls_found, queries_seeded, contents_seen = apply_extraction(self._frontier, extraction)
|
||||||
|
logger.debug(
|
||||||
|
"extraction endpoint=%s kind=%s target=%r urls=%s query_seeds=%d content_texts=%d",
|
||||||
|
endpoint,
|
||||||
|
item.kind,
|
||||||
|
item.value,
|
||||||
|
list(extraction.urls),
|
||||||
|
len(extraction.query_seeds),
|
||||||
|
len(extraction.content_texts),
|
||||||
|
)
|
||||||
|
outcome = WorkOutcome(
|
||||||
|
item=item,
|
||||||
|
endpoint=endpoint,
|
||||||
|
success=True,
|
||||||
|
cache_hit=cache_hit,
|
||||||
|
urls_found=urls_found,
|
||||||
|
queries_seeded=queries_seeded,
|
||||||
|
contents_seen=contents_seen,
|
||||||
|
)
|
||||||
|
logger.info(
|
||||||
|
"request done endpoint=%s kind=%s target=%r status=ok cache_hit=%s urls=%d queries=%d contents=%d",
|
||||||
|
endpoint,
|
||||||
|
item.kind,
|
||||||
|
item.value,
|
||||||
|
cache_hit,
|
||||||
|
urls_found,
|
||||||
|
queries_seeded,
|
||||||
|
contents_seen,
|
||||||
|
)
|
||||||
|
return outcome
|
||||||
|
|
||||||
|
async def run(self, item_source: AsyncIterator[WorkItem]) -> PipelineReport:
|
||||||
|
logger.info("worker pool size=%d max_concurrency=%d", self._pool_size, self._config.max_concurrency)
|
||||||
|
queue: asyncio.Queue[WorkItem | None] = asyncio.Queue(maxsize=self._pool_size * 4)
|
||||||
|
outcomes: list[WorkOutcome] = []
|
||||||
|
|
||||||
|
async def produce() -> None:
|
||||||
|
try:
|
||||||
|
async for item in item_source:
|
||||||
|
await queue.put(item)
|
||||||
|
finally:
|
||||||
|
for _ in range(self._pool_size):
|
||||||
|
await queue.put(None)
|
||||||
|
|
||||||
|
async def consume() -> None:
|
||||||
|
while True:
|
||||||
|
item = await queue.get()
|
||||||
|
if item is None:
|
||||||
|
return
|
||||||
|
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
|
||||||
|
outcomes.append(outcome)
|
||||||
|
|
||||||
|
producer_task = asyncio.create_task(produce())
|
||||||
|
worker_tasks = [asyncio.create_task(consume()) for _ in range(self._pool_size)]
|
||||||
|
try:
|
||||||
|
await producer_task
|
||||||
|
except Exception as exc:
|
||||||
|
logger.error("item source failed error=%s", exc)
|
||||||
|
await asyncio.gather(*worker_tasks)
|
||||||
|
report = self._build_report(outcomes)
|
||||||
|
logger.info(
|
||||||
|
"pipeline finished requests_succeeded=%d requests_failed=%d urls_found=%d queries_seeded=%d contents_seen=%d",
|
||||||
|
report.requests_succeeded,
|
||||||
|
report.requests_failed,
|
||||||
|
report.urls_found,
|
||||||
|
report.queries_seeded,
|
||||||
|
report.contents_seen,
|
||||||
|
)
|
||||||
|
return report
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _build_report(outcomes: list[WorkOutcome]) -> PipelineReport:
|
||||||
|
report = PipelineReport(outcomes=outcomes)
|
||||||
|
for outcome in outcomes:
|
||||||
|
if outcome.success:
|
||||||
|
report.requests_succeeded += 1
|
||||||
|
else:
|
||||||
|
report.requests_failed += 1
|
||||||
|
report.urls_found += outcome.urls_found
|
||||||
|
report.queries_seeded += outcome.queries_seeded
|
||||||
|
report.contents_seen += outcome.contents_seen
|
||||||
|
return report
|
||||||
|
|
||||||
365
tests/test_research_scheduling.py
Normal file
365
tests/test_research_scheduling.py
Normal file
@ -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()
|
||||||
|
|
||||||
Loading…
Reference in New Issue
Block a user