websearch-kit 0.1.0__py3-none-any.whl

websearch_kit/__init__.py

@@ -0,0 +1,100 @@
+"""websearch-kit — web search, fetch, and research pipeline for LLMs.
+
+Public API surface (everything re-exported here is SemVer-protected; see
+VERSIONING.md). Usable three ways:
+
+* SDK:        ``from websearch_kit import SearchKit``
+* MCP server: ``websearch-kit-mcp`` (requires the ``[mcp]`` extra)
+* Open WebUI: the single-file filter under ``adapters/owui/``
+"""
+
+from typing import TYPE_CHECKING, Any
+
+from ._version import __version__
+from .config import WebSearchConfig
+from .errors import (
+    CacheError,
+    ConfigError,
+    DeadlineExceededError,
+    ExpansionError,
+    ExtractionError,
+    FetchError,
+    GuardError,
+    MissingDependencyError,
+    ProviderError,
+    RobotsBlockedError,
+    SSRFBlockedError,
+    WebSearchKitError,
+)
+from .models import (
+    Degradation,
+    FetchOutcome,
+    PageContent,
+    ProgressEvent,
+    ResearchReport,
+    RunStats,
+    SearchResult,
+    Source,
+    Stage,
+    SystemHealth,
+)
+from .protocols import (
+    Cache,
+    CallbackSink,
+    ProgressSink,
+    ProviderCapabilities,
+    QueryExpander,
+    SearchProvider,
+)
+
+if TYPE_CHECKING:  # pragma: no cover - import surface for type checkers only
+    from .kit import SearchKit, SyncSearchKit
+
+__all__ = [
+    "Cache",
+    "CacheError",
+    "CallbackSink",
+    "ConfigError",
+    "DeadlineExceededError",
+    "Degradation",
+    "ExpansionError",
+    "ExtractionError",
+    "FetchError",
+    "FetchOutcome",
+    "GuardError",
+    "MissingDependencyError",
+    "PageContent",
+    "ProgressEvent",
+    "ProgressSink",
+    "ProviderCapabilities",
+    "ProviderError",
+    "QueryExpander",
+    "ResearchReport",
+    "RobotsBlockedError",
+    "RunStats",
+    "SSRFBlockedError",
+    "SearchKit",
+    "SearchProvider",
+    "SearchResult",
+    "Source",
+    "Stage",
+    "SyncSearchKit",
+    "SystemHealth",
+    "WebSearchConfig",
+    "WebSearchKitError",
+    "__version__",
+]
+
+#: Engine names resolved lazily: importing ``kit`` pulls the extraction chain,
+#: whose package import loads trafilatura (heavyweight). Deferring keeps
+#: ``import websearch_kit`` cheap for dep-light consumers (models, config,
+#: sanitize) while ``from websearch_kit import SearchKit`` still just works.
+_LAZY_EXPORTS = frozenset({"SearchKit", "SyncSearchKit"})
+
+
+def __getattr__(name: str) -> Any:
+    if name in _LAZY_EXPORTS:
+        from . import kit
+
+        return getattr(kit, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

websearch_kit/_version.py

@@ -0,0 +1,3 @@
+"""Single source of version truth (read by hatchling and exported from __init__)."""
+
+__version__ = "0.1.0"

websearch_kit/assembly/__init__.py

@@ -0,0 +1,33 @@
+"""Context assembly: numbered source blocks and their 1:1 citation list.
+
+WHY a package: the pipeline's final products — the ``[N]``-numbered context
+string handed to an LLM and the ``Source`` citations a UI maps those markers
+back to — must be generated from the same records in the same order or the
+model cites things the UI cannot resolve. ``citations`` owns the numbering,
+``context_builder`` owns the (reference-verbatim) block formats, and both are
+pure functions over :class:`SourceDraft` so the lockstep invariant is golden-
+testable with no pipeline machinery.
+"""
+
+from __future__ import annotations
+
+from .citations import SourceDraft, number_sources
+from .context_builder import (
+    SNIPPET_POOL_HEADER,
+    TRUNCATION_SUFFIX,
+    build_context,
+    render_pool_block,
+    render_source_block,
+    truncate_to,
+)
+
+__all__ = [
+    "SNIPPET_POOL_HEADER",
+    "TRUNCATION_SUFFIX",
+    "SourceDraft",
+    "build_context",
+    "number_sources",
+    "render_pool_block",
+    "render_source_block",
+    "truncate_to",
+]

websearch_kit/assembly/citations.py

@@ -0,0 +1,70 @@
+"""Citation numbering: assembled source records -> public ``Source`` models.
+
+WHY a dedicated module: the context block and the citation list MUST stay in
+lockstep — every ``[N]`` marker the LLM can emit needs exactly one
+:class:`~websearch_kit.models.Source` with the same ``n``, and the numbering
+must be contiguous across the two segments (fetched sources first, then the
+snippet pool). Splitting "what goes in" (the pipeline's ranked
+:class:`SourceDraft` records) from "how it is numbered" (here) and "how it is
+rendered" (``context_builder``) keeps that 1:1 invariant testable as a pure
+function: same drafts in, same numbering out, no pipeline state involved.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+from ..models import Source
+
+__all__ = ["SourceDraft", "number_sources"]
+
+
+@dataclass(slots=True)
+class SourceDraft:
+    """One assembled source before numbering: the pipeline's working record.
+
+    Mutable on purpose — the ranking/budget stage truncates ``content`` and
+    assigns ``score`` in place. ``kind`` records *what the content is*
+    (``snippet_only`` when the search snippet substituted for a failed or
+    low-quality fetch), independent of *where* the draft renders (primary
+    block vs snippet pool) — that placement is positional, decided by which
+    sequence the draft is passed in.
+    """
+
+    title: str
+    url: str
+    snippet: str
+    content: str
+    kind: Literal["fetched", "snippet_only"]
+    score: float | None = None
+
+
+def number_sources(
+    primary: Sequence[SourceDraft],
+    pool: Sequence[SourceDraft],
+) -> list[Source]:
+    """Assign contiguous 1-based ``[N]`` numbers across both segments.
+
+    ``primary`` (fetched/snippet-fallback sources, already in ranked order)
+    numbers first; ``pool`` (relevance-filtered snippet-only extras) continues
+    the sequence — exactly the reference's ``source_id`` counter that ran
+    uninterrupted from the fetched blocks into the additional-sources section.
+    The returned list is 1:1 with the ``[N]`` markers ``build_context`` renders
+    for the same two sequences.
+    """
+    sources: list[Source] = []
+    for n, draft in enumerate((*primary, *pool), start=1):
+        sources.append(
+            Source(
+                n=n,
+                title=draft.title,
+                url=draft.url,
+                snippet=draft.snippet,
+                kind=draft.kind,
+                score=draft.score,
+                content_chars=len(draft.content),
+            )
+        )
+    return sources

websearch_kit/assembly/context_builder.py

@@ -0,0 +1,90 @@
+"""Context-block rendering: numbered ``--- [N] Title ---`` segments for the LLM.
+
+WHY the formats are frozen: the answer prompt (``prompts.build_answer_prompt``)
+instructs the model to cite with inline ``[N]`` markers and to fall back to the
+``Summary (Snippet)`` line when ``Full Content`` is poor — so the literal field
+labels in these blocks are part of the prompt contract, not cosmetics. All
+three templates (fetched block, snippet-pool block, pool header) are ported
+verbatim from the reference ``_process_results`` Phase C so the golden tests
+pin byte-for-byte parity. Everything here is a pure string function; numbering
+stays in lockstep with ``citations.number_sources`` because both iterate the
+same ``(primary, pool)`` sequences in the same order.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from .citations import SourceDraft
+
+__all__ = [
+    "SNIPPET_POOL_HEADER",
+    "TRUNCATION_SUFFIX",
+    "build_context",
+    "render_pool_block",
+    "render_source_block",
+    "truncate_to",
+]
+
+#: Appended when a source's content is cut to its budget allocation (verbatim
+#: reference suffix — the visible, in-context signal that text was dropped).
+TRUNCATION_SUFFIX = "... [TRUNCATED]"
+
+#: Separator announcing the snippet-only pool segment (verbatim reference).
+SNIPPET_POOL_HEADER = "\n--- ADDITIONAL SOURCES (snippet only, same [N] citation format) ---"
+
+
+def render_source_block(n: int, draft: SourceDraft) -> str:
+    """One primary source block (fetched or snippet-fallback content)."""
+    return (
+        f"--- [{n}] {draft.title} ---\n"
+        f"URL: {draft.url}\n"
+        f"Summary (Snippet): {draft.snippet}\n"
+        f"Full Content:\n{draft.content}\n"
+    )
+
+
+def render_pool_block(n: int, draft: SourceDraft) -> str:
+    """One snippet-pool block (never fetched; the snippet *is* the content)."""
+    return (
+        f"--- [{n}] {draft.title} (snippet only) ---\nURL: {draft.url}\nContent: {draft.snippet}\n"
+    )
+
+
+def truncate_to(content: str, alloc: int) -> tuple[str, bool]:
+    """Cut ``content`` to ``alloc`` chars, marking the cut with the suffix.
+
+    Returns ``(text, truncated)``. The suffix is appended *after* the cut
+    (reference behavior: ``content[:alloc] + "... [TRUNCATED]"``), so the
+    rendered block may slightly exceed ``alloc`` — the budget governs content,
+    the marker is overhead the model needs to see.
+    """
+    if len(content) <= alloc:
+        return content, False
+    return content[:alloc] + TRUNCATION_SUFFIX, True
+
+
+def build_context(
+    primary: Sequence[SourceDraft],
+    pool: Sequence[SourceDraft],
+) -> str:
+    """Assemble the full ``<search_results>`` payload from both segments.
+
+    Primary blocks render first (ranked order preserved), then — only when the
+    pool is non-empty — the pool header and the snippet-only blocks, with
+    ``[N]`` numbering running contiguously across the boundary. Blocks are
+    joined with ``"\\n"`` (each block already ends in a newline, yielding the
+    reference's blank-line separation). Empty input renders an empty string —
+    the caller decides what a sourceless run means.
+    """
+    parts: list[str] = []
+    n = 1
+    for draft in primary:
+        parts.append(render_source_block(n, draft))
+        n += 1
+    if pool:
+        parts.append(SNIPPET_POOL_HEADER)
+        for draft in pool:
+            parts.append(render_pool_block(n, draft))
+            n += 1
+    return "\n".join(parts)

websearch_kit/caching/__init__.py

@@ -0,0 +1,89 @@
+"""Caching subsystem: backends, key builders, and the :func:`make_cache` factory.
+
+WHY this package: the pipeline caches three things — provider searches, fetched
+pages, and query expansions — each with its own TTL (see ``WebSearchConfig``).
+All three flow through the single :class:`~websearch_kit.protocols.Cache`
+protocol, so the choice of backend is a configuration detail resolved once, here,
+by :func:`make_cache`. Centralizing construction keeps backend wiring (the sqlite
+directory resolution, the fail-soft :class:`CacheGuard` wrapping) out of the
+engine.
+
+The factory always wraps a real backend in :class:`CacheGuard` so a backend fault
+degrades to a cache miss instead of failing the run (no-fail-silent: the guard
+logs and counts every swallowed error). ``cache_backend="none"`` returns ``None``
+— the engine treats a ``None`` cache as "caching disabled" and skips lookups
+entirely, which is cheaper than routing every call through a no-op object.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from collections.abc import Callable
+from pathlib import Path
+
+from ..config import WebSearchConfig
+from ..protocols import Cache
+from .keys import content_key, expansion_key, search_key
+from .memory import MemoryTTLCache
+from .sqlite_cache import CacheGuard, SqliteCache
+
+__all__ = [
+    "CacheGuard",
+    "MemoryTTLCache",
+    "SqliteCache",
+    "content_key",
+    "default_cache_dir",
+    "expansion_key",
+    "make_cache",
+    "search_key",
+]
+
+
+def default_cache_dir() -> Path:
+    """Resolve the default sqlite-cache directory using the XDG base-dir spec.
+
+    Honors ``$XDG_CACHE_HOME`` when set to an absolute path (per the spec, a
+    relative or empty value is ignored), otherwise falls back to
+    ``~/.cache``; the websearch-kit cache lives in a ``websearch-kit``
+    subdirectory of that base. Implemented with the standard library only
+    (``os``/``pathlib``) so caching never drags in ``platformdirs``.
+    """
+    xdg_cache = os.environ.get("XDG_CACHE_HOME")
+    # The XDG spec mandates absolute paths; ignore a relative/empty value.
+    if xdg_cache and os.path.isabs(xdg_cache):
+        base = Path(xdg_cache)
+    else:
+        base = Path.home() / ".cache"
+    return base / "websearch-kit"
+
+
+def make_cache(
+    config: WebSearchConfig,
+    clock: Callable[[], float] = time.monotonic,
+) -> Cache | None:
+    """Build the configured cache backend, or ``None`` when caching is disabled.
+
+    * ``"memory"`` -> :class:`MemoryTTLCache` (process-local, default).
+    * ``"sqlite"`` -> :class:`SqliteCache` rooted at ``config.cache_dir`` (or
+      :func:`default_cache_dir` when unset).
+    * ``"none"``   -> ``None`` (the engine skips caching entirely).
+
+    Any real backend is wrapped in :class:`CacheGuard` so backend faults become
+    logged, counted cache misses rather than run failures.
+
+    The ``clock`` is threaded into the memory backend (whose TTL uses a monotonic
+    source) so tests can drive expiry deterministically. The sqlite backend keeps
+    its own wall-clock default because its TTLs must survive process restarts.
+    """
+    backend = config.cache_backend
+    if backend == "none":
+        return None
+    if backend == "memory":
+        return CacheGuard(MemoryTTLCache(max_entries=512, clock=clock))
+    if backend == "sqlite":
+        directory = Path(config.cache_dir) if config.cache_dir else default_cache_dir()
+        return CacheGuard(SqliteCache(directory=directory))
+    # CacheBackend is a closed Literal; an unhandled member is a programming
+    # error, not a runtime input — surface it loudly rather than fail silent.
+    raise AssertionError(f"unhandled cache backend: {backend!r}")  # pragma: no cover

websearch_kit/caching/keys.py

@@ -0,0 +1,93 @@
+"""Deterministic cache-key builders for the three cacheable pipeline stages.
+
+WHY a dedicated module: cache correctness hinges on *stable, collision-resistant*
+keys derived from exactly the inputs that change a result — and nothing else. A
+key built ad-hoc at each call site drifts (parameter order, optional-None
+handling, forgetting to normalize a URL) and silently poisons the cache with
+stale or cross-contaminated entries. Centralizing the builders makes the keyspace
+auditable and the canonicalization rules a single source of truth.
+
+Design choices:
+
+* **sha256 hex of a canonical string.** Keys are fixed-length, opaque, and safe
+  to use as a sqlite ``TEXT PRIMARY KEY`` regardless of how long/odd the inputs
+  are. sha256 is not used here for any security property — only for a uniform,
+  low-collision digest — so the lint that flags weak hashes does not apply.
+* **Explicit stage prefixes** (``search:``/``content:``/``expand:``) keep the
+  three keyspaces disjoint inside a shared backend, so a content URL can never
+  alias a search query that happens to hash the same canonical bytes.
+* **``content_key`` applies :func:`sanitize_url` itself.** Stripping tracking
+  parameters (utm_*, gclid, ...) before hashing is what makes two links to the
+  same page share a cache entry. Doing it *inside* the key builder means a caller
+  physically cannot forget to normalize first — the dedup win is structural, not
+  a convention someone must remember.
+"""
+
+from __future__ import annotations
+
+import hashlib
+
+from ..security.sanitize import sanitize_url
+
+__all__ = [
+    "content_key",
+    "expansion_key",
+    "search_key",
+]
+
+#: Field separator for the canonical string. A control character (unit
+#: separator, 0x1f) cannot appear in a normal query/URL/prompt, so it cannot be
+#: forged by input to merge two distinct tuples into one canonical string.
+_SEP = "\x1f"
+
+
+def _digest(prefix: str, *parts: str) -> str:
+    """Return ``"{prefix}{sha256-hex}"`` over the unit-separator-joined parts."""
+    canonical = _SEP.join(parts)
+    hexdigest = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
+    return f"{prefix}{hexdigest}"
+
+
+def search_key(
+    provider: str,
+    query: str,
+    count: int,
+    lang: str | None,
+    time_range: str | None,
+) -> str:
+    """Key for one provider search call.
+
+    Every argument that changes the returned hits participates: the provider
+    (different engines, different results), the query, the requested ``count``,
+    and the optional language / time-range filters. ``None`` is canonicalized to
+    the empty string so an unset filter is stable and distinct from any real
+    value a provider would accept.
+    """
+    return _digest(
+        "search:",
+        provider,
+        query,
+        str(count),
+        lang or "",
+        time_range or "",
+    )
+
+
+def content_key(sanitized_url: str) -> str:
+    """Key for one fetched+extracted page, keyed on the *normalized* URL.
+
+    The argument name documents intent, but this builder does not trust the
+    caller to have normalized: it runs :func:`sanitize_url` here so two URLs that
+    differ only by tracking parameters (``?utm_source=...``) collapse to the same
+    cache entry. This is deliberate — the strip-before-hash cannot be skipped.
+    """
+    return _digest("content:", sanitize_url(sanitized_url))
+
+
+def expansion_key(prompt_text: str) -> str:
+    """Key for a query-expansion result, keyed on the full prompt text.
+
+    The expansion output is a pure function of the prompt handed to the LLM, so
+    the entire prompt string is the cache identity.
+    """
+    return _digest("expand:", prompt_text)

websearch_kit/caching/memory.py

@@ -0,0 +1,121 @@
+"""In-process TTL cache with LRU eviction — the zero-config default backend.
+
+WHY in-memory by default: the common deployment is a single long-lived process
+(an OWUI worker, an MCP server) where a process-local cache eliminates duplicate
+provider calls and page fetches within a session at zero operational cost — no
+file, no extra dependency. It is intentionally *not* shared across processes;
+callers needing cross-process persistence opt into :class:`SqliteCache`.
+
+Concurrency model: every mutation (and the lazy-expiry bookkeeping on ``get``)
+runs under a single :class:`asyncio.Lock`. The cache is only safe within one
+event loop — which is the contract, since the pipeline is async and single-loop.
+The lock makes the read-modify-write on the backing :class:`OrderedDict` atomic
+with respect to other awaiting tasks, so concurrent ``get``/``set`` from
+``asyncio.gather`` can never observe or produce a half-updated structure.
+
+Eviction policy:
+
+* **TTL** — each entry carries an absolute monotonic expiry timestamp. Expired
+  entries are dropped *lazily* on ``get`` (the cheap, always-correct path) and
+  *opportunistically* swept on ``set`` only when the cache is over capacity
+  (amortizing the sweep cost instead of scanning on every write).
+* **LRU** — recency is tracked by position in the ``OrderedDict``: a hit or a
+  write moves the key to the most-recently-used end. When the map exceeds
+  ``max_entries`` after a write, the least-recently-used keys are popped from the
+  front until back within bounds.
+
+The clock is injectable (default :func:`time.monotonic`) so tests drive TTL
+expiry deterministically without sleeping. Monotonic time is used because TTL is
+a duration; it must be immune to wall-clock jumps (NTP steps, DST).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+__all__ = ["MemoryTTLCache"]
+
+
+@dataclass(slots=True)
+class _Entry:
+    """A cached value plus its absolute monotonic expiry timestamp."""
+
+    value: Any
+    expires_at: float
+
+
+class MemoryTTLCache:
+    """Async, LRU-bounded, per-entry-TTL cache implementing :class:`protocols.Cache`.
+
+    Args:
+        max_entries: Hard cap on retained entries; LRU keys are evicted past it.
+        clock: Monotonic time source returning seconds; injected for tests.
+    """
+
+    def __init__(
+        self,
+        max_entries: int = 512,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self._max_entries = max_entries
+        self._clock = clock
+        # Insertion/access order IS the LRU order: front = least recently used.
+        self._store: OrderedDict[str, _Entry] = OrderedDict()
+        self._lock = asyncio.Lock()
+
+    async def get(self, key: str) -> Any | None:
+        """Return the live value for ``key``, or ``None`` if absent or expired.
+
+        A live hit is promoted to most-recently-used. An expired entry is dropped
+        lazily here (and reported as a miss) so stale data never escapes, even if
+        no write has triggered an opportunistic sweep.
+        """
+        async with self._lock:
+            entry = self._store.get(key)
+            if entry is None:
+                return None
+            if self._is_expired(entry):
+                # Lazy expiry: drop the stale row and report a miss. Not
+                # fail-silent — an expired entry is definitionally not present.
+                del self._store[key]
+                return None
+            self._store.move_to_end(key)
+            return entry.value
+
+    async def set(self, key: str, value: Any, ttl: float) -> None:
+        """Insert/replace ``key`` with ``value``, expiring ``ttl`` seconds hence.
+
+        The key becomes most-recently-used. When the store is over capacity after
+        the write, expired entries are swept first (reclaiming space without
+        discarding live data), then the least-recently-used keys are evicted until
+        back within ``max_entries``.
+        """
+        async with self._lock:
+            expires_at = self._clock() + ttl
+            if key in self._store:
+                self._store.move_to_end(key)
+            self._store[key] = _Entry(value=value, expires_at=expires_at)
+            if len(self._store) > self._max_entries:
+                self._sweep_expired()
+                self._evict_lru()
+
+    def _is_expired(self, entry: _Entry) -> bool:
+        return entry.expires_at <= self._clock()
+
+    def _sweep_expired(self) -> None:
+        """Drop every currently-expired entry. Caller must hold the lock."""
+        now = self._clock()
+        stale = [key for key, entry in self._store.items() if entry.expires_at <= now]
+        for key in stale:
+            del self._store[key]
+
+    def _evict_lru(self) -> None:
+        """Pop least-recently-used keys until within capacity. Caller holds lock."""
+        while len(self._store) > self._max_entries:
+            # popitem(last=False) removes the front == least-recently-used entry.
+            self._store.popitem(last=False)