relvector 0.1.0__py3-none-any.whl

relvector/indexing.py

@@ -0,0 +1,138 @@
+"""The indexing pipeline: extract a batch -> build records -> embed changed ones -> write.
+
+Recovery model:
+
+* Each batch is written in one store transaction together with its checkpoint, so a crash
+  leaves the index at a batch boundary. ``resume=True`` continues an unfinished run from its
+  last committed key; re-processing a batch is harmless because writes are upserts keyed by a
+  deterministic vector id. Resuming requires the mapping to be unchanged. Checkpoint keys keep
+  their types (dates, UUIDs, decimals), so a resumed scan compares them as the column does.
+* Starting a run takes over the entity's checkpoint. An older run that is still going (another
+  process, or one that was abandoned) fails with ``RunSupersededError`` at its next batch and
+  never prunes; the newest run is the one that completes.
+* Each ``run()`` call is a writer with its own owner token. Resuming a run takes it over, so
+  if two processes resume the same run, the one that resumed first is rejected with
+  ``RunSupersededError`` at its next batch. There is never more than one writer that can
+  advance the checkpoint or prune.
+* Records whose content hash is unchanged are not re-embedded, only marked as seen.
+* When a run finishes, records of the entity that the run did not see are deleted (the source
+  rows were deleted or no longer map). Interrupted runs never prune.
+
+This is a full re-scan ("manual re-indexing"); it is not change-data-capture.
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from relvector.embeddings import EmbeddingProvider
+from relvector.errors import RelVectorError
+from relvector.extraction import Extractor
+from relvector.mapping import build_record
+from relvector.store import Checkpoint, VectorStore
+
+
+@dataclass
+class IndexStats:
+    run_id: str
+    rows: int = 0
+    embedded: int = 0
+    unchanged: int = 0
+    deleted: int = 0
+    batches: int = 0
+    truncated_records: int = 0
+    resumed: bool = False
+    timings: dict[str, float] = field(
+        default_factory=lambda: {"extract": 0.0, "embed": 0.0, "store": 0.0})
+
+    @property
+    def rows_per_second(self) -> float:
+        total = sum(self.timings.values())
+        return self.rows / total if total else 0.0
+
+    def __str__(self) -> str:
+        t = ", ".join(f"{k} {v:.2f}s" for k, v in self.timings.items())
+        return (f"{self.rows} rows in {self.batches} batches: {self.embedded} embedded, "
+                f"{self.unchanged} unchanged, {self.deleted} deleted ({t}; "
+                f"{self.rows_per_second:.0f} rows/s)")
+
+
+class Indexer:
+    def __init__(self, extractor: Extractor, provider: EmbeddingProvider,
+                 store: VectorStore) -> None:
+        self.extractor = extractor
+        self.provider = provider
+        self.store = store
+
+    def run(
+        self,
+        batch_size: int = 500,
+        *,
+        resume: bool = False,
+        prune: bool = True,
+        progress: Callable[[IndexStats], None] | None = None,
+    ) -> IndexStats:
+        resolved = self.extractor.resolved
+        entity = resolved.entity
+        self.store.ensure_compatible(self.provider.model_id, self.provider.dimension)
+
+        owner = uuid.uuid4().hex
+        previous = self.store.checkpoint(entity)
+        if resume and previous is not None and previous.status == "running":
+            if previous.fingerprint != resolved.fingerprint:
+                raise RelVectorError(
+                    f"cannot resume run {previous.run_id} for {entity!r}: the mapping changed "
+                    "since it started, so resuming would mix old and new documents; run "
+                    "index() without resume=True for a full re-index"
+                )
+            stats = IndexStats(run_id=previous.run_id, rows=previous.processed, resumed=True)
+            after = previous.last_key
+            self.store.resume_run(Checkpoint(entity, stats.run_id, after, "running",
+                                             stats.rows, resolved.fingerprint, owner))
+        else:
+            stats = IndexStats(run_id=uuid.uuid4().hex)
+            after = None
+            self.store.start_run(Checkpoint(entity, stats.run_id, None, "running", 0,
+                                            resolved.fingerprint, owner))
+
+        batches = self.extractor.iter_batches(batch_size, after)
+        while True:
+            t0 = time.perf_counter()
+            batch = next(batches, None)
+            t1 = time.perf_counter()
+            stats.timings["extract"] += t1 - t0
+            if batch is None:
+                break
+            records = [build_record(resolved, b.row, b.related) for b in batch.bundles]
+            known = self.store.content_hashes([r.vector_id for r in records])
+            changed = [r for r in records if known.get(r.vector_id) != r.content_hash]
+            unchanged = [r.vector_id for r in records if known.get(r.vector_id) == r.content_hash]
+            t2 = time.perf_counter()
+            stats.timings["extract"] += t2 - t1  # record building is part of extraction
+            vectors = (self.provider.embed([r.text for r in changed]) if changed
+                       else np.zeros((0, self.provider.dimension), dtype=np.float32))
+            t3 = time.perf_counter()
+            stats.timings["embed"] += t3 - t2
+            stats.rows += len(records)
+            stats.embedded += len(changed)
+            stats.unchanged += len(unchanged)
+            stats.batches += 1
+            stats.truncated_records += sum(1 for r in records if r.truncated)
+            self.store.write_batch(changed, vectors, unchanged,
+                                   Checkpoint(entity, stats.run_id, batch.last_key, "running",
+                                              stats.rows, resolved.fingerprint, owner))
+            stats.timings["store"] += time.perf_counter() - t3
+            if progress is not None:
+                progress(stats)
+
+        t4 = time.perf_counter()
+        stats.deleted = self.store.finish_run(
+            Checkpoint(entity, stats.run_id, None, "running", stats.rows, resolved.fingerprint,
+                       owner), prune=prune)
+        stats.timings["store"] += time.perf_counter() - t4
+        return stats

relvector/integrations/__init__.py

@@ -0,0 +1 @@
+"""Optional integrations with other projects (each imports its dependency lazily)."""

relvector/integrations/sqlguard.py

@@ -0,0 +1,33 @@
+"""Optional SQLGuard Ninja integration: validate every extraction query against a policy.
+
+``pip install sqlguard-ninja``. Only the public API (``Guard.validate``) is used. A denial (or
+an approval that requires human review) blocks extraction with :class:`QueryRejectedError`.
+
+SQLGuard approval means the *query* is permitted for the role; it is not permission to embed
+and store the returned values. Sensitive-field authorisation remains the mapping's job
+(``exclude_fields`` / ``allow_sensitive``).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from relvector.errors import QueryRejectedError
+
+
+class SQLGuardValidator:
+    """Callable for ``RelVector(sql_validator=...)``."""
+
+    def __init__(self, guard: Any, *, role: str | None = None, dialect: str | None = None,
+                 context: Mapping[str, Any] | None = None) -> None:
+        self.guard = guard
+        self.role = role
+        self.dialect = dialect
+        self.context = dict(context or {})
+
+    def __call__(self, sql: str) -> None:
+        result = self.guard.validate(sql, role=self.role, dialect=self.dialect,
+                                     context=self.context)
+        if not result.allowed:
+            raise QueryRejectedError(f"SQLGuard blocked an extraction query: {result.explain()}")

relvector/mapping.py

@@ -0,0 +1,405 @@
+"""Entity mappings: which relational data becomes one semantic record, and how it is worded.
+
+A mapping names a root table (one semantic record per root row), directly related tables,
+and the exact columns that may be embedded (``text_fields``) or stored alongside the vector
+(``metadata_fields``). Columns that are not listed are never selected from the database.
+
+Text is produced by deterministic templates, never by a language model, so the same rows and
+the same mapping always produce byte-identical text and the same content hash.
+"""
+
+from __future__ import annotations
+
+import datetime as _dt
+import hashlib
+import json
+import math
+import re
+import string
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from decimal import Decimal
+from typing import Any, Literal
+
+from relvector.errors import MappingError, SensitiveFieldError
+from relvector.schema import Schema, TableInfo
+
+NULL_TEXT = "not recorded"
+
+# Adapted from sqlguard_ninja.privacy.suggest (Apache-2.0, same authors). Name heuristics are
+# unreliable, so they only *block* by default; a data owner authorises via ``allow_sensitive``.
+_SENSITIVE_PATTERNS: list[tuple[str, re.Pattern[str]]] = [
+    ("payment_card", re.compile(r"(card_?(number|no|num)|\bpan\b|cc_?num|cvv)")),
+    ("financial", re.compile(r"(iban|bank|account_?number|salary|wage|income)")),
+    ("government_id", re.compile(r"(ssn|national_?id|passport|tax_?id|\bnin\b)")),
+    ("contact", re.compile(r"(e_?mail|phone|mobile|address|postcode|zip)")),
+    ("personal_identifier", re.compile(r"(full_?name|first_?name|last_?name|birth|\bdob\b)")),
+    ("special_category", re.compile(r"(health|medical|diagnos|religio|ethnic|biometric)")),
+    ("credentials", re.compile(r"(password|passwd|secret|token|api_?key)")),
+]
+
+
+def sensitive_label(column: str) -> str | None:
+    """Return a sensitivity label if the column *name* looks sensitive, else ``None``."""
+    name = column.lower()
+    for label, pattern in _SENSITIVE_PATTERNS:
+        if pattern.search(name):
+            return label
+    return None
+
+
+@dataclass(frozen=True)
+class Include:
+    """A table directly related to the root table.
+
+    ``via`` makes the join explicit as ``("root.col", "related.col")``; without it the single
+    declared foreign key between the two tables is used.
+    """
+
+    table: str
+    via: tuple[str, str] | None = None
+    max_items: int | None = None
+
+
+@dataclass(frozen=True)
+class EntityMapping:
+    root_table: str
+    text_fields: Sequence[str]
+    include: Sequence[str | Include] = ()
+    metadata_fields: Sequence[str] = ()
+    exclude_fields: Sequence[str] = ()
+    allow_sensitive: Sequence[str] = ()
+    templates: Mapping[str, str] = field(default_factory=dict)
+    labels: Mapping[str, str] = field(default_factory=dict)
+    name: str | None = None
+    version: str = "1"
+    max_children: int = 20
+
+    @property
+    def entity(self) -> str:
+        return self.name or self.root_table
+
+
+@dataclass(frozen=True)
+class Relation:
+    table: str
+    kind: Literal["many_to_one", "one_to_many"]
+    root_column: str
+    related_column: str
+    primary_key: tuple[str, ...]
+    text_columns: tuple[str, ...]
+    metadata_columns: tuple[str, ...]
+    max_items: int
+
+    @property
+    def select_columns(self) -> tuple[str, ...]:
+        return _unique((self.related_column, *self.primary_key, *self.text_columns,
+                        *self.metadata_columns))
+
+
+@dataclass(frozen=True)
+class ResolvedMapping:
+    """A mapping validated against a concrete schema."""
+
+    spec: EntityMapping
+    primary_key: tuple[str, ...]
+    text_columns: tuple[str, ...]
+    metadata_columns: tuple[str, ...]
+    relations: tuple[Relation, ...]
+    fingerprint: str
+
+    @property
+    def entity(self) -> str:
+        return self.spec.entity
+
+    @property
+    def root_table(self) -> str:
+        return self.spec.root_table
+
+    @property
+    def root_select_columns(self) -> tuple[str, ...]:
+        return _unique((*self.primary_key, *self.text_columns, *self.metadata_columns,
+                        *(r.root_column for r in self.relations)))
+
+
+@dataclass(frozen=True)
+class SemanticRecord:
+    vector_id: str
+    entity: str
+    source_table: str
+    source_pk: dict[str, Any]
+    text: str
+    metadata: dict[str, Any]
+    references: list[dict[str, Any]]
+    truncated: dict[str, int]
+    mapping_version: str
+    content_hash: str
+
+
+def _unique(items: Sequence[str]) -> tuple[str, ...]:
+    return tuple(dict.fromkeys(items))
+
+
+def _split(ref: str, root: str) -> tuple[str, str]:
+    if "." in ref:
+        table, _, column = ref.partition(".")
+        if not table or not column or "." in column:
+            raise MappingError(f"field {ref!r} must be 'table.column'")
+        return table, column
+    return root, ref
+
+
+def resolve(spec: EntityMapping, schema: Schema) -> ResolvedMapping:
+    """Validate ``spec`` against ``schema``. Raises :class:`MappingError` with a clear reason."""
+    if not spec.text_fields:
+        raise MappingError("text_fields must name at least one column to embed")
+    if spec.max_children < 1:
+        raise MappingError("max_children must be >= 1")
+    root = schema.table(spec.root_table)
+    if not root.primary_key:
+        raise MappingError(f"root table {root.name!r} has no primary key; a stable identity "
+                           "is required for idempotent indexing")
+
+    includes = [Include(i) if isinstance(i, str) else i for i in spec.include]
+    tables = {root.name: root}
+    for inc in includes:
+        if inc.table == root.name:
+            raise MappingError("self-referencing includes are not supported in this version")
+        if inc.table in tables:
+            raise MappingError(f"table {inc.table!r} is included more than once")
+        tables[inc.table] = schema.table(inc.table)
+
+    excluded = {_split(f, root.name) for f in spec.exclude_fields}
+    allowed_sensitive = {_split(f, root.name) for f in spec.allow_sensitive}
+
+    def fields(refs: Sequence[str], kind: str) -> dict[str, list[str]]:
+        out: dict[str, list[str]] = {t: [] for t in tables}
+        for ref in refs:
+            table, column = _split(ref, root.name)
+            if table not in tables:
+                raise MappingError(f"{kind} field {ref!r} uses table {table!r}, which is neither "
+                                   "the root table nor an included table")
+            tables[table].column(column)
+            if (table, column) in excluded:
+                raise MappingError(f"{kind} field {table}.{column} is also in exclude_fields")
+            label = sensitive_label(column)
+            if label and (table, column) not in allowed_sensitive:
+                raise SensitiveFieldError(
+                    f"{kind} field {table}.{column} looks sensitive ({label}); exclude it or "
+                    "list it in allow_sensitive after confirming it may be embedded and stored"
+                )
+            if column not in out[table]:
+                out[table].append(column)
+        return out
+
+    text = fields(spec.text_fields, "text")
+    meta = fields(spec.metadata_fields, "metadata")
+
+    for tname, template in spec.templates.items():
+        if tname not in tables:
+            raise MappingError(f"template given for table {tname!r}, which is not mapped")
+        allowed = {*text[tname], *tables[tname].primary_key}
+        for name in _template_fields(template):
+            if name not in allowed:
+                raise MappingError(f"template for {tname!r} uses {{{name}}}; templates may only "
+                                   "use that table's text_fields and primary key")
+
+    relations = tuple(
+        _resolve_relation(inc, root, tables[inc.table], schema, text[inc.table],
+                          meta[inc.table], spec.max_children)
+        for inc in includes
+    )
+    resolved_payload = {
+        "entity": spec.entity, "root": root.name, "version": spec.version,
+        "text": text, "meta": meta, "templates": dict(sorted(spec.templates.items())),
+        "labels": dict(sorted(spec.labels.items())),
+        "relations": [r.__dict__ for r in relations],
+    }
+    fingerprint = hashlib.sha256(
+        json.dumps(resolved_payload, sort_keys=True, default=list).encode()
+    ).hexdigest()[:16]
+    return ResolvedMapping(
+        spec=spec,
+        primary_key=root.primary_key,
+        text_columns=tuple(text[root.name]),
+        metadata_columns=tuple(meta[root.name]),
+        relations=relations,
+        fingerprint=fingerprint,
+    )
+
+
+def _resolve_relation(
+    inc: Include, root: TableInfo, related: TableInfo, schema: Schema,
+    text: list[str], meta: list[str], default_max: int,
+) -> Relation:
+    if inc.via is not None:
+        (lt, lc), (rt, rc) = (_split(inc.via[0], root.name), _split(inc.via[1], related.name))
+        if lt != root.name or rt != related.name:
+            raise MappingError(f"via for {related.name!r} must be ('{root.name}.col', "
+                               f"'{related.name}.col'), got {inc.via!r}")
+        root.column(lc)
+        related.column(rc)
+        root_col, rel_col = lc, rc
+        kind: Literal["many_to_one", "one_to_many"] = (
+            "many_to_one" if related.primary_key == (rc,) else "one_to_many"
+        )
+    else:
+        fks = [fk for fk in schema.relationships_between(root.name, related.name)
+               if not fk.composite]
+        if not fks:
+            raise MappingError(
+                f"no single-column foreign key links {root.name!r} and {related.name!r}; "
+                "declare one in the database or pass Include(..., via=(root_col, related_col))"
+            )
+        if len(fks) > 1:
+            listed = "; ".join(str(fk) for fk in fks)
+            raise MappingError(f"ambiguous relationship between {root.name!r} and "
+                               f"{related.name!r} ({listed}); pass Include(..., via=...)")
+        fk = fks[0]
+        if fk.table == root.name:
+            kind, root_col, rel_col = "many_to_one", fk.columns[0], fk.ref_columns[0]
+        else:
+            kind, root_col, rel_col = "one_to_many", fk.ref_columns[0], fk.columns[0]
+    if not related.primary_key:
+        raise MappingError(f"included table {related.name!r} has no primary key; it is needed "
+                           "for deterministic ordering and source references")
+    if not text and not meta:
+        raise MappingError(f"included table {related.name!r} contributes no text or metadata "
+                           "fields")
+    max_items = inc.max_items if inc.max_items is not None else default_max
+    if max_items < 1:
+        raise MappingError("max_items must be >= 1")
+    return Relation(
+        table=related.name, kind=kind, root_column=root_col, related_column=rel_col,
+        primary_key=related.primary_key, text_columns=tuple(text),
+        metadata_columns=tuple(meta), max_items=max_items,
+    )
+
+
+# --- deterministic serialisation ------------------------------------------------------------
+
+
+def _template_fields(template: str) -> list[str]:
+    names: list[str] = []
+    for _, name, spec, conv in string.Formatter().parse(template):
+        if name is None:
+            continue
+        if not name.isidentifier() or spec or conv:
+            raise MappingError(f"template placeholder {{{name}}} must be a plain column name "
+                               "(no attributes, indexes, conversions or format specs)")
+        names.append(name)
+    return names
+
+
+def jsonable(value: Any) -> Any:
+    if value is None or isinstance(value, bool | int | str):
+        return value
+    if isinstance(value, float):
+        return None if math.isnan(value) else value  # NaN is not valid JSON
+    if isinstance(value, Decimal):
+        return str(value)
+    if isinstance(value, _dt.date | _dt.time):
+        return value.isoformat()
+    if isinstance(value, bytes | bytearray | memoryview):
+        return f"<{len(bytes(value))} bytes>"
+    return str(value)
+
+
+def _text(value: Any) -> str:
+    if value is None:
+        return NULL_TEXT
+    if isinstance(value, bool):
+        return "yes" if value else "no"
+    if isinstance(value, bytes | bytearray | memoryview):
+        return "[binary value omitted]"
+    return " ".join(str(jsonable(value)).split())
+
+
+def _humanize(column: str) -> str:
+    return column.replace("_", " ").strip().capitalize()
+
+
+def _singular(table: str) -> str:
+    word = table.replace("_", " ")
+    if word.endswith("ies"):
+        word = word[:-3] + "y"
+    elif word.endswith("s") and not word.endswith("ss"):
+        word = word[:-1]
+    return word.capitalize()
+
+
+def _key_text(row: Mapping[str, Any], pk: Sequence[str]) -> str:
+    return "/".join(_text(row[c]) for c in pk)
+
+
+def _render(row: Mapping[str, Any], table: str, pk: Sequence[str], columns: Sequence[str],
+            spec: EntityMapping) -> str:
+    template = spec.templates.get(table)
+    if template is not None:
+        names = _template_fields(template)
+        # A null would make a template read as a false statement; fall back to explicit text.
+        if all(row[n] is not None for n in names):
+            return " ".join(template.format(**{n: _text(row[n]) for n in names}).split())
+    label = spec.labels.get(table, _singular(table))
+    parts = [f"{label} {_key_text(row, pk)}."]
+    parts.extend(f"{_humanize(c)}: {_text(row[c])}." for c in columns)
+    return " ".join(parts)
+
+
+def build_record(
+    resolved: ResolvedMapping,
+    row: Mapping[str, Any],
+    related: Mapping[str, tuple[list[Mapping[str, Any]], int]],
+) -> SemanticRecord:
+    """Serialise one root row plus its related rows. ``related[table] = (rows, total_count)``."""
+    spec = resolved.spec
+    pk = {c: jsonable(row[c]) for c in resolved.primary_key}
+    sentences = [_render(row, spec.root_table, resolved.primary_key,
+                         resolved.text_columns, spec)]
+    metadata: dict[str, Any] = {f"{spec.root_table}.{c}": jsonable(row[c])
+                                for c in resolved.metadata_columns}
+    references: list[dict[str, Any]] = []
+    truncated: dict[str, int] = {}
+
+    for rel in resolved.relations:
+        rows, total = related.get(rel.table, ([], 0))
+        noun = spec.labels.get(rel.table, _singular(rel.table)).lower()
+        if rel.kind == "many_to_one":
+            if not rows:
+                key = row[rel.root_column]
+                sentences.append(f"No related {noun} is recorded." if key is None else
+                                 f"Related {noun} {_text(key)} was not found.")
+            for r in rows:
+                sentences.append(_render(r, rel.table, rel.primary_key, rel.text_columns, spec))
+        else:
+            plural = rel.table.replace("_", " ")
+            if total == 0:
+                sentences.append(f"No related {plural}.")
+            elif total > len(rows):
+                sentences.append(f"Related {plural}: {total} in total, showing the first "
+                                 f"{len(rows)}.")
+                truncated[rel.table] = total - len(rows)
+            sentences.extend(_render(r, rel.table, rel.primary_key, rel.text_columns, spec)
+                             for r in rows)
+        for c in rel.metadata_columns:
+            values = [jsonable(r[c]) for r in rows]
+            metadata[f"{rel.table}.{c}"] = (values[0] if values else None) \
+                if rel.kind == "many_to_one" else values
+        references.extend(
+            {"table": rel.table, "relation": rel.kind,
+             "pk": {c: jsonable(r[c]) for c in rel.primary_key}}
+            for r in rows
+        )
+
+    text = " ".join(sentences)
+    vector_id = f"{spec.entity}:" + json.dumps(list(pk.values()), separators=(",", ":"))
+    content_hash = hashlib.sha256(json.dumps(
+        {"text": text, "metadata": metadata, "refs": references, "truncated": truncated,
+         "mapping": resolved.fingerprint},
+        sort_keys=True, separators=(",", ":"),
+    ).encode()).hexdigest()
+    return SemanticRecord(
+        vector_id=vector_id, entity=spec.entity, source_table=spec.root_table, source_pk=pk,
+        text=text, metadata=metadata, references=references, truncated=truncated,
+        mapping_version=spec.version, content_hash=content_hash,
+    )