relvector 0.1.0__py3-none-any.whl

relvector/__init__.py

@@ -0,0 +1,48 @@
+"""RelVector — schema-aware semantic indexes over existing relational databases.
+
+RelVector does not turn a SQL database into a vector database. It builds and maintains a
+semantic index over *selected* relational data, keeping source references back to the rows;
+the original database remains the authoritative source and is only ever read.
+"""
+
+from relvector._version import __version__
+from relvector.api import RelVector, connect
+from relvector.embeddings import EmbeddingProvider, HashingEmbedding, SentenceTransformerEmbedding
+from relvector.errors import (
+    IndexMismatchError,
+    MappingError,
+    QueryRejectedError,
+    RelVectorError,
+    RunSupersededError,
+    SchemaError,
+    SensitiveFieldError,
+)
+from relvector.indexing import IndexStats
+from relvector.mapping import EntityMapping, Include, SemanticRecord
+from relvector.schema import Schema, inspect_schema
+from relvector.store import LocalVectorStore, SearchResult, VectorStore
+
+__all__ = [
+    "EmbeddingProvider",
+    "EntityMapping",
+    "HashingEmbedding",
+    "Include",
+    "IndexMismatchError",
+    "IndexStats",
+    "LocalVectorStore",
+    "MappingError",
+    "QueryRejectedError",
+    "RelVector",
+    "RelVectorError",
+    "RunSupersededError",
+    "Schema",
+    "SchemaError",
+    "SearchResult",
+    "SemanticRecord",
+    "SensitiveFieldError",
+    "SentenceTransformerEmbedding",
+    "VectorStore",
+    "__version__",
+    "connect",
+    "inspect_schema",
+]

relvector/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

relvector/api.py

@@ -0,0 +1,187 @@
+"""The :class:`RelVector` facade: connect, inspect, configure, index, search."""
+
+from __future__ import annotations
+
+import sqlite3
+from collections.abc import Callable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine import Engine, make_url
+
+from relvector.embeddings import EmbeddingProvider, get_provider
+from relvector.errors import MappingError, RelVectorError
+from relvector.extraction import Extractor, SqlValidator
+from relvector.indexing import Indexer, IndexStats
+from relvector.mapping import EntityMapping, Include, ResolvedMapping, resolve
+from relvector.schema import Schema, inspect_schema
+from relvector.store import LocalVectorStore, SearchResult, VectorStore
+
+SUPPORTED_DIALECTS = ("sqlite", "postgresql")
+
+
+def connect(source_url: str) -> Engine:
+    """Create a read-only engine for a SQLite file or a PostgreSQL database.
+
+    SQLite files are opened with ``mode=ro`` (a missing file is an error, never silently
+    created); PostgreSQL sessions are set to ``READ ONLY`` transactions. Read-only mode is a
+    safety net; the database role should also be read-only.
+    """
+    url = make_url(source_url)
+    backend = url.get_backend_name()
+    if backend == "sqlite":
+        if not url.database or url.database == ":memory:":
+            raise RelVectorError("an in-memory SQLite URL has no data to index; pass engine=")
+        path = Path(url.database).resolve()
+        if not path.is_file():
+            raise RelVectorError(f"SQLite database not found: {path}")
+        uri = f"{path.as_uri()}?mode=ro"
+        return create_engine(
+            "sqlite://",
+            creator=lambda: sqlite3.connect(uri, uri=True, check_same_thread=False),
+        )
+    if backend == "postgresql":
+        # REPEATABLE READ gives each batch's root and related-row queries one snapshot.
+        return create_engine(source_url, execution_options={
+            "postgresql_readonly": True, "isolation_level": "REPEATABLE READ"})
+    raise RelVectorError(f"unsupported database {backend!r}; supported: "
+                         f"{', '.join(SUPPORTED_DIALECTS)}")
+
+
+class RelVector:
+    """Semantic index over selected relational data; the source database stays authoritative.
+
+    Example::
+
+        rv = RelVector("sqlite:///shop.db", embedding_provider="local",
+                       index_path="shop.rvindex")
+        rv.configure(root_table="orders", include_tables=["shipments"],
+                     text_fields=["orders.status", "shipments.delay_reason"],
+                     metadata_fields=["orders.status"])
+        rv.index(batch_size=500)
+        for hit in rv.search("orders delayed by transport problems", top_k=5):
+            print(hit.score, hit.source_pk, hit.text)
+    """
+
+    def __init__(
+        self,
+        source_url: str | None = None,
+        *,
+        engine: Engine | None = None,
+        embedding_provider: str | EmbeddingProvider = "local",
+        vector_store: str | VectorStore = "local",
+        index_path: str | Path = "relvector_index.sqlite",
+        schema: str | None = None,
+        sql_validator: SqlValidator | None = None,
+        allow_external_embeddings: bool = False,
+    ) -> None:
+        if (source_url is None) == (engine is None):
+            raise RelVectorError("pass exactly one of source_url or engine")
+        self.engine = engine if engine is not None else connect(source_url)  # type: ignore[arg-type]
+        self.source = (make_url(source_url).render_as_string(hide_password=True) if source_url
+                       else engine.url.render_as_string(hide_password=True))  # type: ignore[union-attr]
+        self.provider = get_provider(embedding_provider)
+        if self.provider.sends_data_externally and not allow_external_embeddings:
+            raise RelVectorError(
+                f"embedding provider {self.provider.model_id!r} sends source text to an external "
+                "service; pass allow_external_embeddings=True to opt in explicitly"
+            )
+        if isinstance(vector_store, str):
+            if vector_store != "local":
+                raise RelVectorError(f"unknown vector store {vector_store!r}; use 'local'")
+            self.store: VectorStore = LocalVectorStore(index_path)
+        else:
+            self.store = vector_store
+        self.db_schema = schema
+        self.sql_validator = sql_validator
+        self._schema: Schema | None = None
+        self._mapping: ResolvedMapping | None = None
+
+    def __repr__(self) -> str:
+        return f"RelVector(source={self.source!r}, provider={self.provider.model_id!r})"
+
+    def __enter__(self) -> RelVector:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self.store.close()
+        self.engine.dispose()
+
+    def inspect_schema(self, *, refresh: bool = False) -> Schema:
+        if self._schema is None or refresh:
+            self._schema = inspect_schema(self.engine, schema=self.db_schema)
+        return self._schema
+
+    def configure(
+        self,
+        mapping: EntityMapping | None = None,
+        *,
+        root_table: str | None = None,
+        include_tables: Sequence[str | Include] = (),
+        text_fields: Sequence[str] = (),
+        metadata_fields: Sequence[str] = (),
+        exclude_fields: Sequence[str] = (),
+        allow_sensitive: Sequence[str] = (),
+        templates: Mapping[str, str] | None = None,
+        labels: Mapping[str, str] | None = None,
+        name: str | None = None,
+        version: str = "1",
+        max_children: int = 20,
+    ) -> ResolvedMapping:
+        """Set the entity mapping (pass an :class:`EntityMapping` or keyword arguments)."""
+        if mapping is None:
+            if root_table is None:
+                raise MappingError("root_table is required")
+            mapping = EntityMapping(
+                root_table=root_table, include=tuple(include_tables),
+                text_fields=tuple(text_fields), metadata_fields=tuple(metadata_fields),
+                exclude_fields=tuple(exclude_fields), allow_sensitive=tuple(allow_sensitive),
+                templates=dict(templates or {}), labels=dict(labels or {}), name=name,
+                version=version, max_children=max_children,
+            )
+        self._mapping = resolve(mapping, self.inspect_schema())
+        return self._mapping
+
+    @property
+    def mapping(self) -> ResolvedMapping:
+        if self._mapping is None:
+            raise MappingError("call configure() before indexing or searching")
+        return self._mapping
+
+    def extractor(self) -> Extractor:
+        return Extractor(self.engine, self.mapping, schema=self.db_schema,
+                         validator=self.sql_validator)
+
+    def index(
+        self,
+        batch_size: int = 500,
+        *,
+        resume: bool = False,
+        prune: bool = True,
+        progress: Callable[[IndexStats], None] | None = None,
+    ) -> IndexStats:
+        """(Re-)index the configured entity. Idempotent; see :mod:`relvector.indexing`."""
+        return Indexer(self.extractor(), self.provider, self.store).run(
+            batch_size, resume=resume, prune=prune, progress=progress)
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        *,
+        filters: Mapping[str, Any] | None = None,
+        entity: str | None = None,
+    ) -> list[SearchResult]:
+        """Nearest records by cosine similarity. Similarity is not evidence of causation.
+
+        ``filters`` match stored metadata exactly, e.g. ``{"orders.status": "delayed"}``.
+        ``entity`` defaults to the configured mapping's entity, if any.
+        """
+        if entity is None and self._mapping is not None:
+            entity = self._mapping.entity
+        vector = self.provider.embed([query])
+        return self.store.search(vector[0], top_k, entity=entity, filters=filters)

relvector/embeddings.py

@@ -0,0 +1,146 @@
+"""Embedding providers.
+
+A provider is any object with ``model_id``, ``dimension``, ``sends_data_externally`` and
+``embed(texts) -> ndarray[(n, dimension), float32]``. Two are built in:
+
+* :class:`SentenceTransformerEmbedding` — a local ``sentence-transformers`` model
+  (``pip install relvector[local]``). Nothing leaves the machine; the model is downloaded by
+  sentence-transformers on first use unless ``local_files_only=True``.
+* :class:`HashingEmbedding` — dependency-free, deterministic lexical feature hashing. It is not a
+  semantic model; it exists for tests, offline demos and as a keyword-ish baseline.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from collections.abc import Sequence
+from typing import Any, Protocol, runtime_checkable
+
+import numpy as np
+from numpy.typing import NDArray
+
+from relvector.errors import RelVectorError
+
+Vectors = NDArray[np.float32]
+
+
+@runtime_checkable
+class EmbeddingProvider(Protocol):
+    model_id: str
+    sends_data_externally: bool
+
+    @property
+    def dimension(self) -> int: ...
+
+    def embed(self, texts: Sequence[str]) -> Vectors: ...
+
+
+def normalize(vectors: Any) -> Vectors:
+    arr = np.asarray(vectors, dtype=np.float32)
+    if arr.ndim != 2:
+        raise RelVectorError(f"embeddings must be 2-dimensional, got shape {arr.shape}")
+    norms = np.linalg.norm(arr, axis=1, keepdims=True)
+    norms[norms == 0] = 1.0
+    out: Vectors = arr / norms
+    return out
+
+
+_TOKEN = re.compile(r"[a-z0-9]+")
+
+
+class HashingEmbedding:
+    """Signed feature hashing of words and character trigrams (stable across processes)."""
+
+    sends_data_externally = False
+
+    def __init__(self, dimension: int = 512) -> None:
+        if dimension < 8:
+            raise ValueError("dimension must be >= 8")
+        self._dimension = dimension
+        self.model_id = f"hashing-v1-{dimension}"
+
+    @property
+    def dimension(self) -> int:
+        return self._dimension
+
+    def _features(self, text: str) -> list[str]:
+        feats: list[str] = []
+        for word in _TOKEN.findall(text.lower()):
+            feats.append(f"w:{word}")
+            padded = f"#{word}#"
+            feats.extend(f"c:{padded[i:i + 3]}" for i in range(len(padded) - 2))
+        return feats
+
+    def embed(self, texts: Sequence[str]) -> Vectors:
+        out = np.zeros((len(texts), self._dimension), dtype=np.float32)
+        for row, text in enumerate(texts):
+            for feat in self._features(text):
+                h = int.from_bytes(hashlib.blake2b(feat.encode(), digest_size=8).digest(), "little")
+                out[row, h % self._dimension] += 1.0 if (h >> 63) & 1 else -1.0
+        return normalize(out)
+
+
+class SentenceTransformerEmbedding:
+    """A local sentence-transformers model, loaded lazily on first use."""
+
+    sends_data_externally = False
+
+    def __init__(
+        self,
+        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+        *,
+        device: str | None = None,
+        batch_size: int = 64,
+        local_files_only: bool = False,
+    ) -> None:
+        self.model_name = model_name
+        self.model_id = f"sentence-transformers:{model_name}"
+        self.device = device
+        self.batch_size = batch_size
+        self.local_files_only = local_files_only
+        self._model: Any = None
+
+    def _load(self) -> Any:
+        if self._model is None:
+            try:
+                from sentence_transformers import SentenceTransformer  # noqa: PLC0415
+            except ImportError as exc:  # pragma: no cover - depends on the environment
+                raise RelVectorError(
+                    f"the 'local' embedding provider could not import sentence-transformers "
+                    f"({exc}); pip install 'relvector[local]' or use "
+                    "embedding_provider='hashing'"
+                ) from exc
+            self._model = SentenceTransformer(
+                self.model_name, device=self.device, local_files_only=self.local_files_only
+            )
+        return self._model
+
+    @property
+    def dimension(self) -> int:
+        dim = self._load().get_sentence_embedding_dimension()
+        if not dim:
+            raise RelVectorError(f"model {self.model_name!r} does not report a dimension")
+        return int(dim)
+
+    def embed(self, texts: Sequence[str]) -> Vectors:
+        if not texts:
+            return np.zeros((0, self.dimension), dtype=np.float32)
+        vectors = self._load().encode(
+            list(texts), batch_size=self.batch_size, convert_to_numpy=True,
+            normalize_embeddings=True, show_progress_bar=False,
+        )
+        return normalize(vectors)
+
+
+def get_provider(spec: str | EmbeddingProvider, **kwargs: Any) -> EmbeddingProvider:
+    """Resolve ``"local"``/``"sentence-transformers"``, ``"hashing"`` or a provider instance."""
+    if not isinstance(spec, str):
+        if not isinstance(spec, EmbeddingProvider):
+            raise RelVectorError("embedding_provider must be a name or an EmbeddingProvider")
+        return spec
+    if spec in ("local", "sentence-transformers"):
+        return SentenceTransformerEmbedding(**kwargs)
+    if spec == "hashing":
+        return HashingEmbedding(**kwargs)
+    raise RelVectorError(f"unknown embedding provider {spec!r}; use 'local' or 'hashing'")

relvector/errors.py

@@ -0,0 +1,31 @@
+"""Exception hierarchy: every error RelVector raises on purpose is a :class:`RelVectorError`."""
+
+from __future__ import annotations
+
+
+class RelVectorError(Exception):
+    """Base class for RelVector errors."""
+
+
+class SchemaError(RelVectorError):
+    """A table, column or relationship does not exist or cannot be resolved."""
+
+
+class MappingError(RelVectorError):
+    """An entity mapping is invalid (unknown field, ambiguous join, conflicting options)."""
+
+
+class SensitiveFieldError(MappingError):
+    """A mapping would embed or store a column that looks sensitive and was not authorised."""
+
+
+class IndexMismatchError(RelVectorError):
+    """A persisted index was built with a different embedding model or dimension."""
+
+
+class RunSupersededError(RelVectorError):
+    """Another indexing run for the same entity started; this run stopped without pruning."""
+
+
+class QueryRejectedError(RelVectorError):
+    """An optional SQL validator (for example SQLGuard Ninja) rejected an extraction query."""

relvector/extraction.py

@@ -0,0 +1,141 @@
+"""Bounded, read-only extraction of root rows and their related rows.
+
+* Root rows are read in primary-key order with keyset pagination (``WHERE pk > :last``), never
+  ``OFFSET``, so each batch costs an index range scan regardless of how deep the run is.
+* Related rows are fetched per batch with ``IN (...)`` on the join key: one query per relation,
+  no joins, so a one-to-many relation can never multiply root rows (no Cartesian products).
+* One-to-many relations are capped per parent with ``ROW_NUMBER()`` in the database and the true
+  total is counted separately, so memory is bounded by ``batch_size * max_items`` and any
+  truncation is reported rather than silent.
+
+Only ``SELECT`` statements built with SQLAlchemy Core are executed; all values are bound
+parameters. Identifiers come from the inspected schema, not from free-form user SQL.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from sqlalchemy import Select, column, func, select, table, tuple_
+from sqlalchemy.engine import Connection, Engine
+from sqlalchemy.sql.expression import TableClause
+
+from relvector.mapping import Relation, ResolvedMapping
+
+SqlValidator = Callable[[str], None]
+"""Called with each extraction query (named bind parameters) before it runs; raise to block."""
+
+Row = Mapping[str, Any]
+
+
+@dataclass(frozen=True)
+class RowBundle:
+    row: Row
+    related: dict[str, tuple[list[Row], int]]
+
+
+@dataclass(frozen=True)
+class Batch:
+    bundles: list[RowBundle]
+    last_key: list[Any]
+
+
+class Extractor:
+    def __init__(
+        self,
+        engine: Engine,
+        resolved: ResolvedMapping,
+        *,
+        schema: str | None = None,
+        validator: SqlValidator | None = None,
+    ) -> None:
+        self.engine = engine
+        self.resolved = resolved
+        self.schema = schema
+        self.validator = validator
+
+    def _table(self, name: str, columns: Sequence[str]) -> TableClause:
+        return table(name, *(column(c) for c in columns), schema=self.schema)
+
+    def _execute(self, conn: Connection, stmt: Select[Any]) -> list[Row]:
+        if self.validator is not None:
+            self.validator(self.render(stmt))
+        return [dict(r) for r in conn.execute(stmt).mappings()]
+
+    def render(self, stmt: Select[Any]) -> str:
+        """The statement as SQL text with named parameters (what a validator sees)."""
+        dialect = type(self.engine.dialect)(paramstyle="named")  # type: ignore[call-arg]
+        return str(stmt.compile(dialect=dialect, compile_kwargs={"render_postcompile": True}))
+
+    def root_query(self, batch_size: int, after: Sequence[Any] | None) -> Select[Any]:
+        r = self.resolved
+        t = self._table(r.root_table, r.root_select_columns)
+        pk = [t.c[c] for c in r.primary_key]
+        # SQLite permits NULL in non-INTEGER primary keys; such rows have no stable identity.
+        stmt = (select(*(t.c[c] for c in r.root_select_columns))
+                .where(*(c.is_not(None) for c in pk)).order_by(*pk).limit(batch_size))
+        if after is not None:
+            if len(pk) == 1:
+                stmt = stmt.where(pk[0] > after[0])
+            else:
+                stmt = stmt.where(tuple_(*pk) > tuple_(*after))
+        return stmt
+
+    def iter_batches(
+        self, batch_size: int = 500, after: Sequence[Any] | None = None
+    ) -> Iterator[Batch]:
+        """Yield batches of at most ``batch_size`` root rows, resuming after key ``after``."""
+        if batch_size < 1:
+            raise ValueError("batch_size must be >= 1")
+        last = list(after) if after is not None else None
+        while True:
+            with self.engine.connect() as conn:
+                rows = self._execute(conn, self.root_query(batch_size, last))
+                if not rows:
+                    return
+                related = {rel.table: self._fetch_related(conn, rel, rows)
+                           for rel in self.resolved.relations}
+            bundles = [
+                RowBundle(row, {t: per_key.get(row[rel.root_column], ([], 0))
+                                for t, (rel, per_key) in related.items()})
+                for row in rows
+            ]
+            last = [rows[-1][c] for c in self.resolved.primary_key]
+            yield Batch(bundles, last)
+            if len(rows) < batch_size:
+                return
+
+    def _fetch_related(
+        self, conn: Connection, rel: Relation, rows: list[Row]
+    ) -> tuple[Relation, dict[Any, tuple[list[Row], int]]]:
+        keys = sorted({r[rel.root_column] for r in rows if r[rel.root_column] is not None},
+                      key=lambda k: (str(type(k)), k))
+        if not keys:
+            return rel, {}
+        t = self._table(rel.table, rel.select_columns)
+        join_col = t.c[rel.related_column]
+        cols = [t.c[c] for c in rel.select_columns]
+        out: dict[Any, tuple[list[Row], int]] = {}
+        if rel.kind == "many_to_one":
+            for r in self._execute(conn, select(*cols).where(join_col.in_(keys))):
+                out[r[rel.related_column]] = ([r], 1)
+            return rel, out
+
+        rn = func.row_number().over(
+            partition_by=join_col, order_by=[t.c[c] for c in rel.primary_key]
+        ).label("_rv_rn")
+        ranked = select(*cols, rn).where(join_col.in_(keys)).subquery()
+        capped = (select(*(ranked.c[c] for c in rel.select_columns))
+                  .where(ranked.c["_rv_rn"] <= rel.max_items)
+                  .order_by(ranked.c[rel.related_column], ranked.c["_rv_rn"]))
+        counts = select(join_col, func.count().label("_rv_n")).where(
+            join_col.in_(keys)).group_by(join_col)
+        totals = {r[rel.related_column]: r["_rv_n"] for r in self._execute(conn, counts)}
+        grouped: dict[Any, list[Row]] = {}
+        for r in self._execute(conn, capped):
+            grouped.setdefault(r[rel.related_column], []).append(r)
+        for key, total in totals.items():
+            out[key] = (grouped.get(key, []), total)
+        return rel, out