src/typosaurus_sandbox/research/__init__.py

@@ -0,0 +1,29 @@
+# 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.envelopes import (
+    ChatResponse,
+    ChatUsage,
+    DeepReport,
+    DescribeResponse,
+    SearchGrade,
+    SearchResponse,
+    SearchResult,
+)
+
+__all__ = [
+    "ChatResponse",
+    "ChatUsage",
+    "DeepReport",
+    "DescribeResponse",
+    "RsearchClient",
+    "RsearchError",
+    "ResearchConfig",
+    "SearchGrade",
+    "SearchResponse",
+    "SearchResult",
+    "TTLCache",
+]
+

src/typosaurus_sandbox/research/cache.py

@@ -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)
+

src/typosaurus_sandbox/research/client.py

@@ -0,0 +1,212 @@
+# 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()
+
+
+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: 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"
+        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
+
+    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
+

src/typosaurus_sandbox/research/config.py

@@ -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),
+        )
+

src/typosaurus_sandbox/research/envelopes.py

@@ -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"),
+        )
+