Compare commits
13 Commits
main
...
typosaurus
| Author | SHA1 | Date | |
|---|---|---|---|
| 84ac3029fc | |||
| dab3bb2437 | |||
| 472a59e122 | |||
| 808d6b4f83 | |||
| 32a17e3f5c | |||
| 14a60ef77f | |||
| bc11dc18b2 | |||
| 3ff5fc686a | |||
| c0a59b4138 | |||
| 7977217013 | |||
| be31a2774f | |||
| 31f2c6451f | |||
| e00a2db81b |
@ -3,6 +3,8 @@ name: CI
|
|||||||
on:
|
on:
|
||||||
push:
|
push:
|
||||||
branches: [main, master]
|
branches: [main, master]
|
||||||
|
pull_request:
|
||||||
|
branches: [main]
|
||||||
jobs:
|
jobs:
|
||||||
test:
|
test:
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
@ -19,3 +21,4 @@ jobs:
|
|||||||
- name: Run tests
|
- name: Run tests
|
||||||
run: make verify
|
run: make verify
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
140
deepresearch.md
Normal file
140
deepresearch.md
Normal file
@ -0,0 +1,140 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
# Deep Research Engine — Design, Optimality Argument and Verification Evidence
|
||||||
|
|
||||||
|
This document describes the exhaustive deep research engine in `src/typosaurus_sandbox/research/`,
|
||||||
|
the mathematical argument that recursive query expansion with URL/content deduplication and
|
||||||
|
closure detection is the most aggressive feasible research strategy over the rsearch API, and the
|
||||||
|
four recursive verification passes executed against it. Every claim is traceable to the run's
|
||||||
|
verified nodes (fact sheet node d5d9e290; optimality node b042b1d23; tester nodes f7f10c64,
|
||||||
|
fde105db, 2e6bd38b; engine node 1b0176bf) and to source path:line references.
|
||||||
|
|
||||||
|
## 1. Scope and constraints
|
||||||
|
|
||||||
|
- Only search API: `https://rsearch.app.molodetz.nl`; the client issues requests only to the
|
||||||
|
`/search`, `/chat` and `/describe` endpoints (client.py:211). `/search` is GET-only.
|
||||||
|
- Content-type agnostic: web results, image results (`type=images`), describe and chat flow
|
||||||
|
through one asynchronous pipeline; no per-type special casing beyond parameter selection.
|
||||||
|
- Native Python 3.12, standard library only (`asyncio`, `urllib`); no new dependency was added.
|
||||||
|
- No artificial depth cap, page cap or time budget stops a run before closure; the engine stops
|
||||||
|
only when a full round adds zero new URLs and zero new queries (least fixed point).
|
||||||
|
|
||||||
|
## 2. Architecture (module map)
|
||||||
|
|
||||||
|
| Module | Public symbol | Path:line |
|
||||||
|
|---|---|---|
|
||||||
|
| config | `ResearchConfig` (base_url, TTLs, `max_concurrency=8`, default_count) | `src/typosaurus_sandbox/research/config.py:12` |
|
||||||
|
| client | `RsearchClient`, `RsearchError` (search/chat/describe, `_request`) | `src/typosaurus_sandbox/research/client.py:72` |
|
||||||
|
| cache | `TTLCache`, `CacheEntry` (thread-safe, monotonic expiry) | `src/typosaurus_sandbox/research/cache.py:20` |
|
||||||
|
| envelopes | `SearchResponse`, `SearchResult`, `DeepReport`, `ChatResponse`, `DescribeResponse` | `src/typosaurus_sandbox/research/envelopes.py:103` |
|
||||||
|
| frontier | `QueryFrontier`, `DedupStats`, URL normalization, content fingerprint | `src/typosaurus_sandbox/research/frontier.py:102` |
|
||||||
|
| pipeline | `ResearchPipeline`, `WorkItem`, `PipelineReport` (bounded worker pool) | `src/typosaurus_sandbox/research/pipeline.py:125` |
|
||||||
|
| engine | `ResearchEngine`, `ResearchReport`, `RoundSummary` (closure loop) | `src/typosaurus_sandbox/research/engine.py:85` |
|
||||||
|
| entry | `main()` CLI | `src/typosaurus_sandbox/research/__main__.py:24` |
|
||||||
|
|
||||||
|
## 3. Concurrency model
|
||||||
|
|
||||||
|
- Bounded asyncio worker pool: `asyncio.Semaphore(pool_size)` with
|
||||||
|
`pool_size = max(1, max_concurrency)` and `max_concurrency = 8`
|
||||||
|
(config.py:18, pipeline.py:126-136).
|
||||||
|
- `run()` drains the frontier through a bounded queue (pool * 4) with pool-size workers and
|
||||||
|
`None` sentinels; every request runs via `asyncio.to_thread` over `urllib` (no extra deps).
|
||||||
|
- Pool size is logged at INFO; every request outcome (endpoint, query/url, status, cache hit)
|
||||||
|
at INFO, every extraction at DEBUG.
|
||||||
|
|
||||||
|
## 4. Deduplication and closure strategy
|
||||||
|
|
||||||
|
- Query dedup key: whitespace-collapsed `casefold` (frontier.py:28); length window 2-200 chars.
|
||||||
|
- URL dedup: `normalize_url` lowercases scheme/host, applies IDNA, strips default port,
|
||||||
|
userinfo and fragment, collapses slashes (frontier.py:28).
|
||||||
|
- Content dedup: SHA-256 fingerprint of whitespace-normalized text (frontier.py:61).
|
||||||
|
- One `threading.Lock` guards all seen-sets and counters for concurrent worker access
|
||||||
|
(frontier.py:103).
|
||||||
|
- Closure rule: a round that adds 0 new URLs and 0 new queries halts the run
|
||||||
|
(engine.py:178-183). The engine is closed-loop verified: a fixed-fixture fake client closed
|
||||||
|
in 3 rounds with all four content types, and a 4-level chain client closed in 5 rounds,
|
||||||
|
proving no depth cap (engine node 1b0176bf).
|
||||||
|
|
||||||
|
## 5. Content-type agnosticism
|
||||||
|
|
||||||
|
- One worker path serves all kinds: `web` -> `search(content=True)`, `images` ->
|
||||||
|
`search(type="images")`, `describe` -> GET `/describe?url=`, `chat` -> POST `/chat`
|
||||||
|
(pipeline.py:138-143, engine.py:106).
|
||||||
|
- Extraction yields new URLs and new query seeds from titles, descriptions and `extra` fields
|
||||||
|
of every content type (frontier.py:66).
|
||||||
|
|
||||||
|
## 6. Optimality argument
|
||||||
|
|
||||||
|
Let `R(q)` be the set of result URLs returned by the aggregator for query `q`, `gen(u)` the
|
||||||
|
query variants generated from URL/content `u`, and `S` the set of collected URLs.
|
||||||
|
|
||||||
|
- Completeness: the process is coverage-complete for subject `t` iff it halts at the least
|
||||||
|
fixed point `S* = lfp(F)` with `F(S) = S ∪ ⋃_{u∈S, q∈gen(u)} R(q)`; the halt condition is
|
||||||
|
"a full round adds 0 new URLs and 0 new queries" (node b042b1d23).
|
||||||
|
- Dominance: depth-`d` iteration reaches `F^d(S0) ⊆ S*`; the inclusion is strict whenever the
|
||||||
|
discovery chain exceeds `d`, so every fixed-depth strategy is incomplete. Closure iterates
|
||||||
|
`F` to its unique least fixed point (Knaster-Tarski), attaining the maximum reachable
|
||||||
|
coverage; any strategy that stops before the fixed point is strictly dominated.
|
||||||
|
- Cost model: `Cost = Σ_{q∈Q_issued} c(q) + Σ_{u∈F_issued} c_c(u)`. Search (5 min) and content
|
||||||
|
(24 h) caches (config.py:16-17) make repeat queries near-free; the dominant cost is
|
||||||
|
`|Q_issued| + |F_issued|`, and query/URL dedup touches each element exactly once.
|
||||||
|
- Stated assumptions and limits: single aggregator (rsearch only), no pagination API,
|
||||||
|
documented count bound 1-100 with the provider capping at 10, and content retrieval only
|
||||||
|
through the aggregator. Optimality is proven within these constraints.
|
||||||
|
- Dated references (tier): rsearch docs https://rsearch.app.molodetz.nl/about (2026-08-07, 1);
|
||||||
|
Gemini https://blog.google/products-and-platforms/products/gemini/google-gemini-deep-research/
|
||||||
|
(2024-12-11, 1); OpenAI https://openai.com/index/introducing-deep-research/ (Feb-2025, 1) +
|
||||||
|
https://techcrunch.com/2025/02/02/openai-unveils-a-new-chatgpt-agent-for-deep-research/ (4);
|
||||||
|
Ntoulas 2005 ACM JCDL 10.1145/1065385.1065407 (3); Chakrabarti 1999 Computer Networks
|
||||||
|
10.1016/S1389-1286(99)00052-3 (3); Olston & Najork 2010 FnTIR 10.1561/1500000017 (3).
|
||||||
|
|
||||||
|
## 7. Four recursive verification passes
|
||||||
|
|
||||||
|
Each pass re-checks the previous pass's optimality claim ("recursive closure over the rsearch
|
||||||
|
aggregator is the most aggressive feasible strategy") and records its own evidence. All four
|
||||||
|
passes passed.
|
||||||
|
|
||||||
|
- Pass 1 — Optimality argument: formal completeness criterion, cost model and Knaster-Tarski
|
||||||
|
dominance proof produced with seven dated, tiered sources (node b042b1d23, 2026-08-07).
|
||||||
|
- Pass 2 — Engine matches the argument: all eight engine acceptance criteria executed with
|
||||||
|
pass verdicts and exact commands (node f7f10c64): rsearch-only source, bounded pool at
|
||||||
|
max_concurrency=8, one web/images/describe/chat pipeline, URL+content dedup (64 concurrent
|
||||||
|
same-query pushes -> 1 enqueued, 63 skipped), closure decision (NullClient probe closed in 1
|
||||||
|
round with 0 new URLs and 0 new queries), logging/annotations, no deferred markers, and
|
||||||
|
`make verify` -> "Ran 199 tests in 2.168s OK verification passed".
|
||||||
|
- Pass 3 — Live probe coverage/cost (node fde105db, 2026-08-07): subject "python asyncio",
|
||||||
|
max_concurrency=8, count=10, 240 s guard: queries_issued=86, urls_seen=754, contents_seen=281,
|
||||||
|
164 network requests (search 105 / chat 46 / describe 13), X-AI-Cost-USD sum $0.002075, wall
|
||||||
|
elapsed 264.91 s. Adversarial subjects ("", spaces, tabs) raised ValueError
|
||||||
|
"research subject must not be empty" (engine.py:125) before any API call; urlopen delta 0.
|
||||||
|
- Pass 4 — Closure and determinism (node 1b0176bf, confirmed by fact sheet d5d9e290):
|
||||||
|
fixed-fixture fake client closed in 3 rounds with all 4 content types; 4-level chain closed
|
||||||
|
in 5 rounds (no depth cap); live `python -m typosaurus_sandbox.research` logged INFO rounds,
|
||||||
|
closure and the typed report JSON; final gate `make verify` green (199 tests OK, git clean).
|
||||||
|
|
||||||
|
## 8. Usage
|
||||||
|
|
||||||
|
```sh
|
||||||
|
python -m typosaurus_sandbox.research [subject]
|
||||||
|
```
|
||||||
|
|
||||||
|
Run a research session on `subject` (default "typosaurus sandbox") until closure; rounds and
|
||||||
|
closure decisions are logged at INFO, and the typed `ResearchReport` JSON is logged at the end.
|
||||||
|
Configuration (base_url, TTLs, max_concurrency, default_count) is loaded from `.env.json` under
|
||||||
|
the `research` key with plug-and-play defaults (config.py:22).
|
||||||
|
|
||||||
|
Verification gate:
|
||||||
|
|
||||||
|
```sh
|
||||||
|
make verify
|
||||||
|
```
|
||||||
|
|
||||||
|
## 9. Verification status
|
||||||
|
|
||||||
|
- `make verify`: exit 0, "Ran 199 tests OK verification passed" (2026-08-07); only the
|
||||||
|
pre-existing Starlette deprecation warning from the FastAPI test client remains, none
|
||||||
|
introduced by the research package.
|
||||||
|
- Re-run at document time: `make verify` exit 0, "Ran 217 tests in 2.162s OK", verification
|
||||||
|
passed; same pre-existing Starlette deprecation warning only.
|
||||||
|
- The research package contains no deferred markers (grep verified, node f7f10c64).
|
||||||
|
|
||||||
|
|
||||||
68
src/typosaurus_sandbox/research/__init__.py
Normal file
68
src/typosaurus_sandbox/research/__init__.py
Normal file
@ -0,0 +1,68 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
from typosaurus_sandbox.research.cache import TTLCache
|
||||||
|
from typosaurus_sandbox.research.client import RsearchClient, RsearchError
|
||||||
|
from typosaurus_sandbox.research.config import ResearchConfig
|
||||||
|
from typosaurus_sandbox.research.engine import ResearchEngine, ResearchReport, RoundSummary
|
||||||
|
from typosaurus_sandbox.research.envelopes import (
|
||||||
|
ChatResponse,
|
||||||
|
ChatUsage,
|
||||||
|
DeepReport,
|
||||||
|
DescribeResponse,
|
||||||
|
SearchGrade,
|
||||||
|
SearchResponse,
|
||||||
|
SearchResult,
|
||||||
|
)
|
||||||
|
from typosaurus_sandbox.research.frontier import (
|
||||||
|
DedupStats,
|
||||||
|
QueryFrontier,
|
||||||
|
fingerprint_text,
|
||||||
|
normalize_url,
|
||||||
|
query_variants_from_result,
|
||||||
|
)
|
||||||
|
from typosaurus_sandbox.research.pipeline import (
|
||||||
|
ContentKind,
|
||||||
|
Extraction,
|
||||||
|
PipelineReport,
|
||||||
|
ResearchPipeline,
|
||||||
|
WorkItem,
|
||||||
|
WorkOutcome,
|
||||||
|
apply_extraction,
|
||||||
|
extract_response,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"ChatResponse",
|
||||||
|
"ChatUsage",
|
||||||
|
"ContentKind",
|
||||||
|
"DedupStats",
|
||||||
|
"DeepReport",
|
||||||
|
"DescribeResponse",
|
||||||
|
"Extraction",
|
||||||
|
"PipelineReport",
|
||||||
|
"QueryFrontier",
|
||||||
|
"ResearchEngine",
|
||||||
|
"ResearchPipeline",
|
||||||
|
"ResearchReport",
|
||||||
|
"RoundSummary",
|
||||||
|
"RsearchClient",
|
||||||
|
"RsearchError",
|
||||||
|
"ResearchConfig",
|
||||||
|
"SearchGrade",
|
||||||
|
"SearchResponse",
|
||||||
|
"SearchResult",
|
||||||
|
"TTLCache",
|
||||||
|
"WorkItem",
|
||||||
|
"WorkOutcome",
|
||||||
|
"apply_extraction",
|
||||||
|
"extract_response",
|
||||||
|
"fingerprint_text",
|
||||||
|
"normalize_url",
|
||||||
|
"query_variants_from_result",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
42
src/typosaurus_sandbox/research/__main__.py
Normal file
42
src/typosaurus_sandbox/research/__main__.py
Normal file
@ -0,0 +1,42 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from typosaurus_sandbox.core.logging import setup_logging
|
||||||
|
from typosaurus_sandbox.research.client import RsearchClient
|
||||||
|
from typosaurus_sandbox.research.config import ResearchConfig
|
||||||
|
from typosaurus_sandbox.research.engine import ResearchEngine
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _enable_console_logging() -> None:
|
||||||
|
root = logging.getLogger()
|
||||||
|
console = logging.StreamHandler()
|
||||||
|
console.setLevel(logging.INFO)
|
||||||
|
console.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s"))
|
||||||
|
root.addHandler(console)
|
||||||
|
|
||||||
|
|
||||||
|
def main(argv: list[str] | None = None) -> None:
|
||||||
|
setup_logging()
|
||||||
|
_enable_console_logging()
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
prog="typosaurus-sandbox-research",
|
||||||
|
description="Exhaustive deep research over the rsearch API until closure",
|
||||||
|
)
|
||||||
|
parser.add_argument("subject", nargs="?", default="typosaurus sandbox", help="subject to research until closure")
|
||||||
|
args = parser.parse_args(argv)
|
||||||
|
config = ResearchConfig.load()
|
||||||
|
client = RsearchClient(config)
|
||||||
|
logger.info("research session starting subject=%r base_url=%s", args.subject, config.base_url)
|
||||||
|
report = asyncio.run(ResearchEngine(client=client).run(args.subject))
|
||||||
|
logger.info("research report %s", json.dumps(report.to_dict(), indent=2))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
|
||||||
50
src/typosaurus_sandbox/research/cache.py
Normal file
50
src/typosaurus_sandbox/research/cache.py
Normal file
@ -0,0 +1,50 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Generic, TypeVar
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
T = TypeVar("T")
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class CacheEntry(Generic[T]):
|
||||||
|
value: T
|
||||||
|
expires_at: float
|
||||||
|
|
||||||
|
|
||||||
|
class TTLCache(Generic[T]):
|
||||||
|
def __init__(self, name: str, ttl_seconds: float) -> None:
|
||||||
|
self._name = name
|
||||||
|
self._ttl_seconds = ttl_seconds
|
||||||
|
self._entries: dict[str, CacheEntry[T]] = {}
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
|
||||||
|
def get(self, key: str) -> T | None:
|
||||||
|
with self._lock:
|
||||||
|
entry = self._entries.get(key)
|
||||||
|
if entry is None:
|
||||||
|
logger.debug("cache %s miss key=%s", self._name, key)
|
||||||
|
return None
|
||||||
|
if time.monotonic() >= entry.expires_at:
|
||||||
|
del self._entries[key]
|
||||||
|
logger.debug("cache %s expired key=%s", self._name, key)
|
||||||
|
return None
|
||||||
|
logger.debug("cache %s hit key=%s", self._name, key)
|
||||||
|
return entry.value
|
||||||
|
|
||||||
|
def set(self, key: str, value: T) -> None:
|
||||||
|
with self._lock:
|
||||||
|
self._entries[key] = CacheEntry(value=value, expires_at=time.monotonic() + self._ttl_seconds)
|
||||||
|
logger.debug("cache %s set key=%s ttl=%.0fs", self._name, key, self._ttl_seconds)
|
||||||
|
|
||||||
|
def clear(self) -> None:
|
||||||
|
with self._lock:
|
||||||
|
count = len(self._entries)
|
||||||
|
self._entries.clear()
|
||||||
|
logger.debug("cache %s cleared %d entries", self._name, count)
|
||||||
|
|
||||||
251
src/typosaurus_sandbox/research/client.py
Normal file
251
src/typosaurus_sandbox/research/client.py
Normal file
@ -0,0 +1,251 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import hashlib
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import secrets
|
||||||
|
import urllib.error
|
||||||
|
import urllib.parse
|
||||||
|
import urllib.request
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from typosaurus_sandbox.research.cache import TTLCache
|
||||||
|
from typosaurus_sandbox.research.config import ResearchConfig
|
||||||
|
from typosaurus_sandbox.research.envelopes import ChatResponse, DescribeResponse, SearchResponse
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
MAX_ERROR_LENGTH = 200
|
||||||
|
|
||||||
|
|
||||||
|
class RsearchError(RuntimeError):
|
||||||
|
def __init__(self, message: str, status_code: int | None = None) -> None:
|
||||||
|
super().__init__(message)
|
||||||
|
self.status_code = status_code
|
||||||
|
|
||||||
|
|
||||||
|
def _multipart_body(field_name: str, filename: str, mime_type: str, payload: bytes) -> tuple[bytes, str]:
|
||||||
|
boundary = "----rsearch-" + secrets.token_hex(8)
|
||||||
|
head = (
|
||||||
|
f"--{boundary}\r\n".encode()
|
||||||
|
+ f'Content-Disposition: form-data; name="{field_name}"; filename="{filename}"\r\n'.encode()
|
||||||
|
+ f"Content-Type: {mime_type}\r\n\r\n".encode()
|
||||||
|
)
|
||||||
|
tail = b"\r\n--" + boundary.encode() + b"--\r\n"
|
||||||
|
return head + payload + tail, f"multipart/form-data; boundary={boundary}"
|
||||||
|
|
||||||
|
|
||||||
|
def _content_hash(image_bytes: bytes) -> str:
|
||||||
|
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:
|
||||||
|
def __init__(self, config: ResearchConfig | None = None) -> None:
|
||||||
|
self._config = config if config is not None else ResearchConfig()
|
||||||
|
self._search_cache = TTLCache[SearchResponse]("search", self._config.search_cache_ttl_seconds)
|
||||||
|
self._content_cache = TTLCache[str]("content", self._config.content_cache_ttl_seconds)
|
||||||
|
self._describe_cache = TTLCache[DescribeResponse]("describe", self._config.content_cache_ttl_seconds)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def config(self) -> ResearchConfig:
|
||||||
|
return self._config
|
||||||
|
|
||||||
|
def get_cached_content(self, url: str) -> str | None:
|
||||||
|
return self._content_cache.get(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:
|
||||||
|
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()))
|
||||||
|
if cache:
|
||||||
|
cached_response = self._search_cache.get(key)
|
||||||
|
if cached_response is not None:
|
||||||
|
return cached_response
|
||||||
|
timeout = self._config.deep_timeout_seconds if deep else self._config.request_timeout_seconds
|
||||||
|
status, data = await asyncio.to_thread(self._request, "GET", "/search", params, None, None, timeout)
|
||||||
|
response = SearchResponse.from_dict(data)
|
||||||
|
if cache:
|
||||||
|
self._search_cache.set(key, response)
|
||||||
|
if content:
|
||||||
|
for result in response.results:
|
||||||
|
if result.content:
|
||||||
|
self._content_cache.set(result.url, result.content)
|
||||||
|
logger.info(
|
||||||
|
"search query=%r source=%s count=%s deep=%s ai=%s results=%d",
|
||||||
|
query,
|
||||||
|
response.source,
|
||||||
|
response.count,
|
||||||
|
deep,
|
||||||
|
ai,
|
||||||
|
len(response.results),
|
||||||
|
)
|
||||||
|
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(
|
||||||
|
self,
|
||||||
|
prompt: str,
|
||||||
|
*,
|
||||||
|
system: str | None = None,
|
||||||
|
json_mode: bool = False,
|
||||||
|
cache: bool = True,
|
||||||
|
) -> ChatResponse:
|
||||||
|
payload: dict[str, Any] = {"prompt": prompt}
|
||||||
|
if system is not None:
|
||||||
|
payload["system"] = system
|
||||||
|
if json_mode:
|
||||||
|
payload["json"] = True
|
||||||
|
if not cache:
|
||||||
|
payload["cache"] = False
|
||||||
|
body = json.dumps(payload).encode()
|
||||||
|
headers = {"Content-Type": "application/json"}
|
||||||
|
status, data = await asyncio.to_thread(self._request, "POST", "/chat", None, body, headers, None)
|
||||||
|
response = ChatResponse.from_dict(data)
|
||||||
|
logger.info("chat prompt=%r cached=%s", prompt, response.cached)
|
||||||
|
return response
|
||||||
|
|
||||||
|
async def describe(self, url: str) -> DescribeResponse:
|
||||||
|
key = f"url:{url}"
|
||||||
|
cached = self._describe_cache.get(key)
|
||||||
|
if cached is not None:
|
||||||
|
return cached
|
||||||
|
status, data = await asyncio.to_thread(self._request, "GET", "/describe", {"url": url}, None, None, None)
|
||||||
|
response = DescribeResponse.from_dict(data)
|
||||||
|
self._describe_cache.set(key, response)
|
||||||
|
logger.info("describe url=%s", url)
|
||||||
|
return response
|
||||||
|
|
||||||
|
async def describe_upload(self, image_bytes: bytes, *, filename: str, mime_type: str) -> DescribeResponse:
|
||||||
|
body, content_type = _multipart_body("file", filename, mime_type, image_bytes)
|
||||||
|
headers = {"Content-Type": content_type}
|
||||||
|
return await self._describe_post(image_bytes, body, headers)
|
||||||
|
|
||||||
|
async def describe_raw(self, image_bytes: bytes, *, mime_type: str) -> DescribeResponse:
|
||||||
|
headers = {"Content-Type": mime_type}
|
||||||
|
return await self._describe_post(image_bytes, image_bytes, headers)
|
||||||
|
|
||||||
|
async def _describe_post(self, image_bytes: bytes, body: bytes, headers: dict[str, str]) -> DescribeResponse:
|
||||||
|
key = "hash:" + _content_hash(image_bytes)
|
||||||
|
cached = self._describe_cache.get(key)
|
||||||
|
if cached is not None:
|
||||||
|
return cached
|
||||||
|
status, data = await asyncio.to_thread(self._request, "POST", "/describe", None, body, headers, None)
|
||||||
|
response = DescribeResponse.from_dict(data)
|
||||||
|
self._describe_cache.set(key, response)
|
||||||
|
logger.info("describe post size=%d", len(image_bytes))
|
||||||
|
return response
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _error_message(data: dict[str, Any]) -> str:
|
||||||
|
error = data.get("error")
|
||||||
|
if isinstance(error, str) and error:
|
||||||
|
return error
|
||||||
|
detail = data.get("detail")
|
||||||
|
if isinstance(detail, str) and detail:
|
||||||
|
return detail
|
||||||
|
title = data.get("title")
|
||||||
|
if isinstance(title, str) and title:
|
||||||
|
return title
|
||||||
|
return json.dumps(data)[:MAX_ERROR_LENGTH]
|
||||||
|
|
||||||
|
def _request(
|
||||||
|
self,
|
||||||
|
method: str,
|
||||||
|
path: str,
|
||||||
|
params: dict[str, str] | None = None,
|
||||||
|
payload: bytes | None = None,
|
||||||
|
headers: dict[str, str] | None = None,
|
||||||
|
timeout: float | None = None,
|
||||||
|
) -> tuple[int, dict[str, Any]]:
|
||||||
|
timeout_seconds = timeout if timeout is not None else self._config.request_timeout_seconds
|
||||||
|
base_url = self._config.base_url
|
||||||
|
if base_url.endswith("/"):
|
||||||
|
base_url = base_url[:-1]
|
||||||
|
url = base_url + path
|
||||||
|
if params:
|
||||||
|
url = url + "?" + urllib.parse.urlencode(params)
|
||||||
|
request = urllib.request.Request(url, data=payload, method=method, headers=headers or {})
|
||||||
|
try:
|
||||||
|
with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
|
||||||
|
status = response.status
|
||||||
|
body = response.read()
|
||||||
|
except urllib.error.HTTPError as exc:
|
||||||
|
status = exc.code
|
||||||
|
body = exc.read()
|
||||||
|
except urllib.error.URLError as exc:
|
||||||
|
raise RsearchError(f"connection failure for {method} {path}: {exc.reason}") from exc
|
||||||
|
if not body:
|
||||||
|
raise RsearchError(f"empty response for {method} {path}", status)
|
||||||
|
try:
|
||||||
|
data = json.loads(body)
|
||||||
|
except (json.JSONDecodeError, UnicodeDecodeError) as exc:
|
||||||
|
raise RsearchError(f"invalid JSON for {method} {path}: {exc}", status) from exc
|
||||||
|
if not isinstance(data, dict):
|
||||||
|
raise RsearchError(f"unexpected response shape for {method} {path}", status)
|
||||||
|
if status >= 400 or data.get("success") is False:
|
||||||
|
raise RsearchError(self._error_message(data), status)
|
||||||
|
return status, data
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
40
src/typosaurus_sandbox/research/config.py
Normal file
40
src/typosaurus_sandbox/research/config.py
Normal file
@ -0,0 +1,40 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ResearchConfig:
|
||||||
|
base_url: str = "https://rsearch.app.molodetz.nl"
|
||||||
|
request_timeout_seconds: float = 30.0
|
||||||
|
deep_timeout_seconds: float = 180.0
|
||||||
|
search_cache_ttl_seconds: float = 300.0
|
||||||
|
content_cache_ttl_seconds: float = 86400.0
|
||||||
|
max_concurrency: int = 8
|
||||||
|
default_count: int = 10
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def load(cls) -> "ResearchConfig":
|
||||||
|
config_path = Path(".env.json")
|
||||||
|
if not config_path.exists():
|
||||||
|
logger.info("no .env.json found, using default research config")
|
||||||
|
return cls()
|
||||||
|
with config_path.open() as f:
|
||||||
|
data = json.load(f)
|
||||||
|
research = data.get("research", {})
|
||||||
|
logger.info("loaded research config from .env.json")
|
||||||
|
return cls(
|
||||||
|
base_url=research.get("base_url", cls.base_url),
|
||||||
|
request_timeout_seconds=research.get("request_timeout_seconds", cls.request_timeout_seconds),
|
||||||
|
deep_timeout_seconds=research.get("deep_timeout_seconds", cls.deep_timeout_seconds),
|
||||||
|
search_cache_ttl_seconds=research.get("search_cache_ttl_seconds", cls.search_cache_ttl_seconds),
|
||||||
|
content_cache_ttl_seconds=research.get("content_cache_ttl_seconds", cls.content_cache_ttl_seconds),
|
||||||
|
max_concurrency=research.get("max_concurrency", cls.max_concurrency),
|
||||||
|
default_count=research.get("default_count", cls.default_count),
|
||||||
|
)
|
||||||
|
|
||||||
210
src/typosaurus_sandbox/research/engine.py
Normal file
210
src/typosaurus_sandbox/research/engine.py
Normal file
@ -0,0 +1,210 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import AsyncIterator
|
||||||
|
|
||||||
|
from typosaurus_sandbox.research.client import RsearchClient
|
||||||
|
from typosaurus_sandbox.research.config import ResearchConfig
|
||||||
|
from typosaurus_sandbox.research.frontier import QueryFrontier
|
||||||
|
from typosaurus_sandbox.research.pipeline import PipelineReport, ResearchPipeline, WorkItem
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class RoundSummary:
|
||||||
|
number: int = 0
|
||||||
|
items_processed: int = 0
|
||||||
|
requests_succeeded: int = 0
|
||||||
|
requests_failed: int = 0
|
||||||
|
cache_hits: int = 0
|
||||||
|
content_types: dict[str, int] = field(default_factory=dict)
|
||||||
|
new_urls: int = 0
|
||||||
|
new_queries: int = 0
|
||||||
|
new_contents: int = 0
|
||||||
|
closed: bool = False
|
||||||
|
|
||||||
|
def to_dict(self) -> dict[str, int | dict[str, int] | bool]:
|
||||||
|
return {
|
||||||
|
"number": self.number,
|
||||||
|
"items_processed": self.items_processed,
|
||||||
|
"requests_succeeded": self.requests_succeeded,
|
||||||
|
"requests_failed": self.requests_failed,
|
||||||
|
"cache_hits": self.cache_hits,
|
||||||
|
"content_types": self.content_types,
|
||||||
|
"new_urls": self.new_urls,
|
||||||
|
"new_queries": self.new_queries,
|
||||||
|
"new_contents": self.new_contents,
|
||||||
|
"closed": self.closed,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ResearchReport:
|
||||||
|
subject: str
|
||||||
|
rounds: list[RoundSummary] = field(default_factory=list)
|
||||||
|
total_rounds: int = 0
|
||||||
|
queries_generated: int = 0
|
||||||
|
queries_enqueued: int = 0
|
||||||
|
queries_issued: int = 0
|
||||||
|
queries_duplicates_skipped: int = 0
|
||||||
|
urls_collected: int = 0
|
||||||
|
urls_duplicates_skipped: int = 0
|
||||||
|
contents_seen: int = 0
|
||||||
|
content_duplicates_skipped: int = 0
|
||||||
|
content_types: dict[str, int] = field(default_factory=dict)
|
||||||
|
requests_succeeded: int = 0
|
||||||
|
requests_failed: int = 0
|
||||||
|
cache_hits: int = 0
|
||||||
|
cache_misses: int = 0
|
||||||
|
closed: bool = False
|
||||||
|
|
||||||
|
def to_dict(self) -> dict[str, object]:
|
||||||
|
return {
|
||||||
|
"subject": self.subject,
|
||||||
|
"rounds": [round_summary.to_dict() for round_summary in self.rounds],
|
||||||
|
"total_rounds": self.total_rounds,
|
||||||
|
"queries_generated": self.queries_generated,
|
||||||
|
"queries_enqueued": self.queries_enqueued,
|
||||||
|
"queries_issued": self.queries_issued,
|
||||||
|
"queries_duplicates_skipped": self.queries_duplicates_skipped,
|
||||||
|
"urls_collected": self.urls_collected,
|
||||||
|
"urls_duplicates_skipped": self.urls_duplicates_skipped,
|
||||||
|
"contents_seen": self.contents_seen,
|
||||||
|
"content_duplicates_skipped": self.content_duplicates_skipped,
|
||||||
|
"content_types": self.content_types,
|
||||||
|
"requests_succeeded": self.requests_succeeded,
|
||||||
|
"requests_failed": self.requests_failed,
|
||||||
|
"cache_hits": self.cache_hits,
|
||||||
|
"cache_misses": self.cache_misses,
|
||||||
|
"closed": self.closed,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class ResearchEngine:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
client: RsearchClient | None = None,
|
||||||
|
frontier: QueryFrontier | None = None,
|
||||||
|
pipeline: ResearchPipeline | None = None,
|
||||||
|
) -> None:
|
||||||
|
self._client = client if client is not None else RsearchClient()
|
||||||
|
self._config: ResearchConfig = self._client.config
|
||||||
|
self._frontier = frontier if frontier is not None else QueryFrontier()
|
||||||
|
self._pipeline = pipeline if pipeline is not None else ResearchPipeline(self._client, self._frontier)
|
||||||
|
self._described_marker = 0
|
||||||
|
|
||||||
|
@property
|
||||||
|
def frontier(self) -> QueryFrontier:
|
||||||
|
return self._frontier
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pipeline(self) -> ResearchPipeline:
|
||||||
|
return self._pipeline
|
||||||
|
|
||||||
|
async def _round_items(self, pending_queries: int, urls_to_describe: list[str]) -> AsyncIterator[WorkItem]:
|
||||||
|
for _ in range(pending_queries):
|
||||||
|
query = self._frontier.pop_query()
|
||||||
|
if query is None:
|
||||||
|
break
|
||||||
|
yield WorkItem("web", query, deep=True, ai=True)
|
||||||
|
yield WorkItem("images", query)
|
||||||
|
yield WorkItem("chat", query)
|
||||||
|
for url in urls_to_describe:
|
||||||
|
yield WorkItem("describe", url)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _round_summary(number: int, pipeline_report: PipelineReport) -> RoundSummary:
|
||||||
|
summary = RoundSummary(number=number, items_processed=len(pipeline_report.outcomes))
|
||||||
|
for outcome in pipeline_report.outcomes:
|
||||||
|
if outcome.success:
|
||||||
|
summary.requests_succeeded += 1
|
||||||
|
else:
|
||||||
|
summary.requests_failed += 1
|
||||||
|
if outcome.cache_hit:
|
||||||
|
summary.cache_hits += 1
|
||||||
|
kind = outcome.item.kind
|
||||||
|
summary.content_types[kind] = summary.content_types.get(kind, 0) + 1
|
||||||
|
return summary
|
||||||
|
|
||||||
|
async def run(self, subject: str) -> ResearchReport:
|
||||||
|
cleaned_subject = " ".join(subject.split())
|
||||||
|
if not cleaned_subject:
|
||||||
|
raise ValueError("research subject must not be empty")
|
||||||
|
self._frontier.seed(cleaned_subject)
|
||||||
|
report = ResearchReport(subject=cleaned_subject)
|
||||||
|
round_number = 0
|
||||||
|
while True:
|
||||||
|
round_start = self._frontier.snapshot()
|
||||||
|
pending_queries = round_start.queries_enqueued - round_start.queries_issued
|
||||||
|
urls_to_describe = self._frontier.urls_since(self._described_marker)
|
||||||
|
self._described_marker = round_start.urls_seen
|
||||||
|
if pending_queries == 0 and not urls_to_describe:
|
||||||
|
logger.info("research closed, no pending queries or urls after round %d", round_number)
|
||||||
|
break
|
||||||
|
round_number += 1
|
||||||
|
logger.info(
|
||||||
|
"round %d start pending_queries=%d urls_to_describe=%d",
|
||||||
|
round_number,
|
||||||
|
pending_queries,
|
||||||
|
len(urls_to_describe),
|
||||||
|
)
|
||||||
|
pipeline_report = await self._pipeline.run(self._round_items(pending_queries, urls_to_describe))
|
||||||
|
summary = self._round_summary(round_number, pipeline_report)
|
||||||
|
round_end = self._frontier.snapshot()
|
||||||
|
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
|
||||||
|
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",
|
||||||
|
round_number,
|
||||||
|
summary.new_urls,
|
||||||
|
summary.new_queries,
|
||||||
|
summary.new_contents,
|
||||||
|
summary.closed,
|
||||||
|
)
|
||||||
|
if summary.closed:
|
||||||
|
break
|
||||||
|
report.total_rounds = round_number
|
||||||
|
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 closed=%s",
|
||||||
|
report.subject,
|
||||||
|
report.total_rounds,
|
||||||
|
report.queries_issued,
|
||||||
|
report.urls_collected,
|
||||||
|
report.contents_seen,
|
||||||
|
report.cache_hits,
|
||||||
|
report.closed,
|
||||||
|
)
|
||||||
|
return report
|
||||||
|
|
||||||
|
def _finalize(self, report: ResearchReport) -> None:
|
||||||
|
stats = self._frontier.snapshot()
|
||||||
|
report.queries_generated = stats.queries_generated
|
||||||
|
report.queries_enqueued = stats.queries_enqueued
|
||||||
|
report.queries_issued = stats.queries_issued
|
||||||
|
report.queries_duplicates_skipped = stats.queries_duplicates_skipped
|
||||||
|
report.urls_collected = stats.urls_seen
|
||||||
|
report.urls_duplicates_skipped = stats.urls_duplicates_skipped
|
||||||
|
report.contents_seen = stats.content_seen
|
||||||
|
report.content_duplicates_skipped = stats.content_duplicates_skipped
|
||||||
|
total_items = 0
|
||||||
|
for summary in report.rounds:
|
||||||
|
total_items += summary.items_processed
|
||||||
|
report.requests_succeeded += summary.requests_succeeded
|
||||||
|
report.requests_failed += summary.requests_failed
|
||||||
|
report.cache_hits += summary.cache_hits
|
||||||
|
for kind, count in summary.content_types.items():
|
||||||
|
report.content_types[kind] = report.content_types.get(kind, 0) + count
|
||||||
|
report.cache_misses = total_items - report.cache_hits
|
||||||
|
|
||||||
|
|
||||||
203
src/typosaurus_sandbox/research/envelopes.py
Normal file
203
src/typosaurus_sandbox/research/envelopes.py
Normal file
@ -0,0 +1,203 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
|
||||||
|
def _as_float(value: Any) -> float | None:
|
||||||
|
if value is None:
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
return float(value)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class SearchGrade:
|
||||||
|
overall: float = 0.0
|
||||||
|
relevance: float = 0.0
|
||||||
|
depth: float = 0.0
|
||||||
|
authority: float = 0.0
|
||||||
|
freshness: float = 0.0
|
||||||
|
word_count: int = 0
|
||||||
|
intent_hits: int = 0
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, data: dict[str, Any] | None) -> "SearchGrade | None":
|
||||||
|
if data is None:
|
||||||
|
return None
|
||||||
|
return cls(
|
||||||
|
overall=float(data.get("overall", 0.0) or 0.0),
|
||||||
|
relevance=float(data.get("relevance", 0.0) or 0.0),
|
||||||
|
depth=float(data.get("depth", 0.0) or 0.0),
|
||||||
|
authority=float(data.get("authority", 0.0) or 0.0),
|
||||||
|
freshness=float(data.get("freshness", 0.0) or 0.0),
|
||||||
|
word_count=int(data.get("word_count", 0) or 0),
|
||||||
|
intent_hits=int(data.get("intent_hits", 0) or 0),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class SearchResult:
|
||||||
|
title: str = ""
|
||||||
|
url: str = ""
|
||||||
|
description: str = ""
|
||||||
|
source: str = ""
|
||||||
|
content: str | None = None
|
||||||
|
extra: dict[str, Any] = field(default_factory=dict)
|
||||||
|
index: int | None = None
|
||||||
|
grade: SearchGrade | None = None
|
||||||
|
query_origin: str | None = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, data: dict[str, Any]) -> "SearchResult":
|
||||||
|
return cls(
|
||||||
|
title=data.get("title", ""),
|
||||||
|
url=data.get("url", ""),
|
||||||
|
description=data.get("description", ""),
|
||||||
|
source=data.get("source", ""),
|
||||||
|
content=data.get("content"),
|
||||||
|
extra=data.get("extra", {}),
|
||||||
|
index=data.get("index"),
|
||||||
|
grade=SearchGrade.from_dict(data.get("grade")),
|
||||||
|
query_origin=data.get("query_origin"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DeepReport:
|
||||||
|
query: str = ""
|
||||||
|
markdown: str = ""
|
||||||
|
sources: list[SearchResult] = field(default_factory=list)
|
||||||
|
graded_count: int = 0
|
||||||
|
total_count: int = 0
|
||||||
|
model: str = ""
|
||||||
|
elapsed: float = 0.0
|
||||||
|
cache_hit: bool = False
|
||||||
|
rounds: int = 0
|
||||||
|
queries_tried: list[str] = field(default_factory=list)
|
||||||
|
error: str | None = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, data: dict[str, Any] | None) -> "DeepReport | None":
|
||||||
|
if data is None:
|
||||||
|
return None
|
||||||
|
sources = [SearchResult.from_dict(item) for item in data.get("sources", [])]
|
||||||
|
return cls(
|
||||||
|
query=data.get("query", ""),
|
||||||
|
markdown=data.get("markdown", ""),
|
||||||
|
sources=sources,
|
||||||
|
graded_count=int(data.get("graded_count", 0) or 0),
|
||||||
|
total_count=int(data.get("total_count", 0) or 0),
|
||||||
|
model=data.get("model", ""),
|
||||||
|
elapsed=_as_float(data.get("elapsed")) or 0.0,
|
||||||
|
cache_hit=bool(data.get("cache_hit", False)),
|
||||||
|
rounds=int(data.get("rounds", 0) or 0),
|
||||||
|
queries_tried=list(data.get("queries_tried", [])),
|
||||||
|
error=data.get("error"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class SearchResponse:
|
||||||
|
query: str = ""
|
||||||
|
source: str = ""
|
||||||
|
count: int = 0
|
||||||
|
results: list[SearchResult] = field(default_factory=list)
|
||||||
|
success: bool = False
|
||||||
|
error: str | None = None
|
||||||
|
ai_response: str | None = None
|
||||||
|
ai_error: str | None = None
|
||||||
|
deep: DeepReport | None = None
|
||||||
|
timestamp: str | None = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, data: dict[str, Any]) -> "SearchResponse":
|
||||||
|
results = [SearchResult.from_dict(item) for item in data.get("results", [])]
|
||||||
|
return cls(
|
||||||
|
query=data.get("query", ""),
|
||||||
|
source=data.get("source", ""),
|
||||||
|
count=int(data.get("count", 0) or 0),
|
||||||
|
results=results,
|
||||||
|
success=bool(data.get("success", False)),
|
||||||
|
error=data.get("error"),
|
||||||
|
ai_response=data.get("ai_response"),
|
||||||
|
ai_error=data.get("ai_error"),
|
||||||
|
deep=DeepReport.from_dict(data.get("deep")),
|
||||||
|
timestamp=data.get("timestamp"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ChatUsage:
|
||||||
|
prompt_tokens: int = 0
|
||||||
|
completion_tokens: int = 0
|
||||||
|
total_tokens: int = 0
|
||||||
|
cost_usd: float = 0.0
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, data: dict[str, Any] | None) -> "ChatUsage | None":
|
||||||
|
if data is None:
|
||||||
|
return None
|
||||||
|
return cls(
|
||||||
|
prompt_tokens=int(data.get("prompt_tokens", 0) or 0),
|
||||||
|
completion_tokens=int(data.get("completion_tokens", 0) or 0),
|
||||||
|
total_tokens=int(data.get("total_tokens", 0) or 0),
|
||||||
|
cost_usd=float(data.get("cost_usd", 0.0) or 0.0),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ChatResponse:
|
||||||
|
response: str = ""
|
||||||
|
prompt: str = ""
|
||||||
|
json_mode: bool = False
|
||||||
|
cached: bool = False
|
||||||
|
usage: ChatUsage | None = None
|
||||||
|
error: str | None = None
|
||||||
|
max_context_window: int | None = None
|
||||||
|
max_output_tokens: int | None = None
|
||||||
|
elapsed: float | None = None
|
||||||
|
timestamp: str | None = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, data: dict[str, Any]) -> "ChatResponse":
|
||||||
|
return cls(
|
||||||
|
response=data.get("response", ""),
|
||||||
|
prompt=data.get("prompt", ""),
|
||||||
|
json_mode=bool(data.get("json_mode", False)),
|
||||||
|
cached=bool(data.get("cached", False)),
|
||||||
|
usage=ChatUsage.from_dict(data.get("usage")),
|
||||||
|
error=data.get("error"),
|
||||||
|
max_context_window=data.get("max_context_window"),
|
||||||
|
max_output_tokens=data.get("max_output_tokens"),
|
||||||
|
elapsed=_as_float(data.get("elapsed")),
|
||||||
|
timestamp=data.get("timestamp"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DescribeResponse:
|
||||||
|
description: str = ""
|
||||||
|
url: str | None = None
|
||||||
|
mime_type: str | None = None
|
||||||
|
size: int | None = None
|
||||||
|
elapsed: float | None = None
|
||||||
|
timestamp: str | None = None
|
||||||
|
success: bool = True
|
||||||
|
error: str | None = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, data: dict[str, Any]) -> "DescribeResponse":
|
||||||
|
return cls(
|
||||||
|
description=data.get("description", ""),
|
||||||
|
url=data.get("url"),
|
||||||
|
mime_type=data.get("mime_type"),
|
||||||
|
size=data.get("size"),
|
||||||
|
elapsed=_as_float(data.get("elapsed")),
|
||||||
|
timestamp=data.get("timestamp"),
|
||||||
|
success=bool(data.get("success", True)),
|
||||||
|
error=data.get("error"),
|
||||||
|
)
|
||||||
|
|
||||||
234
src/typosaurus_sandbox/research/frontier.py
Normal file
234
src/typosaurus_sandbox/research/frontier.py
Normal file
@ -0,0 +1,234 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import hashlib
|
||||||
|
import logging
|
||||||
|
import re
|
||||||
|
import threading
|
||||||
|
import urllib.parse
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
from typosaurus_sandbox.research.envelopes import SearchResult
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
MIN_QUERY_LENGTH = 2
|
||||||
|
MAX_QUERY_LENGTH = 200
|
||||||
|
DEFAULT_PORTS: dict[str, int] = {"http": 80, "https": 443}
|
||||||
|
|
||||||
|
|
||||||
|
def _clean_text(value: str) -> str:
|
||||||
|
return " ".join(value.split())
|
||||||
|
|
||||||
|
|
||||||
|
def _query_key(query: str) -> str:
|
||||||
|
return _clean_text(query).casefold()
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_url(url: str) -> str:
|
||||||
|
cleaned = _clean_text(url)
|
||||||
|
try:
|
||||||
|
parsed = urllib.parse.urlsplit(cleaned)
|
||||||
|
except ValueError:
|
||||||
|
return cleaned
|
||||||
|
scheme = parsed.scheme.lower()
|
||||||
|
if scheme not in DEFAULT_PORTS:
|
||||||
|
return cleaned
|
||||||
|
host = (parsed.hostname or "").lower()
|
||||||
|
if not host:
|
||||||
|
return cleaned
|
||||||
|
try:
|
||||||
|
host = host.encode("idna").decode("ascii")
|
||||||
|
except UnicodeError:
|
||||||
|
pass
|
||||||
|
port: int | None = None
|
||||||
|
try:
|
||||||
|
port = parsed.port
|
||||||
|
except ValueError:
|
||||||
|
port = None
|
||||||
|
if port is not None and DEFAULT_PORTS.get(scheme) == port:
|
||||||
|
port = None
|
||||||
|
display_host = f"[{host}]" if ":" in host else host
|
||||||
|
netloc = display_host if port is None else f"{display_host}:{port}"
|
||||||
|
path = re.sub(r"/{2,}", "/", parsed.path)
|
||||||
|
if len(path) > 1 and path.endswith("/"):
|
||||||
|
path = path[:-1]
|
||||||
|
if parsed.query:
|
||||||
|
return f"{scheme}://{netloc}{path}?{parsed.query}"
|
||||||
|
return f"{scheme}://{netloc}{path}"
|
||||||
|
|
||||||
|
|
||||||
|
def fingerprint_text(text: str) -> str:
|
||||||
|
normalized = _clean_text(text)
|
||||||
|
return hashlib.sha256(normalized.encode("utf-8")).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
def query_variants_from_result(result: SearchResult) -> list[tuple[str, str]]:
|
||||||
|
variants: list[tuple[str, str]] = []
|
||||||
|
if result.title:
|
||||||
|
variants.append((result.title, "title"))
|
||||||
|
if result.description:
|
||||||
|
variants.append((result.description, "description"))
|
||||||
|
for value in result.extra.values():
|
||||||
|
if isinstance(value, str) and value:
|
||||||
|
variants.append((value, "extra"))
|
||||||
|
return variants
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class DedupStats:
|
||||||
|
queries_generated: int = 0
|
||||||
|
queries_enqueued: int = 0
|
||||||
|
queries_issued: int = 0
|
||||||
|
queries_duplicates_skipped: int = 0
|
||||||
|
urls_seen: int = 0
|
||||||
|
urls_duplicates_skipped: int = 0
|
||||||
|
content_seen: int = 0
|
||||||
|
content_duplicates_skipped: int = 0
|
||||||
|
|
||||||
|
def to_dict(self) -> dict[str, int]:
|
||||||
|
return {
|
||||||
|
"queries_generated": self.queries_generated,
|
||||||
|
"queries_enqueued": self.queries_enqueued,
|
||||||
|
"queries_issued": self.queries_issued,
|
||||||
|
"queries_duplicates_skipped": self.queries_duplicates_skipped,
|
||||||
|
"urls_seen": self.urls_seen,
|
||||||
|
"urls_duplicates_skipped": self.urls_duplicates_skipped,
|
||||||
|
"content_seen": self.content_seen,
|
||||||
|
"content_duplicates_skipped": self.content_duplicates_skipped,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class QueryFrontier:
|
||||||
|
def __init__(self, subject: str | None = None) -> None:
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
self._seen_queries: set[str] = set()
|
||||||
|
self._seen_urls: set[str] = set()
|
||||||
|
self._seen_url_order: list[str] = []
|
||||||
|
self._seen_content: set[str] = set()
|
||||||
|
self._origins: dict[str, str] = {}
|
||||||
|
self._pending: asyncio.Queue[str] = asyncio.Queue()
|
||||||
|
self._queries_generated = 0
|
||||||
|
self._queries_enqueued = 0
|
||||||
|
self._queries_issued = 0
|
||||||
|
self._queries_duplicates_skipped = 0
|
||||||
|
self._urls_seen = 0
|
||||||
|
self._urls_duplicates_skipped = 0
|
||||||
|
self._content_seen = 0
|
||||||
|
self._content_duplicates_skipped = 0
|
||||||
|
if subject:
|
||||||
|
self.seed(subject)
|
||||||
|
|
||||||
|
def seed(self, subject: str) -> None:
|
||||||
|
cleaned = _clean_text(subject)
|
||||||
|
if cleaned:
|
||||||
|
self.push_query(cleaned, "seed")
|
||||||
|
logger.info("frontier seeded subject=%r", cleaned)
|
||||||
|
|
||||||
|
def push_query(self, query: str, origin: str = "manual") -> bool:
|
||||||
|
cleaned = _clean_text(query)
|
||||||
|
if not MIN_QUERY_LENGTH <= len(cleaned) <= MAX_QUERY_LENGTH:
|
||||||
|
logger.debug("query variant invalid length=%d query=%r", len(cleaned), cleaned)
|
||||||
|
return False
|
||||||
|
key = _query_key(cleaned)
|
||||||
|
with self._lock:
|
||||||
|
self._queries_generated += 1
|
||||||
|
if key in self._seen_queries:
|
||||||
|
self._queries_duplicates_skipped += 1
|
||||||
|
logger.debug("query duplicate skipped origin=%s query=%r", origin, cleaned)
|
||||||
|
return False
|
||||||
|
self._seen_queries.add(key)
|
||||||
|
self._origins[key] = origin
|
||||||
|
self._queries_enqueued += 1
|
||||||
|
self._pending.put_nowait(cleaned)
|
||||||
|
logger.info("query enqueued origin=%s query=%r", origin, cleaned)
|
||||||
|
return True
|
||||||
|
|
||||||
|
def push_variants_from_result(self, result: SearchResult) -> int:
|
||||||
|
new_queries = 0
|
||||||
|
for text, origin in query_variants_from_result(result):
|
||||||
|
if self.push_query(text, origin):
|
||||||
|
new_queries += 1
|
||||||
|
return new_queries
|
||||||
|
|
||||||
|
def register_url(self, url: str) -> bool:
|
||||||
|
if not url:
|
||||||
|
return False
|
||||||
|
normalized = normalize_url(url)
|
||||||
|
with self._lock:
|
||||||
|
if normalized in self._seen_urls:
|
||||||
|
self._urls_duplicates_skipped += 1
|
||||||
|
logger.debug("url duplicate skipped url=%s", normalized)
|
||||||
|
return False
|
||||||
|
self._seen_urls.add(normalized)
|
||||||
|
self._seen_url_order.append(normalized)
|
||||||
|
self._urls_seen += 1
|
||||||
|
logger.info("url registered url=%s", normalized)
|
||||||
|
return True
|
||||||
|
|
||||||
|
def urls_since(self, seen_count: int) -> list[str]:
|
||||||
|
with self._lock:
|
||||||
|
return list(self._seen_url_order[seen_count:])
|
||||||
|
|
||||||
|
def register_content(self, text: str) -> bool:
|
||||||
|
if not text.strip():
|
||||||
|
return False
|
||||||
|
fingerprint = fingerprint_text(text)
|
||||||
|
with self._lock:
|
||||||
|
if fingerprint in self._seen_content:
|
||||||
|
self._content_duplicates_skipped += 1
|
||||||
|
logger.debug("content duplicate skipped fingerprint=%s", fingerprint)
|
||||||
|
return False
|
||||||
|
self._seen_content.add(fingerprint)
|
||||||
|
self._content_seen += 1
|
||||||
|
logger.info("content registered fingerprint=%s", fingerprint)
|
||||||
|
return True
|
||||||
|
|
||||||
|
def register_result(self, result: SearchResult) -> bool:
|
||||||
|
is_new = self.register_url(result.url)
|
||||||
|
if result.content:
|
||||||
|
self.register_content(result.content)
|
||||||
|
return is_new
|
||||||
|
|
||||||
|
async def get_query(self) -> str:
|
||||||
|
query = await self._pending.get()
|
||||||
|
with self._lock:
|
||||||
|
self._queries_issued += 1
|
||||||
|
logger.info("query issued query=%r", query)
|
||||||
|
return query
|
||||||
|
|
||||||
|
def pop_query(self) -> str | None:
|
||||||
|
try:
|
||||||
|
query = self._pending.get_nowait()
|
||||||
|
except asyncio.QueueEmpty:
|
||||||
|
return None
|
||||||
|
with self._lock:
|
||||||
|
self._queries_issued += 1
|
||||||
|
logger.info("query issued query=%r", query)
|
||||||
|
return query
|
||||||
|
|
||||||
|
def pending_count(self) -> int:
|
||||||
|
return self._pending.qsize()
|
||||||
|
|
||||||
|
def has_pending(self) -> bool:
|
||||||
|
return not self._pending.empty()
|
||||||
|
|
||||||
|
def origin_of(self, query: str) -> str | None:
|
||||||
|
with self._lock:
|
||||||
|
return self._origins.get(_query_key(query))
|
||||||
|
|
||||||
|
def snapshot(self) -> DedupStats:
|
||||||
|
with self._lock:
|
||||||
|
return DedupStats(
|
||||||
|
queries_generated=self._queries_generated,
|
||||||
|
queries_enqueued=self._queries_enqueued,
|
||||||
|
queries_issued=self._queries_issued,
|
||||||
|
queries_duplicates_skipped=self._queries_duplicates_skipped,
|
||||||
|
urls_seen=self._urls_seen,
|
||||||
|
urls_duplicates_skipped=self._urls_duplicates_skipped,
|
||||||
|
content_seen=self._content_seen,
|
||||||
|
content_duplicates_skipped=self._content_duplicates_skipped,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
345
src/typosaurus_sandbox/research/pipeline.py
Normal file
345
src/typosaurus_sandbox/research/pipeline.py
Normal file
@ -0,0 +1,345 @@
|
|||||||
|
# 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<>\"']+")
|
||||||
|
|
||||||
|
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:
|
||||||
|
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)
|
||||||
|
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,
|
||||||
|
failure.status_code,
|
||||||
|
cache_hit,
|
||||||
|
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",
|
||||||
|
)
|
||||||
|
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 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())
|
||||||
|
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
|
||||||
|
|
||||||
700
tests/test_research_client.py
Normal file
700
tests/test_research_client.py
Normal file
@ -0,0 +1,700 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import io
|
||||||
|
import json
|
||||||
|
import unittest
|
||||||
|
import urllib.error
|
||||||
|
from typing import Any
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
from typosaurus_sandbox.research.client import RsearchClient, RsearchError
|
||||||
|
from typosaurus_sandbox.research.envelopes import (
|
||||||
|
ChatResponse,
|
||||||
|
ChatUsage,
|
||||||
|
DeepReport,
|
||||||
|
DescribeResponse,
|
||||||
|
SearchGrade,
|
||||||
|
SearchResponse,
|
||||||
|
SearchResult,
|
||||||
|
)
|
||||||
|
|
||||||
|
WEB_RESPONSE: dict[str, Any] = {
|
||||||
|
"query": "asyncio python",
|
||||||
|
"source": "duckduckgo",
|
||||||
|
"count": 3,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"timestamp": "2026-08-07T12:00:00Z",
|
||||||
|
"results": [
|
||||||
|
{
|
||||||
|
"title": "asyncio documentation",
|
||||||
|
"url": "https://docs.python.org/3/library/asyncio.html",
|
||||||
|
"description": "Asynchronous I/O event loop.",
|
||||||
|
"source": "docs.python.org",
|
||||||
|
"extra": {"rank": 1},
|
||||||
|
"index": 0,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"title": "asyncio in Python",
|
||||||
|
"url": "https://example.com/asyncio",
|
||||||
|
"description": "Tutorial on asyncio.",
|
||||||
|
"source": "example.com",
|
||||||
|
"extra": {"rank": 2},
|
||||||
|
"index": 1,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
AI_MEMORY_RESPONSE: dict[str, Any] = {
|
||||||
|
"query": "python history",
|
||||||
|
"source": "ai",
|
||||||
|
"count": 0,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [],
|
||||||
|
"ai_response": "From memory: Python was released in 1991 by Guido van Rossum.",
|
||||||
|
"ai_error": None,
|
||||||
|
}
|
||||||
|
|
||||||
|
AI_PROVIDER_RESPONSE: dict[str, Any] = {
|
||||||
|
"query": "quantum computing",
|
||||||
|
"source": "google",
|
||||||
|
"count": 1,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [
|
||||||
|
{
|
||||||
|
"title": "Quantum computing overview",
|
||||||
|
"url": "https://example.com/quantum",
|
||||||
|
"description": "Overview of quantum computing.",
|
||||||
|
"source": "example.com",
|
||||||
|
"extra": {},
|
||||||
|
"index": 0,
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"ai_response": "Quantum computing uses qubits. [citation:1]",
|
||||||
|
"ai_error": None,
|
||||||
|
}
|
||||||
|
|
||||||
|
GRADED_RESPONSE: dict[str, Any] = {
|
||||||
|
"query": "deep research",
|
||||||
|
"source": "google",
|
||||||
|
"count": 1,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [
|
||||||
|
{
|
||||||
|
"title": "Deep research systems",
|
||||||
|
"url": "https://example.com/deep-research",
|
||||||
|
"description": "Survey of deep research systems.",
|
||||||
|
"source": "example.com",
|
||||||
|
"extra": {},
|
||||||
|
"index": 0,
|
||||||
|
"grade": {
|
||||||
|
"overall": 9.2,
|
||||||
|
"relevance": 8.8,
|
||||||
|
"depth": 9.0,
|
||||||
|
"authority": 9.5,
|
||||||
|
"freshness": 7.0,
|
||||||
|
"word_count": 1200,
|
||||||
|
"intent_hits": 4,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
DEEP_RESPONSE: dict[str, Any] = {
|
||||||
|
"query": "deep research systems",
|
||||||
|
"source": "google",
|
||||||
|
"count": 8,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [
|
||||||
|
{
|
||||||
|
"title": "Deep research systems",
|
||||||
|
"url": "https://example.com/deep-research",
|
||||||
|
"description": "Survey of deep research systems.",
|
||||||
|
"source": "example.com",
|
||||||
|
"extra": {},
|
||||||
|
"index": 0,
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"deep": {
|
||||||
|
"query": "deep research systems",
|
||||||
|
"markdown": "# Deep research\n\nA survey.",
|
||||||
|
"sources": [
|
||||||
|
{
|
||||||
|
"title": "Deep research systems",
|
||||||
|
"url": "https://example.com/deep-research",
|
||||||
|
"description": "Survey of deep research systems.",
|
||||||
|
"source": "example.com",
|
||||||
|
"extra": {},
|
||||||
|
"grade": {
|
||||||
|
"overall": 9.2,
|
||||||
|
"relevance": 8.8,
|
||||||
|
"depth": 9.0,
|
||||||
|
"authority": 9.5,
|
||||||
|
"freshness": 7.0,
|
||||||
|
"word_count": 1200,
|
||||||
|
"intent_hits": 4,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"graded_count": 8,
|
||||||
|
"total_count": 10,
|
||||||
|
"model": "gemma-3-12b-it",
|
||||||
|
"elapsed": 166.96,
|
||||||
|
"cache_hit": False,
|
||||||
|
"rounds": 3,
|
||||||
|
"queries_tried": ["deep research systems", "deep research architecture"],
|
||||||
|
"error": None,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
IMAGES_RESPONSE: dict[str, Any] = {
|
||||||
|
"query": "aurora borealis",
|
||||||
|
"source": "wikimedia",
|
||||||
|
"count": 2,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [
|
||||||
|
{
|
||||||
|
"title": "Aurora borealis over Norway",
|
||||||
|
"url": "https://commons.wikimedia.org/wiki/File:Aurora.jpg",
|
||||||
|
"description": "Photograph of the aurora borealis.",
|
||||||
|
"source": "wikimedia",
|
||||||
|
"extra": {
|
||||||
|
"thumbnail": "https://upload.wikimedia.org/thumb.jpg",
|
||||||
|
"dimensions": {"width": 1920, "height": 1080},
|
||||||
|
"mime": "image/jpeg",
|
||||||
|
"license": "CC BY-SA 4.0",
|
||||||
|
},
|
||||||
|
"index": 0,
|
||||||
|
}
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
CHAT_RESPONSE: dict[str, Any] = {
|
||||||
|
"response": "The answer.",
|
||||||
|
"prompt": "question",
|
||||||
|
"json_mode": True,
|
||||||
|
"cached": False,
|
||||||
|
"error": None,
|
||||||
|
"usage": {
|
||||||
|
"prompt_tokens": 120,
|
||||||
|
"completion_tokens": 80,
|
||||||
|
"total_tokens": 200,
|
||||||
|
"cost_usd": 0.0012,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
DESCRIBE_RESPONSE: dict[str, Any] = {
|
||||||
|
"url": "https://example.com/page",
|
||||||
|
"description": "Page description",
|
||||||
|
"elapsed": 1.23,
|
||||||
|
"timestamp": "2026-08-07T12:00:00Z",
|
||||||
|
}
|
||||||
|
|
||||||
|
SEARCH_EMPTY_OK: dict[str, Any] = {
|
||||||
|
"query": "q",
|
||||||
|
"source": "s",
|
||||||
|
"count": 1,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _recorded_request(fixture: dict[str, Any]) -> tuple[list[tuple[Any, ...]], Any]:
|
||||||
|
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))
|
||||||
|
return 200, fixture
|
||||||
|
|
||||||
|
return recorded, fake
|
||||||
|
|
||||||
|
|
||||||
|
def _raising_request(message: str, status_code: int) -> 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]]:
|
||||||
|
raise RsearchError(message, status_code)
|
||||||
|
|
||||||
|
return fake
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeResponse:
|
||||||
|
def __init__(self, status: int, body: bytes) -> None:
|
||||||
|
self.status = status
|
||||||
|
self._body = body
|
||||||
|
|
||||||
|
def __enter__(self) -> "_FakeResponse":
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, *args: object) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
def read(self) -> bytes:
|
||||||
|
return self._body
|
||||||
|
|
||||||
|
|
||||||
|
class TestSearchResponseParsing(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_web_results_parse_into_search_response(self) -> None:
|
||||||
|
response = SearchResponse.from_dict(WEB_RESPONSE)
|
||||||
|
self.assertEqual(response.query, "asyncio python")
|
||||||
|
self.assertEqual(response.source, "duckduckgo")
|
||||||
|
self.assertEqual(response.count, 3)
|
||||||
|
self.assertTrue(response.success)
|
||||||
|
self.assertIsNone(response.error)
|
||||||
|
self.assertEqual(response.timestamp, "2026-08-07T12:00:00Z")
|
||||||
|
self.assertEqual(len(response.results), 2)
|
||||||
|
first = response.results[0]
|
||||||
|
self.assertIsInstance(first, SearchResult)
|
||||||
|
self.assertEqual(first.title, "asyncio documentation")
|
||||||
|
self.assertEqual(first.url, "https://docs.python.org/3/library/asyncio.html")
|
||||||
|
self.assertEqual(first.description, "Asynchronous I/O event loop.")
|
||||||
|
self.assertEqual(first.source, "docs.python.org")
|
||||||
|
self.assertEqual(first.extra, {"rank": 1})
|
||||||
|
self.assertEqual(first.index, 0)
|
||||||
|
self.assertIsNone(first.content)
|
||||||
|
self.assertIsNone(first.grade)
|
||||||
|
self.assertIsNone(first.query_origin)
|
||||||
|
self.assertIsNone(response.ai_response)
|
||||||
|
self.assertIsNone(response.deep)
|
||||||
|
|
||||||
|
def test_ai_memory_variant_parses(self) -> None:
|
||||||
|
response = SearchResponse.from_dict(AI_MEMORY_RESPONSE)
|
||||||
|
self.assertEqual(response.source, "ai")
|
||||||
|
self.assertEqual(response.results, [])
|
||||||
|
self.assertIn("From memory", response.ai_response)
|
||||||
|
self.assertIsNone(response.ai_error)
|
||||||
|
|
||||||
|
def test_ai_provider_variant_parses(self) -> None:
|
||||||
|
response = SearchResponse.from_dict(AI_PROVIDER_RESPONSE)
|
||||||
|
self.assertEqual(response.source, "google")
|
||||||
|
self.assertEqual(len(response.results), 1)
|
||||||
|
self.assertIn("[citation:1]", response.ai_response)
|
||||||
|
self.assertIsNone(response.ai_error)
|
||||||
|
|
||||||
|
def test_deep_block_parses_into_deep_report(self) -> None:
|
||||||
|
response = SearchResponse.from_dict(DEEP_RESPONSE)
|
||||||
|
self.assertIsNotNone(response.deep)
|
||||||
|
deep = response.deep
|
||||||
|
self.assertIsInstance(deep, DeepReport)
|
||||||
|
self.assertEqual(deep.query, "deep research systems")
|
||||||
|
self.assertEqual(deep.markdown, "# Deep research\n\nA survey.")
|
||||||
|
self.assertEqual(deep.graded_count, 8)
|
||||||
|
self.assertEqual(deep.total_count, 10)
|
||||||
|
self.assertEqual(deep.model, "gemma-3-12b-it")
|
||||||
|
self.assertEqual(deep.elapsed, 166.96)
|
||||||
|
self.assertFalse(deep.cache_hit)
|
||||||
|
self.assertEqual(deep.rounds, 3)
|
||||||
|
self.assertEqual(deep.queries_tried, ["deep research systems", "deep research architecture"])
|
||||||
|
self.assertIsNone(deep.error)
|
||||||
|
self.assertEqual(len(deep.sources), 1)
|
||||||
|
source = deep.sources[0]
|
||||||
|
self.assertIsInstance(source, SearchResult)
|
||||||
|
self.assertEqual(source.url, "https://example.com/deep-research")
|
||||||
|
self.assertIsInstance(source.grade, SearchGrade)
|
||||||
|
self.assertEqual(source.grade.overall, 9.2)
|
||||||
|
|
||||||
|
def test_images_results_parse_extra_metadata(self) -> None:
|
||||||
|
response = SearchResponse.from_dict(IMAGES_RESPONSE)
|
||||||
|
self.assertEqual(response.source, "wikimedia")
|
||||||
|
result = response.results[0]
|
||||||
|
self.assertEqual(result.extra["mime"], "image/jpeg")
|
||||||
|
self.assertEqual(result.extra["dimensions"], {"width": 1920, "height": 1080})
|
||||||
|
self.assertEqual(result.extra["license"], "CC BY-SA 4.0")
|
||||||
|
self.assertIn("thumbnail", result.extra)
|
||||||
|
|
||||||
|
def test_result_grade_parses_into_search_grade(self) -> None:
|
||||||
|
response = SearchResponse.from_dict(GRADED_RESPONSE)
|
||||||
|
grade = response.results[0].grade
|
||||||
|
self.assertIsInstance(grade, SearchGrade)
|
||||||
|
self.assertEqual(grade.overall, 9.2)
|
||||||
|
self.assertEqual(grade.relevance, 8.8)
|
||||||
|
self.assertEqual(grade.depth, 9.0)
|
||||||
|
self.assertEqual(grade.authority, 9.5)
|
||||||
|
self.assertEqual(grade.freshness, 7.0)
|
||||||
|
self.assertEqual(grade.word_count, 1200)
|
||||||
|
self.assertEqual(grade.intent_hits, 4)
|
||||||
|
|
||||||
|
def test_sparse_body_parses_with_defaults(self) -> None:
|
||||||
|
response = SearchResponse.from_dict({"query": "x", "success": True})
|
||||||
|
self.assertEqual(response.source, "")
|
||||||
|
self.assertEqual(response.count, 0)
|
||||||
|
self.assertEqual(response.results, [])
|
||||||
|
self.assertIsNone(response.error)
|
||||||
|
self.assertIsNone(response.deep)
|
||||||
|
self.assertIsNone(response.ai_response)
|
||||||
|
|
||||||
|
|
||||||
|
class TestSearchRequestConstruction(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_search_forwards_all_parameters(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(SEARCH_EMPTY_OK)
|
||||||
|
client._request = fake
|
||||||
|
await client.search(
|
||||||
|
"query text",
|
||||||
|
source="google",
|
||||||
|
count=7,
|
||||||
|
content=True,
|
||||||
|
type="images",
|
||||||
|
deep=True,
|
||||||
|
ai=True,
|
||||||
|
cache=False,
|
||||||
|
)
|
||||||
|
method, path, params, payload, headers, timeout = recorded[0]
|
||||||
|
self.assertEqual(method, "GET")
|
||||||
|
self.assertEqual(path, "/search")
|
||||||
|
self.assertEqual(
|
||||||
|
params,
|
||||||
|
{
|
||||||
|
"query": "query text",
|
||||||
|
"source": "google",
|
||||||
|
"count": "7",
|
||||||
|
"content": "true",
|
||||||
|
"type": "images",
|
||||||
|
"deep": "true",
|
||||||
|
"ai": "true",
|
||||||
|
"cache": "false",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
self.assertIsNone(payload)
|
||||||
|
self.assertIsNone(headers)
|
||||||
|
self.assertEqual(timeout, 180.0)
|
||||||
|
|
||||||
|
async def test_search_without_deep_uses_request_timeout(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(SEARCH_EMPTY_OK)
|
||||||
|
client._request = fake
|
||||||
|
await client.search("q")
|
||||||
|
self.assertEqual(recorded[0][5], 30.0)
|
||||||
|
|
||||||
|
async def test_count_none_omits_count_parameter(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(SEARCH_EMPTY_OK)
|
||||||
|
client._request = fake
|
||||||
|
await client.search("q")
|
||||||
|
self.assertNotIn("count", recorded[0][2])
|
||||||
|
|
||||||
|
async def test_count_zero_forwarded_and_server_clamp_parsed(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(
|
||||||
|
{"query": "q", "source": "s", "count": 1, "success": True, "error": None, "results": []}
|
||||||
|
)
|
||||||
|
client._request = fake
|
||||||
|
response = await client.search("q", count=0)
|
||||||
|
self.assertEqual(recorded[0][2]["count"], "0")
|
||||||
|
self.assertEqual(response.count, 1)
|
||||||
|
|
||||||
|
async def test_count_above_limit_forwarded_and_server_clamp_parsed(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(
|
||||||
|
{"query": "q", "source": "s", "count": 10, "success": True, "error": None, "results": []}
|
||||||
|
)
|
||||||
|
client._request = fake
|
||||||
|
response = await client.search("q", count=25)
|
||||||
|
self.assertEqual(recorded[0][2]["count"], "25")
|
||||||
|
self.assertEqual(response.count, 10)
|
||||||
|
|
||||||
|
async def test_invalid_count_forwarded_and_server_clamp_parsed(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(
|
||||||
|
{"query": "q", "source": "s", "count": 10, "success": True, "error": None, "results": []}
|
||||||
|
)
|
||||||
|
client._request = fake
|
||||||
|
response = await client.search("q", count="not-a-number")
|
||||||
|
self.assertEqual(recorded[0][2]["count"], "not-a-number")
|
||||||
|
self.assertEqual(response.count, 10)
|
||||||
|
|
||||||
|
async def test_search_cache_hit_skips_second_request(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(SEARCH_EMPTY_OK)
|
||||||
|
client._request = fake
|
||||||
|
await client.search("cached query")
|
||||||
|
await client.search("cached query")
|
||||||
|
self.assertEqual(len(recorded), 1)
|
||||||
|
|
||||||
|
async def test_search_cache_disabled_repeats_request(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(SEARCH_EMPTY_OK)
|
||||||
|
client._request = fake
|
||||||
|
await client.search("uncached query", cache=False)
|
||||||
|
await client.search("uncached query", cache=False)
|
||||||
|
self.assertEqual(len(recorded), 2)
|
||||||
|
|
||||||
|
async def test_search_with_content_populates_content_cache(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
fixture = {
|
||||||
|
"query": "q",
|
||||||
|
"source": "s",
|
||||||
|
"count": 1,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [
|
||||||
|
{
|
||||||
|
"title": "t",
|
||||||
|
"url": "https://example.com/a",
|
||||||
|
"description": "d",
|
||||||
|
"source": "s",
|
||||||
|
"extra": {},
|
||||||
|
"content": "full page text",
|
||||||
|
}
|
||||||
|
],
|
||||||
|
}
|
||||||
|
recorded, fake = _recorded_request(fixture)
|
||||||
|
client._request = fake
|
||||||
|
await client.search("q", content=True)
|
||||||
|
self.assertEqual(len(recorded), 1)
|
||||||
|
self.assertEqual(client.get_cached_content("https://example.com/a"), "full page text")
|
||||||
|
|
||||||
|
async def test_search_error_in_body_surfaces_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
client._request = _raising_request("Empty query", 400)
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
await client.search("")
|
||||||
|
self.assertEqual(ctx.exception.status_code, 400)
|
||||||
|
self.assertEqual(str(ctx.exception), "Empty query")
|
||||||
|
|
||||||
|
|
||||||
|
class TestChatResponseParsing(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_chat_response_parses_envelope(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(CHAT_RESPONSE)
|
||||||
|
client._request = fake
|
||||||
|
response = await client.chat("question", json_mode=True)
|
||||||
|
method, path, params, payload, headers, timeout = recorded[0]
|
||||||
|
self.assertEqual(method, "POST")
|
||||||
|
self.assertEqual(path, "/chat")
|
||||||
|
self.assertEqual(json.loads(payload), {"prompt": "question", "json": True})
|
||||||
|
self.assertEqual(headers, {"Content-Type": "application/json"})
|
||||||
|
self.assertIsNone(params)
|
||||||
|
self.assertIsNone(timeout)
|
||||||
|
self.assertIsInstance(response, ChatResponse)
|
||||||
|
self.assertEqual(response.response, "The answer.")
|
||||||
|
self.assertEqual(response.prompt, "question")
|
||||||
|
self.assertTrue(response.json_mode)
|
||||||
|
self.assertFalse(response.cached)
|
||||||
|
self.assertIsNone(response.error)
|
||||||
|
self.assertIsInstance(response.usage, ChatUsage)
|
||||||
|
self.assertEqual(response.usage.prompt_tokens, 120)
|
||||||
|
self.assertEqual(response.usage.completion_tokens, 80)
|
||||||
|
self.assertEqual(response.usage.total_tokens, 200)
|
||||||
|
self.assertEqual(response.usage.cost_usd, 0.0012)
|
||||||
|
|
||||||
|
async def test_chat_request_accepts_system_and_disables_cache(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(CHAT_RESPONSE)
|
||||||
|
client._request = fake
|
||||||
|
await client.chat("q", system="sys", cache=False)
|
||||||
|
body = json.loads(recorded[0][3])
|
||||||
|
self.assertEqual(body, {"prompt": "q", "system": "sys", "cache": False})
|
||||||
|
|
||||||
|
async def test_chat_error_raises_mapped_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
client._request = _raising_request("No prompt provided", 400)
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
await client.chat("")
|
||||||
|
self.assertEqual(ctx.exception.status_code, 400)
|
||||||
|
self.assertEqual(str(ctx.exception), "No prompt provided")
|
||||||
|
|
||||||
|
|
||||||
|
class TestDescribeResponseParsing(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_describe_get_parses_envelope(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(DESCRIBE_RESPONSE)
|
||||||
|
client._request = fake
|
||||||
|
response = await client.describe("https://example.com/page")
|
||||||
|
method, path, params, payload, headers, timeout = recorded[0]
|
||||||
|
self.assertEqual(method, "GET")
|
||||||
|
self.assertEqual(path, "/describe")
|
||||||
|
self.assertEqual(params, {"url": "https://example.com/page"})
|
||||||
|
self.assertIsNone(payload)
|
||||||
|
self.assertIsNone(headers)
|
||||||
|
self.assertIsNone(timeout)
|
||||||
|
self.assertIsInstance(response, DescribeResponse)
|
||||||
|
self.assertEqual(response.description, "Page description")
|
||||||
|
self.assertEqual(response.url, "https://example.com/page")
|
||||||
|
self.assertEqual(response.elapsed, 1.23)
|
||||||
|
self.assertEqual(response.timestamp, "2026-08-07T12:00:00Z")
|
||||||
|
self.assertIsNone(response.mime_type)
|
||||||
|
self.assertIsNone(response.size)
|
||||||
|
self.assertTrue(response.success)
|
||||||
|
|
||||||
|
async def test_describe_get_uses_cache(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(DESCRIBE_RESPONSE)
|
||||||
|
client._request = fake
|
||||||
|
await client.describe("https://example.com/page")
|
||||||
|
await client.describe("https://example.com/page")
|
||||||
|
self.assertEqual(len(recorded), 1)
|
||||||
|
|
||||||
|
async def test_describe_error_raises_mapped_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
client._request = _raising_request("No url provided", 400)
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
await client.describe("")
|
||||||
|
self.assertEqual(ctx.exception.status_code, 400)
|
||||||
|
self.assertEqual(str(ctx.exception), "No url provided")
|
||||||
|
|
||||||
|
async def test_describe_raw_posts_bytes_with_content_type(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(DESCRIBE_RESPONSE)
|
||||||
|
client._request = fake
|
||||||
|
image = b"\x89PNG\r\n\x1a\npayload"
|
||||||
|
await client.describe_raw(image, mime_type="image/png")
|
||||||
|
method, path, params, payload, headers, timeout = recorded[0]
|
||||||
|
self.assertEqual(method, "POST")
|
||||||
|
self.assertEqual(path, "/describe")
|
||||||
|
self.assertEqual(payload, image)
|
||||||
|
self.assertEqual(headers, {"Content-Type": "image/png"})
|
||||||
|
|
||||||
|
async def test_describe_upload_builds_multipart_body(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(DESCRIBE_RESPONSE)
|
||||||
|
client._request = fake
|
||||||
|
image = b"\x89PNGpayload"
|
||||||
|
await client.describe_upload(image, filename="photo.png", mime_type="image/png")
|
||||||
|
method, path, params, payload, headers, timeout = recorded[0]
|
||||||
|
self.assertEqual(method, "POST")
|
||||||
|
self.assertEqual(path, "/describe")
|
||||||
|
self.assertIn(b'name="file"; filename="photo.png"', payload)
|
||||||
|
self.assertIn(b"Content-Type: image/png", payload)
|
||||||
|
self.assertIn(image, payload)
|
||||||
|
self.assertIn("multipart/form-data; boundary=", headers["Content-Type"])
|
||||||
|
|
||||||
|
async def test_describe_raw_reuses_cache_by_hash(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
recorded, fake = _recorded_request(DESCRIBE_RESPONSE)
|
||||||
|
client._request = fake
|
||||||
|
image = b"\x89PNGpayload"
|
||||||
|
await client.describe_raw(image, mime_type="image/png")
|
||||||
|
await client.describe_raw(image, mime_type="image/png")
|
||||||
|
self.assertEqual(len(recorded), 1)
|
||||||
|
|
||||||
|
|
||||||
|
class TestErrorInBodyHandling(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_empty_query_error_in_body_maps_to_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
body = b'{"success": false, "error": "Empty query"}'
|
||||||
|
error = urllib.error.HTTPError(
|
||||||
|
"https://rsearch.app.molodetz.nl/search", 400, "Bad Request", {}, io.BytesIO(body)
|
||||||
|
)
|
||||||
|
with mock.patch("urllib.request.urlopen", side_effect=error):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/search", {"query": ""}, None, None, 30.0)
|
||||||
|
self.assertEqual(ctx.exception.status_code, 400)
|
||||||
|
self.assertEqual(str(ctx.exception), "Empty query")
|
||||||
|
|
||||||
|
def test_providers_exhausted_503_maps_to_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
body = b'{"success": false, "error": "All providers are exhausted, please try again later"}'
|
||||||
|
error = urllib.error.HTTPError(
|
||||||
|
"https://rsearch.app.molodetz.nl/search", 503, "Service Unavailable", {}, io.BytesIO(body)
|
||||||
|
)
|
||||||
|
with mock.patch("urllib.request.urlopen", side_effect=error):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/search", {"query": "q"}, None, None, 30.0)
|
||||||
|
self.assertEqual(ctx.exception.status_code, 503)
|
||||||
|
self.assertEqual(str(ctx.exception), "All providers are exhausted, please try again later")
|
||||||
|
|
||||||
|
def test_success_false_body_with_http_200_raises(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
fake = _FakeResponse(200, b'{"success": false, "error": "Empty query"}')
|
||||||
|
with mock.patch("urllib.request.urlopen", return_value=fake):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/search", {"query": ""}, None, None, 30.0)
|
||||||
|
self.assertEqual(ctx.exception.status_code, 200)
|
||||||
|
self.assertEqual(str(ctx.exception), "Empty query")
|
||||||
|
|
||||||
|
def test_detail_field_falls_back_for_error_message(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
body = b'{"detail": "No url provided"}'
|
||||||
|
error = urllib.error.HTTPError(
|
||||||
|
"https://rsearch.app.molodetz.nl/describe", 400, "Bad Request", {}, io.BytesIO(body)
|
||||||
|
)
|
||||||
|
with mock.patch("urllib.request.urlopen", side_effect=error):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/describe", {"url": "x"}, None, None, 30.0)
|
||||||
|
self.assertEqual(ctx.exception.status_code, 400)
|
||||||
|
self.assertEqual(str(ctx.exception), "No url provided")
|
||||||
|
|
||||||
|
def test_title_field_falls_back_for_error_message(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
body = b'{"title": "Provider error"}'
|
||||||
|
error = urllib.error.HTTPError(
|
||||||
|
"https://rsearch.app.molodetz.nl/search", 502, "Bad Gateway", {}, io.BytesIO(body)
|
||||||
|
)
|
||||||
|
with mock.patch("urllib.request.urlopen", side_effect=error):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/search", {"query": "q"}, None, None, 30.0)
|
||||||
|
self.assertEqual(ctx.exception.status_code, 502)
|
||||||
|
self.assertEqual(str(ctx.exception), "Provider error")
|
||||||
|
|
||||||
|
def test_empty_body_raises_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
fake = _FakeResponse(200, b"")
|
||||||
|
with mock.patch("urllib.request.urlopen", return_value=fake):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/search", {"query": "q"}, None, None, 30.0)
|
||||||
|
self.assertEqual(ctx.exception.status_code, 200)
|
||||||
|
self.assertIn("empty response", str(ctx.exception))
|
||||||
|
|
||||||
|
def test_invalid_json_body_raises_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
fake = _FakeResponse(200, b"<html>not json</html>")
|
||||||
|
with mock.patch("urllib.request.urlopen", return_value=fake):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/search", {"query": "q"}, None, None, 30.0)
|
||||||
|
self.assertEqual(ctx.exception.status_code, 200)
|
||||||
|
self.assertIn("invalid JSON", str(ctx.exception))
|
||||||
|
|
||||||
|
def test_non_dict_body_raises_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
fake = _FakeResponse(200, b'["not", "a", "dict"]')
|
||||||
|
with mock.patch("urllib.request.urlopen", return_value=fake):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/search", {"query": "q"}, None, None, 30.0)
|
||||||
|
self.assertEqual(ctx.exception.status_code, 200)
|
||||||
|
|
||||||
|
def test_connection_failure_raises_rsearch_error(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
with mock.patch("urllib.request.urlopen", side_effect=urllib.error.URLError("connection refused")):
|
||||||
|
with self.assertRaises(RsearchError) as ctx:
|
||||||
|
client._request("GET", "/search", {"query": "q"}, None, None, 30.0)
|
||||||
|
self.assertIn("connection failure", str(ctx.exception))
|
||||||
|
|
||||||
|
def test_successful_request_returns_status_and_body(self) -> None:
|
||||||
|
client = RsearchClient()
|
||||||
|
fake = _FakeResponse(200, b'{"success": true, "query": "q", "count": 1, "results": []}')
|
||||||
|
with mock.patch("urllib.request.urlopen", return_value=fake) as urlopen:
|
||||||
|
status, data = client._request("GET", "/search", {"query": "q"}, None, None, 30.0)
|
||||||
|
self.assertEqual(status, 200)
|
||||||
|
self.assertEqual(data, {"success": True, "query": "q", "count": 1, "results": []})
|
||||||
|
request = urlopen.call_args[0][0]
|
||||||
|
self.assertEqual(request.get_method(), "GET")
|
||||||
|
self.assertEqual(request.get_full_url(), "https://rsearch.app.molodetz.nl/search?query=q")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
|
|
||||||
310
tests/test_research_dedup.py
Normal file
310
tests/test_research_dedup.py
Normal file
@ -0,0 +1,310 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
from typosaurus_sandbox.research.envelopes import SearchResult
|
||||||
|
from typosaurus_sandbox.research.frontier import (
|
||||||
|
MAX_QUERY_LENGTH,
|
||||||
|
MIN_QUERY_LENGTH,
|
||||||
|
DedupStats,
|
||||||
|
QueryFrontier,
|
||||||
|
fingerprint_text,
|
||||||
|
normalize_url,
|
||||||
|
query_variants_from_result,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _round_halts(frontier: QueryFrontier, before: DedupStats) -> bool:
|
||||||
|
after = frontier.snapshot()
|
||||||
|
new_urls = after.urls_seen - before.urls_seen
|
||||||
|
new_queries = after.queries_enqueued - before.queries_enqueued
|
||||||
|
return new_urls == 0 and new_queries == 0
|
||||||
|
|
||||||
|
|
||||||
|
class TestNormalizeUrl(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_lowercases_scheme_and_host_and_strips_default_port(self) -> None:
|
||||||
|
self.assertEqual(
|
||||||
|
normalize_url("HTTPS://Example.COM:443/Path//To//Page/"),
|
||||||
|
"https://example.com/Path/To/Page",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_strips_userinfo_and_fragment_keeps_query(self) -> None:
|
||||||
|
self.assertEqual(
|
||||||
|
normalize_url("https://user:pass@example.com:8443/a?x=1#frag"),
|
||||||
|
"https://example.com:8443/a?x=1",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_fragment_dropped_with_default_port(self) -> None:
|
||||||
|
self.assertEqual(normalize_url("https://example.com/a?x=1#sec"), "https://example.com/a?x=1")
|
||||||
|
|
||||||
|
def test_non_default_port_preserved(self) -> None:
|
||||||
|
self.assertEqual(normalize_url("https://example.com:80/x"), "https://example.com:80/x")
|
||||||
|
|
||||||
|
def test_idna_encodes_non_ascii_host(self) -> None:
|
||||||
|
self.assertEqual(normalize_url("https://MÜNCHEN.example/"), "https://xn--mnchen-3ya.example/")
|
||||||
|
|
||||||
|
def test_http_and_https_remain_distinct(self) -> None:
|
||||||
|
self.assertNotEqual(normalize_url("http://example.com/x"), normalize_url("https://example.com/x"))
|
||||||
|
|
||||||
|
def test_non_http_scheme_returned_cleaned(self) -> None:
|
||||||
|
self.assertEqual(normalize_url("not a url"), "not a url")
|
||||||
|
|
||||||
|
def test_blank_url_normalizes_to_empty(self) -> None:
|
||||||
|
self.assertEqual(normalize_url(" "), "")
|
||||||
|
|
||||||
|
|
||||||
|
class TestFingerprintText(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_whitespace_variants_produce_identical_fingerprint(self) -> None:
|
||||||
|
self.assertEqual(fingerprint_text("identical body\n\n"), fingerprint_text("identical body"))
|
||||||
|
|
||||||
|
def test_distinct_text_produces_distinct_fingerprint(self) -> None:
|
||||||
|
self.assertNotEqual(fingerprint_text("first text"), fingerprint_text("second text"))
|
||||||
|
|
||||||
|
def test_fingerprint_is_sha256_hex(self) -> None:
|
||||||
|
digest = fingerprint_text("sample")
|
||||||
|
self.assertEqual(len(digest), 64)
|
||||||
|
int(digest, 16)
|
||||||
|
|
||||||
|
|
||||||
|
class TestQueryVariantsFromResult(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_title_description_and_string_extra_become_variants(self) -> None:
|
||||||
|
result = SearchResult(
|
||||||
|
title="Deep research",
|
||||||
|
description="Survey of deep research systems",
|
||||||
|
url="https://a.example",
|
||||||
|
extra={"tag": "research methods", "rank": 3},
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
query_variants_from_result(result),
|
||||||
|
[
|
||||||
|
("Deep research", "title"),
|
||||||
|
("Survey of deep research systems", "description"),
|
||||||
|
("research methods", "extra"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_non_string_extra_values_ignored(self) -> None:
|
||||||
|
result = SearchResult(title="t", url="https://a.example", extra={"rank": 3, "ok": True})
|
||||||
|
self.assertEqual(query_variants_from_result(result), [("t", "title")])
|
||||||
|
|
||||||
|
def test_empty_fields_produce_no_variants(self) -> None:
|
||||||
|
result = SearchResult(url="https://a.example")
|
||||||
|
self.assertEqual(query_variants_from_result(result), [])
|
||||||
|
|
||||||
|
|
||||||
|
class TestUrlDeduplication(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_first_registration_records_url_once(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertTrue(frontier.register_url("https://example.com/page"))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.urls_seen, 1)
|
||||||
|
self.assertEqual(stats.urls_duplicates_skipped, 0)
|
||||||
|
|
||||||
|
async def test_same_url_registered_twice_rejects_second(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertTrue(frontier.register_url("https://example.com/page"))
|
||||||
|
self.assertFalse(frontier.register_url("https://example.com/page"))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.urls_seen, 1)
|
||||||
|
self.assertEqual(stats.urls_duplicates_skipped, 1)
|
||||||
|
|
||||||
|
async def test_normalized_variants_of_same_url_rejected(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertTrue(frontier.register_url("HTTPS://Example.COM:443/a//b/"))
|
||||||
|
self.assertFalse(frontier.register_url("https://example.com/a/b"))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.urls_seen, 1)
|
||||||
|
self.assertEqual(stats.urls_duplicates_skipped, 1)
|
||||||
|
|
||||||
|
async def test_duplicate_urls_across_responses_recorded_once(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
first = SearchResult(url="https://example.com/page", title="first title", description="first description")
|
||||||
|
second = SearchResult(url="https://example.com/page", title="second title", description="second description")
|
||||||
|
self.assertTrue(frontier.register_result(first))
|
||||||
|
self.assertFalse(frontier.register_result(second))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.urls_seen, 1)
|
||||||
|
self.assertEqual(stats.urls_duplicates_skipped, 1)
|
||||||
|
|
||||||
|
async def test_empty_url_rejected(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertFalse(frontier.register_url(""))
|
||||||
|
self.assertEqual(frontier.snapshot().urls_seen, 0)
|
||||||
|
|
||||||
|
async def test_whitespace_url_normalized_and_deduplicated(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertTrue(frontier.register_url(" "))
|
||||||
|
self.assertFalse(frontier.register_url(" "))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.urls_seen, 1)
|
||||||
|
self.assertEqual(stats.urls_duplicates_skipped, 1)
|
||||||
|
|
||||||
|
|
||||||
|
class TestContentDeduplication(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_identical_content_different_urls_rejects_second_occurrence(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
first = SearchResult(url="https://a.example/1", title="t1", description="d1", content="identical body")
|
||||||
|
second = SearchResult(url="https://b.example/2", title="t2", description="d2", content="identical body")
|
||||||
|
self.assertTrue(frontier.register_result(first))
|
||||||
|
self.assertTrue(frontier.register_result(second))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.urls_seen, 2)
|
||||||
|
self.assertEqual(stats.content_seen, 1)
|
||||||
|
self.assertEqual(stats.content_duplicates_skipped, 1)
|
||||||
|
self.assertFalse(frontier.register_content("identical body"))
|
||||||
|
|
||||||
|
async def test_near_identical_whitespace_content_rejected(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertTrue(frontier.register_content(" Deep research system \n"))
|
||||||
|
self.assertFalse(frontier.register_content("Deep research system"))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.content_seen, 1)
|
||||||
|
self.assertEqual(stats.content_duplicates_skipped, 1)
|
||||||
|
|
||||||
|
async def test_blank_content_rejected(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertFalse(frontier.register_content(""))
|
||||||
|
self.assertFalse(frontier.register_content(" \n "))
|
||||||
|
self.assertEqual(frontier.snapshot().content_seen, 0)
|
||||||
|
|
||||||
|
async def test_result_without_content_registers_url_only(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
result = SearchResult(url="https://a.example", title="t", description="d")
|
||||||
|
self.assertTrue(frontier.register_result(result))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.urls_seen, 1)
|
||||||
|
self.assertEqual(stats.content_seen, 0)
|
||||||
|
|
||||||
|
|
||||||
|
class TestQueryDeduplication(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_duplicate_query_rejected(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertTrue(frontier.push_query("deep research", "manual"))
|
||||||
|
self.assertFalse(frontier.push_query("deep research", "manual"))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.queries_enqueued, 1)
|
||||||
|
self.assertEqual(stats.queries_duplicates_skipped, 1)
|
||||||
|
|
||||||
|
async def test_query_dedup_ignores_case_and_whitespace(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertTrue(frontier.push_query(" Deep RESEARCH "))
|
||||||
|
self.assertFalse(frontier.push_query("deep research"))
|
||||||
|
self.assertEqual(frontier.snapshot().queries_enqueued, 1)
|
||||||
|
|
||||||
|
async def test_variants_from_result_deduplicated_across_fields(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
result = SearchResult(
|
||||||
|
title="Python asyncio",
|
||||||
|
description="python asyncio",
|
||||||
|
url="https://a.example",
|
||||||
|
extra={"tag": " PYTHON ASYNCIO "},
|
||||||
|
)
|
||||||
|
self.assertEqual(frontier.push_variants_from_result(result), 1)
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.queries_enqueued, 1)
|
||||||
|
self.assertEqual(stats.queries_duplicates_skipped, 2)
|
||||||
|
|
||||||
|
async def test_duplicate_query_never_issued_twice(self) -> None:
|
||||||
|
frontier = QueryFrontier("asyncio python")
|
||||||
|
self.assertEqual(frontier.pop_query(), "asyncio python")
|
||||||
|
self.assertFalse(frontier.push_query("ASYNCIO python"))
|
||||||
|
self.assertIsNone(frontier.pop_query())
|
||||||
|
self.assertEqual(frontier.snapshot().queries_issued, 1)
|
||||||
|
|
||||||
|
async def test_query_length_window_enforced(self) -> None:
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
self.assertFalse(frontier.push_query("a" * (MIN_QUERY_LENGTH - 1)))
|
||||||
|
self.assertTrue(frontier.push_query("a" * MIN_QUERY_LENGTH))
|
||||||
|
self.assertTrue(frontier.push_query("b" * MAX_QUERY_LENGTH))
|
||||||
|
self.assertFalse(frontier.push_query("c" * (MAX_QUERY_LENGTH + 1)))
|
||||||
|
self.assertEqual(frontier.snapshot().queries_enqueued, 2)
|
||||||
|
|
||||||
|
async def test_reseed_same_subject_enqueues_once(self) -> None:
|
||||||
|
frontier = QueryFrontier("subject alpha")
|
||||||
|
frontier.seed("SUBJECT ALPHA")
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.queries_enqueued, 1)
|
||||||
|
self.assertEqual(stats.queries_duplicates_skipped, 1)
|
||||||
|
|
||||||
|
|
||||||
|
class TestClosureDecision(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_round_with_no_new_urls_and_no_new_queries_halts(self) -> None:
|
||||||
|
frontier = QueryFrontier("subject alpha")
|
||||||
|
frontier.pop_query()
|
||||||
|
discovered = SearchResult(
|
||||||
|
url="https://a.example/page", title="alpha discovery", description="alpha details", content="body text"
|
||||||
|
)
|
||||||
|
self.assertTrue(frontier.register_result(discovered))
|
||||||
|
before = frontier.snapshot()
|
||||||
|
self.assertFalse(frontier.register_result(discovered))
|
||||||
|
self.assertFalse(frontier.push_query("SUBJECT ALPHA"))
|
||||||
|
after = frontier.snapshot()
|
||||||
|
self.assertEqual(after.urls_seen - before.urls_seen, 0)
|
||||||
|
self.assertEqual(after.queries_enqueued - before.queries_enqueued, 0)
|
||||||
|
self.assertTrue(_round_halts(frontier, before))
|
||||||
|
|
||||||
|
async def test_round_adding_new_url_continues(self) -> None:
|
||||||
|
frontier = QueryFrontier("subject beta")
|
||||||
|
frontier.pop_query()
|
||||||
|
before = frontier.snapshot()
|
||||||
|
self.assertTrue(frontier.register_url("https://new.example/x"))
|
||||||
|
after = frontier.snapshot()
|
||||||
|
self.assertEqual(after.urls_seen - before.urls_seen, 1)
|
||||||
|
self.assertFalse(_round_halts(frontier, before))
|
||||||
|
|
||||||
|
async def test_round_adding_new_query_continues(self) -> None:
|
||||||
|
frontier = QueryFrontier("subject gamma")
|
||||||
|
frontier.pop_query()
|
||||||
|
before = frontier.snapshot()
|
||||||
|
result = SearchResult(url="https://a.example", title="gamma subtopic", description="")
|
||||||
|
self.assertEqual(frontier.push_variants_from_result(result), 1)
|
||||||
|
after = frontier.snapshot()
|
||||||
|
self.assertEqual(after.queries_enqueued - before.queries_enqueued, 1)
|
||||||
|
self.assertFalse(_round_halts(frontier, before))
|
||||||
|
|
||||||
|
async def test_round_adding_url_and_query_continues(self) -> None:
|
||||||
|
frontier = QueryFrontier("subject delta")
|
||||||
|
frontier.pop_query()
|
||||||
|
before = frontier.snapshot()
|
||||||
|
result = SearchResult(url="https://b.example/page", title="delta subtopic", description="delta details")
|
||||||
|
self.assertTrue(frontier.register_result(result))
|
||||||
|
self.assertEqual(frontier.push_variants_from_result(result), 2)
|
||||||
|
after = frontier.snapshot()
|
||||||
|
self.assertEqual(after.urls_seen - before.urls_seen, 1)
|
||||||
|
self.assertEqual(after.queries_enqueued - before.queries_enqueued, 2)
|
||||||
|
self.assertFalse(_round_halts(frontier, before))
|
||||||
|
|
||||||
|
async def test_empty_result_round_halts_and_exhausts_pending(self) -> None:
|
||||||
|
frontier = QueryFrontier("lonely topic")
|
||||||
|
before = frontier.snapshot()
|
||||||
|
self.assertEqual(frontier.pop_query(), "lonely topic")
|
||||||
|
self.assertFalse(frontier.has_pending())
|
||||||
|
self.assertTrue(_round_halts(frontier, before))
|
||||||
|
|
||||||
|
async def test_snapshot_reports_round_deltas_for_closure_accounting(self) -> None:
|
||||||
|
frontier = QueryFrontier("subject epsilon")
|
||||||
|
before = frontier.snapshot()
|
||||||
|
self.assertEqual(frontier.pop_query(), "subject epsilon")
|
||||||
|
result = SearchResult(
|
||||||
|
url="https://c.example/page", title="epsilon topic", description="epsilon details", content="epsilon body"
|
||||||
|
)
|
||||||
|
self.assertTrue(frontier.register_result(result))
|
||||||
|
after = frontier.snapshot()
|
||||||
|
self.assertEqual(after.urls_seen - before.urls_seen, 1)
|
||||||
|
self.assertEqual(after.content_seen - before.content_seen, 1)
|
||||||
|
self.assertEqual(after.queries_issued - before.queries_issued, 1)
|
||||||
|
self.assertEqual(len(after.to_dict()), 8)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
305
tests/test_research_engine.py
Normal file
305
tests/test_research_engine.py
Normal file
@ -0,0 +1,305 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
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,
|
||||||
|
*,
|
||||||
|
web_results: list[SearchResult] | None = None,
|
||||||
|
web_result_factory: Callable[[str], list[SearchResult]] | None = None,
|
||||||
|
chat_text: str = "",
|
||||||
|
describe_text: str = "",
|
||||||
|
) -> None:
|
||||||
|
self.config = ResearchConfig(max_concurrency=4, default_count=5)
|
||||||
|
self._web_results = web_results if web_results is not None else []
|
||||||
|
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, bool, bool]] = []
|
||||||
|
|
||||||
|
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,
|
||||||
|
) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
def describe_cached(self, url: str) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
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=[])
|
||||||
|
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, False, False))
|
||||||
|
return ChatResponse(response=self._chat_text, prompt=prompt)
|
||||||
|
|
||||||
|
async def describe(self, url: str) -> DescribeResponse:
|
||||||
|
self.calls.append(("describe", url, None, False, False))
|
||||||
|
return DescribeResponse(description=self._describe_text, url=url)
|
||||||
|
|
||||||
|
|
||||||
|
class TestEngineClosureDetection(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_run_closes_after_single_round_when_nothing_new(self) -> None:
|
||||||
|
client = FakeResearchClient(web_results=[], chat_text="", describe_text="")
|
||||||
|
engine = ResearchEngine(client=client)
|
||||||
|
report = await engine.run(" deep research ")
|
||||||
|
self.assertEqual(report.subject, "deep research")
|
||||||
|
self.assertTrue(report.closed)
|
||||||
|
self.assertEqual(report.total_rounds, 1)
|
||||||
|
self.assertEqual(len(report.rounds), 1)
|
||||||
|
first = report.rounds[0]
|
||||||
|
self.assertEqual(first.number, 1)
|
||||||
|
self.assertEqual(first.items_processed, 3)
|
||||||
|
self.assertEqual(first.requests_succeeded, 3)
|
||||||
|
self.assertEqual(first.requests_failed, 0)
|
||||||
|
self.assertEqual(first.new_urls, 0)
|
||||||
|
self.assertEqual(first.new_queries, 0)
|
||||||
|
self.assertTrue(first.closed)
|
||||||
|
self.assertEqual(report.queries_issued, 1)
|
||||||
|
self.assertEqual(report.queries_enqueued, 1)
|
||||||
|
self.assertEqual(report.urls_collected, 0)
|
||||||
|
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({call[0] for call in client.calls}, {"search", "chat"})
|
||||||
|
|
||||||
|
async def test_run_discovery_rounds_then_closes(self) -> None:
|
||||||
|
client = FakeResearchClient(
|
||||||
|
web_results=[
|
||||||
|
SearchResult(
|
||||||
|
title="topic alpha",
|
||||||
|
url="https://example.com/alpha",
|
||||||
|
description="alpha details",
|
||||||
|
content="alpha body",
|
||||||
|
)
|
||||||
|
],
|
||||||
|
chat_text="",
|
||||||
|
describe_text="",
|
||||||
|
)
|
||||||
|
engine = ResearchEngine(client=client)
|
||||||
|
report = await engine.run("deep research")
|
||||||
|
self.assertTrue(report.closed)
|
||||||
|
self.assertEqual(report.total_rounds, 2)
|
||||||
|
first = report.rounds[0]
|
||||||
|
self.assertEqual(first.new_urls, 1)
|
||||||
|
self.assertEqual(first.new_queries, 2)
|
||||||
|
self.assertEqual(first.new_contents, 1)
|
||||||
|
self.assertFalse(first.closed)
|
||||||
|
second = report.rounds[1]
|
||||||
|
self.assertEqual(second.new_urls, 0)
|
||||||
|
self.assertEqual(second.new_queries, 0)
|
||||||
|
self.assertTrue(second.closed)
|
||||||
|
self.assertEqual(report.queries_generated, 7)
|
||||||
|
self.assertEqual(report.queries_enqueued, 3)
|
||||||
|
self.assertEqual(report.queries_issued, 3)
|
||||||
|
self.assertEqual(report.queries_duplicates_skipped, 4)
|
||||||
|
self.assertEqual(report.urls_collected, 1)
|
||||||
|
self.assertEqual(report.urls_duplicates_skipped, 2)
|
||||||
|
self.assertEqual(report.contents_seen, 1)
|
||||||
|
self.assertEqual(report.content_duplicates_skipped, 2)
|
||||||
|
self.assertEqual(report.requests_succeeded, 10)
|
||||||
|
self.assertEqual(report.requests_failed, 0)
|
||||||
|
self.assertEqual(report.cache_hits, 0)
|
||||||
|
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, _, _, _, _ 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]:
|
||||||
|
return [
|
||||||
|
SearchResult(
|
||||||
|
title="dup title",
|
||||||
|
url="https://example.com/dup",
|
||||||
|
description="dup details",
|
||||||
|
content=f"body for {query}",
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
client = FakeResearchClient(web_result_factory=factory, chat_text="", describe_text="")
|
||||||
|
engine = ResearchEngine(client=client)
|
||||||
|
report = await engine.run("subject")
|
||||||
|
self.assertTrue(report.closed)
|
||||||
|
self.assertEqual(report.total_rounds, 2)
|
||||||
|
first = report.rounds[0]
|
||||||
|
self.assertEqual(first.new_urls, 1)
|
||||||
|
self.assertEqual(first.new_queries, 2)
|
||||||
|
self.assertFalse(first.closed)
|
||||||
|
second = report.rounds[1]
|
||||||
|
self.assertEqual(second.new_urls, 0)
|
||||||
|
self.assertEqual(second.new_queries, 0)
|
||||||
|
self.assertEqual(second.new_contents, 2)
|
||||||
|
self.assertTrue(second.closed)
|
||||||
|
self.assertEqual(report.contents_seen, 3)
|
||||||
|
|
||||||
|
async def test_round_summary_dict_is_serialisable(self) -> None:
|
||||||
|
client = FakeResearchClient(web_results=[], chat_text="", describe_text="")
|
||||||
|
engine = ResearchEngine(client=client)
|
||||||
|
report = await engine.run("serialisable subject")
|
||||||
|
summary_dict = report.rounds[0].to_dict()
|
||||||
|
self.assertEqual(summary_dict["number"], 1)
|
||||||
|
self.assertTrue(summary_dict["closed"])
|
||||||
|
report_dict = report.to_dict()
|
||||||
|
self.assertEqual(report_dict["subject"], "serialisable subject")
|
||||||
|
self.assertEqual(report_dict["total_rounds"], 1)
|
||||||
|
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:
|
||||||
|
client = FakeResearchClient(web_results=[], chat_text="", describe_text="")
|
||||||
|
engine = ResearchEngine(client=client)
|
||||||
|
with self.assertRaises(ValueError) as ctx:
|
||||||
|
await engine.run(" \n\t ")
|
||||||
|
self.assertEqual(str(ctx.exception), "research subject must not be empty")
|
||||||
|
self.assertEqual(client.calls, [])
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
112
tests/test_research_integration.py
Normal file
112
tests/test_research_integration.py
Normal file
@ -0,0 +1,112 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import unittest
|
||||||
|
from typing import Any, AsyncIterator
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
from typosaurus_sandbox.research.client import RsearchClient
|
||||||
|
from typosaurus_sandbox.research.config import ResearchConfig
|
||||||
|
from typosaurus_sandbox.research.frontier import QueryFrontier
|
||||||
|
from typosaurus_sandbox.research.pipeline import PipelineReport, ResearchPipeline, WorkItem
|
||||||
|
|
||||||
|
PROBE_SUBJECT = "python asyncio"
|
||||||
|
RSEARCH_BASE_URL = "https://rsearch.app.molodetz.nl"
|
||||||
|
RUN_TIMEOUT_SECONDS = 60.0
|
||||||
|
|
||||||
|
SEARCH_FIXTURE: dict[str, Any] = {
|
||||||
|
"query": PROBE_SUBJECT,
|
||||||
|
"source": "duckduckgo",
|
||||||
|
"count": 2,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [
|
||||||
|
{
|
||||||
|
"title": "asyncio documentation",
|
||||||
|
"url": "https://docs.python.org/3/library/asyncio.html",
|
||||||
|
"description": "Asynchronous I/O event loop.",
|
||||||
|
"source": "docs.python.org",
|
||||||
|
"extra": {},
|
||||||
|
"index": 0,
|
||||||
|
"content": "The asyncio module provides infrastructure for writing single-threaded concurrent code.",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"title": "asyncio in Python",
|
||||||
|
"url": "https://example.com/asyncio",
|
||||||
|
"description": "Tutorial on asyncio.",
|
||||||
|
"source": "example.com",
|
||||||
|
"extra": {},
|
||||||
|
"index": 1,
|
||||||
|
"content": "A tutorial covering the asyncio event loop and coroutines.",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeResponse:
|
||||||
|
def __init__(self, status: int, body: bytes) -> None:
|
||||||
|
self.status = status
|
||||||
|
self._body = body
|
||||||
|
|
||||||
|
def __enter__(self) -> "_FakeResponse":
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, *args: object) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
def read(self) -> bytes:
|
||||||
|
return self._body
|
||||||
|
|
||||||
|
|
||||||
|
class TestBoundedOfflineProbe(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_bounded_probe_runs_against_mocked_transport_only(self) -> None:
|
||||||
|
config = ResearchConfig(
|
||||||
|
base_url=RSEARCH_BASE_URL,
|
||||||
|
max_concurrency=2,
|
||||||
|
default_count=2,
|
||||||
|
request_timeout_seconds=30.0,
|
||||||
|
)
|
||||||
|
self.assertEqual(config.base_url, RSEARCH_BASE_URL)
|
||||||
|
client = RsearchClient(config)
|
||||||
|
frontier = QueryFrontier(PROBE_SUBJECT)
|
||||||
|
requested: list[str] = []
|
||||||
|
|
||||||
|
def fake_urlopen(request: Any, timeout: float | None = None) -> _FakeResponse:
|
||||||
|
requested.append(request.get_full_url())
|
||||||
|
return _FakeResponse(200, json.dumps(SEARCH_FIXTURE).encode())
|
||||||
|
|
||||||
|
with mock.patch("urllib.request.urlopen", side_effect=fake_urlopen):
|
||||||
|
first = asyncio.run(self._bounded_run(client, frontier))
|
||||||
|
first_request_count = len(requested)
|
||||||
|
second = asyncio.run(self._bounded_run(client, frontier))
|
||||||
|
second_request_count = len(requested)
|
||||||
|
|
||||||
|
self.assertGreaterEqual(first.requests_succeeded, 1)
|
||||||
|
self.assertGreaterEqual(first.urls_found, 1)
|
||||||
|
self.assertGreaterEqual(first.contents_seen, 1)
|
||||||
|
self.assertFalse(any(outcome.cache_hit for outcome in first.outcomes))
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertGreaterEqual(stats.urls_seen, 1)
|
||||||
|
self.assertGreaterEqual(stats.content_seen, 1)
|
||||||
|
self.assertGreaterEqual(first_request_count, 1)
|
||||||
|
for url in requested:
|
||||||
|
self.assertTrue(url.startswith(RSEARCH_BASE_URL), url)
|
||||||
|
self.assertTrue(any("/search" in url for url in requested))
|
||||||
|
self.assertEqual(second.requests_succeeded, 1)
|
||||||
|
self.assertTrue(any(outcome.cache_hit for outcome in second.outcomes))
|
||||||
|
self.assertEqual(second_request_count, first_request_count)
|
||||||
|
|
||||||
|
async def _bounded_run(self, client: RsearchClient, frontier: QueryFrontier) -> PipelineReport:
|
||||||
|
async def items() -> AsyncIterator[WorkItem]:
|
||||||
|
yield WorkItem("web", PROBE_SUBJECT)
|
||||||
|
|
||||||
|
pipeline = ResearchPipeline(client, frontier)
|
||||||
|
return await asyncio.wait_for(pipeline.run(items()), timeout=RUN_TIMEOUT_SECONDS)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
|
|
||||||
|
|
||||||
356
tests/test_research_pipeline.py
Normal file
356
tests/test_research_pipeline.py
Normal file
@ -0,0 +1,356 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
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()
|
||||||
|
|
||||||
519
tests/test_research_scheduling.py
Normal file
519
tests/test_research_scheduling.py
Normal file
@ -0,0 +1,519 @@
|
|||||||
|
# retoor <retoor@molodetz.nl>
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import unittest
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
from typing import Any, AsyncIterator
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
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 ChatResponse, DescribeResponse, SearchResponse, SearchResult
|
||||||
|
from typosaurus_sandbox.research.frontier import QueryFrontier
|
||||||
|
from typosaurus_sandbox.research.pipeline import ResearchPipeline, WorkItem
|
||||||
|
|
||||||
|
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",
|
||||||
|
}
|
||||||
|
|
||||||
|
IMAGES_FIXTURE: dict[str, Any] = {
|
||||||
|
"query": "subject",
|
||||||
|
"source": "wikimedia",
|
||||||
|
"count": 1,
|
||||||
|
"success": True,
|
||||||
|
"error": None,
|
||||||
|
"results": [
|
||||||
|
{
|
||||||
|
"title": "Aurora borealis over Norway",
|
||||||
|
"url": "https://example.com/image",
|
||||||
|
"description": "Photograph of the aurora borealis.",
|
||||||
|
"source": "wikimedia",
|
||||||
|
"extra": {},
|
||||||
|
"index": 0,
|
||||||
|
}
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeClient:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.config = ResearchConfig()
|
||||||
|
self.search_calls: list[tuple[str, dict[str, Any]]] = []
|
||||||
|
self.chat_calls: list[str] = []
|
||||||
|
self.describe_calls: list[str] = []
|
||||||
|
self.delay_seconds = 0.0
|
||||||
|
self.active = 0
|
||||||
|
self.peak = 0
|
||||||
|
|
||||||
|
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,
|
||||||
|
) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
def describe_cached(self, url: str) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
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.search_calls.append((query, {"type": type, "content": content, "count": count}))
|
||||||
|
self.active += 1
|
||||||
|
self.peak = max(self.peak, self.active)
|
||||||
|
try:
|
||||||
|
if self.delay_seconds:
|
||||||
|
await asyncio.sleep(self.delay_seconds)
|
||||||
|
if type == "images":
|
||||||
|
return SearchResponse.from_dict(IMAGES_FIXTURE)
|
||||||
|
return SearchResponse.from_dict(SEARCH_FIXTURE)
|
||||||
|
finally:
|
||||||
|
self.active -= 1
|
||||||
|
|
||||||
|
async def chat(self, prompt: str, *, system: str | None = None, json_mode: bool = False, cache: bool = True) -> ChatResponse:
|
||||||
|
self.chat_calls.append(prompt)
|
||||||
|
return ChatResponse.from_dict(CHAT_FIXTURE)
|
||||||
|
|
||||||
|
async def describe(self, url: str) -> DescribeResponse:
|
||||||
|
self.describe_calls.append(url)
|
||||||
|
return DescribeResponse.from_dict(DESCRIBE_FIXTURE)
|
||||||
|
|
||||||
|
|
||||||
|
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}"))
|
||||||
|
|
||||||
|
def test_entry_expires_after_ttl_elapses(self) -> None:
|
||||||
|
cache = TTLCache[str]("expiry", ttl_seconds=10.0)
|
||||||
|
with mock.patch("typosaurus_sandbox.research.cache.time.monotonic", side_effect=[100.0, 100.0, 111.0]):
|
||||||
|
cache.set("key", "value")
|
||||||
|
self.assertEqual(cache.get("key"), "value")
|
||||||
|
self.assertIsNone(cache.get("key"))
|
||||||
|
|
||||||
|
def test_entry_expires_exactly_at_ttl_boundary(self) -> None:
|
||||||
|
cache = TTLCache[str]("boundary", ttl_seconds=10.0)
|
||||||
|
with mock.patch("typosaurus_sandbox.research.cache.time.monotonic", side_effect=[100.0, 100.0, 110.0]):
|
||||||
|
cache.set("key", "value")
|
||||||
|
self.assertEqual(cache.get("key"), "value")
|
||||||
|
self.assertIsNone(cache.get("key"))
|
||||||
|
|
||||||
|
|
||||||
|
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")
|
||||||
|
|
||||||
|
|
||||||
|
class TestResearchPipeline(unittest.IsolatedAsyncioTestCase):
|
||||||
|
|
||||||
|
async def test_pipeline_pool_size_bounded_by_configured_concurrency(self) -> None:
|
||||||
|
client = _FakeClient()
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
default_pipeline = ResearchPipeline(client, frontier)
|
||||||
|
self.assertEqual(default_pipeline.pool_size, 8)
|
||||||
|
self.assertEqual(default_pipeline.pool_size, ResearchConfig().max_concurrency)
|
||||||
|
narrow_pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=3))
|
||||||
|
self.assertEqual(narrow_pipeline.pool_size, 3)
|
||||||
|
floor_pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=0))
|
||||||
|
self.assertEqual(floor_pipeline.pool_size, 1)
|
||||||
|
|
||||||
|
async def test_pipeline_drains_all_work_items_and_dedups(self) -> None:
|
||||||
|
client = _FakeClient()
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
|
||||||
|
async def items() -> AsyncIterator[WorkItem]:
|
||||||
|
yield WorkItem("web", "subject")
|
||||||
|
yield WorkItem("web", "subject")
|
||||||
|
yield WorkItem("images", "subject")
|
||||||
|
yield WorkItem("describe", "https://example.com/page")
|
||||||
|
yield WorkItem("chat", "prompt")
|
||||||
|
|
||||||
|
pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=4))
|
||||||
|
report = await pipeline.run(items())
|
||||||
|
self.assertEqual(report.requests_succeeded, 5)
|
||||||
|
self.assertEqual(report.requests_failed, 0)
|
||||||
|
self.assertEqual(len(report.outcomes), 5)
|
||||||
|
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"})
|
||||||
|
self.assertEqual(len(client.search_calls), 3)
|
||||||
|
self.assertEqual(len(client.describe_calls), 1)
|
||||||
|
self.assertEqual(len(client.chat_calls), 1)
|
||||||
|
stats = frontier.snapshot()
|
||||||
|
self.assertEqual(stats.urls_seen, 2)
|
||||||
|
self.assertEqual(stats.urls_duplicates_skipped, 1)
|
||||||
|
self.assertEqual(stats.queries_enqueued, 6)
|
||||||
|
self.assertEqual(stats.queries_duplicates_skipped, 2)
|
||||||
|
self.assertEqual(stats.content_seen, 2)
|
||||||
|
|
||||||
|
async def test_pipeline_concurrency_bounded_by_pool_size(self) -> None:
|
||||||
|
client = _FakeClient()
|
||||||
|
client.delay_seconds = 0.02
|
||||||
|
frontier = QueryFrontier()
|
||||||
|
|
||||||
|
async def items() -> AsyncIterator[WorkItem]:
|
||||||
|
for i in range(12):
|
||||||
|
yield WorkItem("web", f"subject {i}")
|
||||||
|
|
||||||
|
pipeline = ResearchPipeline(client, frontier, ResearchConfig(max_concurrency=4))
|
||||||
|
report = await pipeline.run(items())
|
||||||
|
self.assertEqual(report.requests_succeeded, 12)
|
||||||
|
self.assertGreaterEqual(client.peak, 2)
|
||||||
|
self.assertLessEqual(client.peak, 4)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
Loading…
Reference in New Issue
Block a user