refindery 0.1.0__py3-none-any.whl

refindery/__init__.py

@@ -0,0 +1,3 @@
+"""Refindery: a local, single-machine retrieval engine over the web pages you read."""
+
+__version__ = "0.1.0"

refindery/__main__.py

@@ -0,0 +1,6 @@
+"""Entry point: ``python -m refindery`` dispatches to the CLI (default: serve)."""
+
+from refindery.cli import main
+
+if __name__ == "__main__":
+    main()

refindery/adapters/__init__.py

@@ -0,0 +1 @@
+"""Adapters: concrete implementations of the application ports."""

refindery/adapters/chunking/__init__.py

@@ -0,0 +1 @@
+"""Canonical chunking adapter (chonkie)."""

refindery/adapters/chunking/chonkie_chunker.py

@@ -0,0 +1,73 @@
+"""Sentence-aware chunking via chonkie with a hard token ceiling.
+
+Chunking is canonical and model-independent: one chunking, every embedding
+model embeds the same spans. The cl100k tokenizer is the canonical counter.
+chonkie targets ``chunk_size`` but can exceed it on pathological input, so a
+post-pass re-splits anything above the hard max (which model registration
+guarantees is within every registered model's budget).
+"""
+
+from chonkie import SentenceChunker, TokenChunker
+
+from refindery.domain.ids import PageId, new_chunk_id
+from refindery.domain.models import Chunk
+
+_TOKENIZER = "cl100k_base"
+
+
+class ChonkieChunker:
+    """Chunker port implementation over chonkie's SentenceChunker."""
+
+    def __init__(
+        self,
+        *,
+        target_tokens: int = 448,
+        overlap_tokens: int = 64,
+        hard_max_tokens: int = 512,
+    ) -> None:
+        self._hard_max = hard_max_tokens
+        self._sentence = SentenceChunker(
+            tokenizer=_TOKENIZER,
+            chunk_size=target_tokens,
+            chunk_overlap=overlap_tokens,
+        )
+        self._splitter = TokenChunker(
+            tokenizer=_TOKENIZER,
+            chunk_size=hard_max_tokens,
+            chunk_overlap=0,
+        )
+
+    def chunk(self, *, page_id: PageId, text: str) -> list[Chunk]:
+        """Split ``text`` into ordered chunks with char offsets and token counts."""
+        if not text.strip():
+            return []
+        spans: list[tuple[str, int, int, int]] = []
+        for piece in self._sentence.chunk(text):
+            if piece.token_count <= self._hard_max:
+                spans.append(
+                    (piece.text, piece.start_index, piece.end_index, piece.token_count)
+                )
+                continue
+            spans.extend(
+                (
+                    sub.text,
+                    piece.start_index + sub.start_index,
+                    piece.start_index + sub.end_index,
+                    sub.token_count,
+                )
+                for sub in self._splitter.chunk(piece.text)
+            )
+        return [
+            Chunk(
+                id=new_chunk_id(),
+                page_id=page_id,
+                ordinal=ordinal,
+                text=chunk_text,
+                token_count=token_count,
+                char_start=char_start,
+                char_end=char_end,
+            )
+            for ordinal, (chunk_text, char_start, char_end, token_count) in enumerate(
+                spans
+            )
+        ]

refindery/adapters/clock.py

@@ -0,0 +1,11 @@
+"""System clock adapter."""
+
+from datetime import UTC, datetime
+
+
+class SystemClock:
+    """Wall clock returning tz-aware UTC datetimes."""
+
+    def now(self) -> datetime:
+        """Return the current UTC time."""
+        return datetime.now(tz=UTC)

refindery/adapters/cluster/__init__.py

@@ -0,0 +1 @@
+"""Clustering engine adapters (UMAP/PCA + HDBSCAN/KMeans in a process pool)."""

refindery/adapters/cluster/engine.py

@@ -0,0 +1,53 @@
+"""ClusterEngine adapter: runs the worker in a spawn ProcessPoolExecutor."""
+
+import asyncio
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
+from multiprocessing import get_context
+
+import numpy as np
+import numpy.typing as npt
+
+from refindery.adapters.cluster.worker import reduce_and_cluster
+from refindery.application.ports.cluster_engine import ClusterFitResult, ClusterParams
+
+
+class ProcessPoolClusterEngine:
+    """CPU-heavy clustering off the event loop."""
+
+    def __init__(self, *, max_workers: int = 1) -> None:
+        self._executor = ProcessPoolExecutor(
+            max_workers=max_workers, mp_context=get_context("spawn")
+        )
+
+    async def fit(
+        self, *, vectors: npt.NDArray[np.float32], params: ClusterParams
+    ) -> ClusterFitResult:
+        """Fit in the pool; only arrays and primitives cross the boundary."""
+        loop = asyncio.get_running_loop()
+        labels, probabilities, reduce_ms, cluster_ms = await loop.run_in_executor(
+            self._executor,
+            partial(
+                reduce_and_cluster,
+                vectors,
+                algorithm=params.algorithm,
+                reducer=params.reducer,
+                n_components=params.n_components,
+                n_neighbors=params.n_neighbors,
+                min_dist=params.min_dist,
+                min_cluster_size=params.min_cluster_size,
+                min_samples=params.min_samples,
+                leiden_resolution=params.leiden_resolution,
+                random_state=params.random_state,
+            ),
+        )
+        return ClusterFitResult(
+            labels=labels,
+            probabilities=probabilities,
+            reduce_ms=reduce_ms,
+            cluster_ms=cluster_ms,
+        )
+
+    def close(self) -> None:
+        """Shut the pool down."""
+        self._executor.shutdown(wait=False, cancel_futures=True)

refindery/adapters/cluster/worker.py

@@ -0,0 +1,208 @@
+"""Process-pool worker: everything CPU-heavy about clustering.
+
+Module-level pure function so the spawn context can pickle the reference.
+Only primitives and numpy arrays cross the boundary; umap/sklearn import
+INSIDE the function so the parent never pays the numba JIT cost. Stage
+timings are returned (spans cannot cross a process boundary).
+"""
+
+import time
+from collections.abc import Sequence
+from importlib import import_module
+from typing import Protocol, cast
+
+import numpy as np
+import numpy.typing as npt
+
+_LEIDEN_EXTRA_HINT = "Leiden clustering requires `uv sync --extra leiden`"
+
+
+class _GraphFactory(Protocol):
+    def __call__(
+        self, *, n: int, edges: list[tuple[int, int]], directed: bool
+    ) -> object:
+        """Build an igraph graph."""
+        ...
+
+
+class _IgraphModule(Protocol):
+    Graph: _GraphFactory
+
+
+class _FindPartition(Protocol):
+    def __call__(
+        self,
+        graph: object,
+        partition_type: object,
+        *,
+        weights: list[float],
+        resolution_parameter: float,
+        seed: int,
+    ) -> Sequence[Sequence[int]]:
+        """Run Leiden partitioning."""
+        ...
+
+
+class _LeidenModule(Protocol):
+    RBConfigurationVertexPartition: object
+    find_partition: _FindPartition
+
+
+def _optional_leiden_modules() -> tuple[_IgraphModule, _LeidenModule]:
+    try:
+        igraph = import_module("igraph")
+        leidenalg = import_module("leidenalg")
+    except ImportError as exc:
+        raise RuntimeError(_LEIDEN_EXTRA_HINT) from exc
+    return cast("_IgraphModule", igraph), cast("_LeidenModule", leidenalg)
+
+
+def _knn_graph_edges(
+    matrix: npt.NDArray[np.float32], *, n_neighbors: int
+) -> tuple[list[tuple[int, int]], list[float]]:
+    from sklearn.neighbors import NearestNeighbors  # noqa: PLC0415 — worker only
+
+    k = min(max(n_neighbors, 1), len(matrix) - 1)
+    model = NearestNeighbors(n_neighbors=k + 1, metric="cosine")
+    distances, indices = model.fit(matrix).kneighbors(matrix)
+
+    edge_weights: dict[tuple[int, int], float] = {}
+    for source, (row_distances, row_indices) in enumerate(
+        zip(distances, indices, strict=True)
+    ):
+        for distance, target in zip(row_distances, row_indices, strict=True):
+            target_int = int(target)
+            if source == target_int:
+                continue
+            edge = (min(source, target_int), max(source, target_int))
+            distance_value = float(distance)
+            if not np.isfinite(distance_value):
+                continue
+            weight = 1.0 / (1.0 + max(distance_value, 0.0))
+            edge_weights[edge] = max(edge_weights.get(edge, 0.0), weight)
+    edges = sorted(edge_weights)
+    return edges, [edge_weights[edge] for edge in edges]
+
+
+def _leiden_cluster(
+    matrix: npt.NDArray[np.float32],
+    *,
+    n_neighbors: int,
+    resolution: float,
+    random_state: int,
+) -> npt.NDArray[np.int64]:
+    if len(matrix) <= 1:
+        return np.zeros(len(matrix), dtype=np.int64)
+    igraph, leidenalg = _optional_leiden_modules()
+    edges, weights = _knn_graph_edges(matrix, n_neighbors=n_neighbors)
+    if not edges:
+        return np.arange(len(matrix), dtype=np.int64)
+
+    graph = igraph.Graph(n=len(matrix), edges=edges, directed=False)
+    partition = leidenalg.find_partition(
+        graph,
+        leidenalg.RBConfigurationVertexPartition,
+        weights=weights,
+        resolution_parameter=resolution,
+        seed=random_state,
+    )
+
+    labels = np.full(len(matrix), -1, dtype=np.int64)
+    for label, members in enumerate(partition):
+        labels[members] = label
+    unassigned_mask = labels < 0
+    n_unassigned = int(np.count_nonzero(unassigned_mask))
+    if n_unassigned > 0:
+        labels[unassigned_mask] = np.arange(
+            len(partition), len(partition) + n_unassigned, dtype=np.int64
+        )
+    return labels
+
+
+def reduce_and_cluster(
+    vectors: npt.NDArray[np.float32],
+    *,
+    algorithm: str,
+    reducer: str,
+    n_components: int,
+    n_neighbors: int,
+    min_dist: float,
+    min_cluster_size: int,
+    min_samples: int,
+    random_state: int,
+    leiden_resolution: float = 1.0,
+) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.float32], float, float]:
+    """Reduce then cluster; returns (labels, probabilities, reduce_ms, cluster_ms)."""
+    if len(vectors) == 0:
+        return (
+            np.array([], dtype=np.int64),
+            np.array([], dtype=np.float32),
+            0.0,
+            0.0,
+        )
+    started = time.perf_counter()
+    reduced = np.ascontiguousarray(vectors, dtype=np.float32)
+    match reducer:
+        case "umap" if len(vectors) > n_components + 2:
+            from umap import UMAP  # noqa: PLC0415 — heavy import, worker only
+
+            reduced = UMAP(
+                n_components=n_components,
+                n_neighbors=min(n_neighbors, max(len(vectors) - 1, 2)),
+                min_dist=min_dist,
+                metric="cosine",
+                random_state=random_state,
+            ).fit_transform(vectors)
+        case "umap":
+            pass
+        case "pca" if len(vectors) > n_components:
+            from sklearn.decomposition import PCA  # noqa: PLC0415 — worker only
+
+            reduced = PCA(
+                n_components=n_components, random_state=random_state
+            ).fit_transform(vectors)
+        case "pca" | "none":
+            pass
+        case _:
+            msg = f"unknown reducer {reducer!r}"
+            raise ValueError(msg)
+    reduce_ms = (time.perf_counter() - started) * 1_000.0
+
+    started = time.perf_counter()
+    match algorithm:
+        case "kmeans":
+            from sklearn.cluster import KMeans  # noqa: PLC0415 — worker only
+
+            k = min(len(vectors), max(1, int(np.sqrt(len(vectors) / 2))))
+            model = KMeans(n_clusters=k, random_state=random_state, n_init="auto")
+            labels = model.fit_predict(reduced).astype(np.int64)
+            probabilities = np.ones(len(vectors), dtype=np.float32)
+        case "hdbscan":
+            from sklearn.cluster import (  # noqa: PLC0415 — worker only
+                HDBSCAN,  # pyrefly: ignore[missing-module-attribute]
+            )
+
+            model = HDBSCAN(
+                min_cluster_size=min_cluster_size,
+                min_samples=min_samples,
+                metric="euclidean",
+                cluster_selection_method="eom",
+            )
+            labels = model.fit_predict(reduced).astype(np.int64)
+            probabilities = np.asarray(
+                getattr(model, "probabilities_", np.ones(len(vectors))),
+                dtype=np.float32,
+            )
+        case "leiden":
+            labels = _leiden_cluster(
+                reduced,
+                n_neighbors=n_neighbors,
+                resolution=leiden_resolution,
+                random_state=random_state,
+            )
+            probabilities = np.ones(len(vectors), dtype=np.float32)
+        case _:
+            msg = f"unknown clustering algorithm {algorithm!r}"
+            raise ValueError(msg)
+    cluster_ms = (time.perf_counter() - started) * 1_000.0
+    return labels, probabilities, reduce_ms, cluster_ms

refindery/adapters/embedding/__init__.py

@@ -0,0 +1 @@
+"""Embedding adapters (catsu multi-provider client)."""

refindery/adapters/embedding/catsu_embedder.py

@@ -0,0 +1,84 @@
+"""Embedder adapter over catsu (Voyage/OpenAI/Cohere/... behind one client).
+
+``dim``/``max_input_tokens`` come from the registry/settings — catsu does not
+guarantee exposing them per model — and the returned vector length is checked
+against ``dim`` on every call.
+"""
+
+import numpy as np
+from catsu import Client
+
+from refindery.domain.rollup import Vector
+
+_PROVIDER_ALIASES = {"voyage": "voyageai"}
+
+
+class EmbeddingDimensionMismatchError(RuntimeError):
+    """The provider returned vectors of an unexpected dimension."""
+
+    def __init__(self, *, model_id: str, expected: int, got: int) -> None:
+        super().__init__(
+            f"model {model_id!r} returned {got}-d vectors, expected {expected}; "
+            f"fix the registered dim"
+        )
+
+
+class CatsuEmbedder:
+    """Embedder port implementation for one registered model."""
+
+    def __init__(
+        self,
+        *,
+        model_id: str,
+        provider: str,
+        model_name: str,
+        dim: int,
+        max_input_tokens: int,
+    ) -> None:
+        self._model_id = model_id
+        self._provider = _PROVIDER_ALIASES.get(provider, provider)
+        self._model_name = model_name
+        self._dim = dim
+        self._max_input_tokens = max_input_tokens
+        self._client = Client()
+
+    @property
+    def model_id(self) -> str:
+        """Registry id of the model this embedder serves."""
+        return self._model_id
+
+    @property
+    def dim(self) -> int:
+        """Dimensionality of produced vectors."""
+        return self._dim
+
+    @property
+    def max_input_tokens(self) -> int:
+        """Maximum tokens the model accepts per input."""
+        return self._max_input_tokens
+
+    async def embed_documents(self, texts: list[str]) -> list[Vector]:
+        """Embed document chunks (storage side)."""
+        return await self._embed(texts, input_type="document")
+
+    async def embed_query(self, text: str) -> Vector:
+        """Embed a query (query side)."""
+        vectors = await self._embed([text], input_type="query")
+        return vectors[0]
+
+    async def _embed(self, texts: list[str], *, input_type: str) -> list[Vector]:
+        response = await self._client.aembed(
+            model=self._model_name,
+            input=texts,
+            provider=self._provider,
+            input_type=input_type,
+        )
+        vectors = [np.asarray(row, dtype=np.float32) for row in response.embeddings]
+        for vector in vectors:
+            if vector.shape != (self._dim,):
+                raise EmbeddingDimensionMismatchError(
+                    model_id=self._model_id,
+                    expected=self._dim,
+                    got=int(vector.shape[0]),
+                )
+        return vectors

refindery/adapters/embedding/surface_forms.py

@@ -0,0 +1,53 @@
+"""Surface-form embedders: model2vec static model (default)."""
+
+from threading import Lock
+from typing import Protocol, cast
+
+import numpy as np
+import numpy.typing as npt
+
+from refindery.domain.rollup import Vector, l2_normalize
+
+_DEFAULT_MODEL = "minishlab/potion-base-8M"
+
+
+class _StaticModelLike(Protocol):
+    def encode(self, sentences: list[str]) -> npt.NDArray[np.float32]:
+        """Encode surface forms."""
+        ...
+
+
+class Model2VecSurfaceEmbedder:
+    """Static-embedding surface-form encoder (no torch, microseconds/call)."""
+
+    def __init__(self, model_name: str = _DEFAULT_MODEL) -> None:
+        self._model_name = model_name
+        self._model: _StaticModelLike | None = None
+        self._lock = Lock()
+        self._id = f"model2vec:{model_name}"
+
+    @property
+    def embedder_id(self) -> str:
+        """Cache key."""
+        return self._id
+
+    def embed(self, forms: list[str]) -> list[Vector]:
+        """Encode and L2-normalize."""
+        matrix = self._load_model().encode(forms)
+        return [l2_normalize(np.asarray(row, dtype=np.float32)) for row in matrix]
+
+    def _load_model(self) -> _StaticModelLike:
+        if self._model is None:
+            with self._lock:
+                if self._model is None:
+                    from model2vec import StaticModel  # noqa: PLC0415
+
+                    self._model = cast(
+                        "_StaticModelLike",
+                        StaticModel.from_pretrained(self._model_name),
+                    )
+        model = self._model
+        if model is None:
+            msg = "surface-form model failed to load"
+            raise RuntimeError(msg)
+        return model

refindery/adapters/extraction/__init__.py

@@ -0,0 +1 @@
+"""Content fetching and extraction adapters."""

refindery/adapters/extraction/http_fetcher.py

@@ -0,0 +1,56 @@
+"""httpx-based fetcher for the fetch_and_index path.
+
+The raw response is an external input: it is validated into the pydantic
+``FetchResult`` model (size cap, content-type normalization) before anything
+downstream touches it.
+"""
+
+import httpx
+from pydantic import ValidationError
+
+from refindery.application.ports.content_extractor import FetchResult
+from refindery.domain.errors import FetchFailedError
+
+
+class HttpFetcher:
+    """Fetcher port implementation."""
+
+    def __init__(self, *, timeout_s: float = 10.0, max_bytes: int = 10_000_000) -> None:
+        self._timeout_s = timeout_s
+        self._max_bytes = max_bytes
+
+    async def fetch(self, url: str) -> FetchResult:
+        """Fetch ``url``; raise FetchFailedError on any failure."""
+        try:
+            async with (
+                httpx.AsyncClient(
+                    follow_redirects=True,
+                    timeout=self._timeout_s,
+                    headers={"User-Agent": "refindery/0.1"},
+                ) as client,
+                client.stream("GET", url) as response,
+            ):
+                response.raise_for_status()
+                body = bytearray()
+                async for chunk in response.aiter_bytes():
+                    body.extend(chunk)
+                    if len(body) > self._max_bytes:
+                        raise FetchFailedError(
+                            url=url,
+                            detail=f"body exceeds {self._max_bytes} bytes",
+                        )
+                result = FetchResult(
+                    url=url,
+                    final_url=str(response.url),
+                    status_code=response.status_code,
+                    content_type=response.headers.get(
+                        "content-type", "application/octet-stream"
+                    ),
+                    charset=response.charset_encoding,
+                    body=bytes(body),
+                )
+        except httpx.HTTPError as exc:
+            raise FetchFailedError(url=url, detail=repr(exc)) from exc
+        except ValidationError as exc:
+            raise FetchFailedError(url=url, detail=str(exc)) from exc
+        return result

refindery/adapters/extraction/pdf_pypdf.py

@@ -0,0 +1,32 @@
+"""PDF text extraction via pypdf (pure Python, BSD)."""
+
+import asyncio
+import io
+
+from pypdf import PdfReader
+
+from refindery.domain.models import ExtractedContent
+
+
+class PypdfExtractor:
+    """ContentExtractor for application/pdf."""
+
+    @property
+    def content_types(self) -> frozenset[str]:
+        """Handled content types."""
+        return frozenset({"application/pdf"})
+
+    async def extract(self, *, raw: bytes, charset: str | None) -> ExtractedContent:  # noqa: ARG002 — port signature; PDFs carry their own encoding
+        """Extract text from all pages; title from PDF metadata when present."""
+
+        def _extract() -> ExtractedContent:
+            reader = PdfReader(io.BytesIO(raw))
+            pages = [page.extract_text() or "" for page in reader.pages]
+            title = None
+            if reader.metadata is not None and reader.metadata.title:
+                title = str(reader.metadata.title)
+            return ExtractedContent(
+                body_text="\n\n".join(p for p in pages if p.strip()), title=title
+            )
+
+        return await asyncio.to_thread(_extract)

refindery/adapters/extraction/pulpie_html.py

@@ -0,0 +1,82 @@
+"""HTML -> markdown extraction via the pulpie model (transformers).
+
+Requires the ``html`` extra (torch, ~2 GB): without it this extractor
+reports itself unavailable and ingest of ``body_html`` / fetched HTML fails
+with a clear install hint. Inference runs in a ProcessPoolExecutor with the
+model loaded once per worker (spawn context; the model object never crosses
+the process boundary).
+"""
+
+import asyncio
+from collections.abc import Callable
+from concurrent.futures import ProcessPoolExecutor
+from importlib import util as importlib_util
+from multiprocessing import get_context
+
+from refindery.domain.errors import ExtractionUnavailableError
+from refindery.domain.models import ExtractedContent
+
+MODEL_NAME = "feyninc/pulpie-orange-small"
+_MAX_HTML_CHARS = 500_000
+
+_pipeline: Callable[[str], object] | None = None
+
+
+def _load_pipeline() -> Callable[[str], object]:
+    """Load the pulpie pipeline once per worker process."""
+    global _pipeline  # noqa: PLW0603 — per-process model cache
+    if _pipeline is None:
+        from transformers import pipeline  # noqa: PLC0415 — worker-only import
+
+        _pipeline = pipeline(model=MODEL_NAME)  # type: ignore[call-overload]  # ty: ignore[no-matching-overload]  # pyrefly: ignore
+    return _pipeline
+
+
+def _extract_markdown(html: str) -> str:
+    """Worker-side: run the model over raw HTML, return markdown."""
+    pipe = _load_pipeline()
+    result = pipe(html[:_MAX_HTML_CHARS])
+    if isinstance(result, str):
+        return result
+    if isinstance(result, list) and result and isinstance(result[0], dict):
+        first = dict(result[0])
+        for key in ("generated_text", "summary_text"):
+            if isinstance(value := first.get(key), str):
+                return value
+    msg = f"unexpected pulpie pipeline output shape: {type(result)!r}"
+    raise RuntimeError(msg)
+
+
+def torch_available() -> bool:
+    """Whether the html extra (torch) is installed."""
+    return importlib_util.find_spec("torch") is not None
+
+
+class PulpieHtmlExtractor:
+    """ContentExtractor for text/html via pulpie."""
+
+    def __init__(self, *, max_workers: int = 1) -> None:
+        if not torch_available():
+            raise ExtractionUnavailableError(content_type="text/html", extra="html")
+        self._executor = ProcessPoolExecutor(
+            max_workers=max_workers, mp_context=get_context("spawn")
+        )
+
+    @property
+    def content_types(self) -> frozenset[str]:
+        """Handled content types."""
+        return frozenset({"text/html", "application/xhtml+xml"})
+
+    async def extract(self, *, raw: bytes, charset: str | None) -> ExtractedContent:
+        """Decode and extract markdown body text from HTML."""
+        try:
+            html = raw.decode(charset or "utf-8", errors="replace")
+        except LookupError:
+            html = raw.decode("utf-8", errors="replace")
+        loop = asyncio.get_running_loop()
+        markdown = await loop.run_in_executor(self._executor, _extract_markdown, html)
+        return ExtractedContent(body_text=markdown)
+
+    def close(self) -> None:
+        """Shut the worker pool down."""
+        self._executor.shutdown(wait=False, cancel_futures=True)