mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/sources/catalog/embedding.py

@@ -0,0 +1,532 @@
+"""Pluggable embedding backends, and the local integer encoder the catalog ships with.
+
+Similarity finds candidates. Facts decide. Nothing in this module ever gates anything, and the
+gating path never imports it.
+
+Three properties are load-bearing.
+
+**A backend is a seam, not a dependency.** ``EmbeddingBackend`` is a protocol and
+``BackendRegistry`` is an allowlist in the same shape as the coordinator's adapter allowlist, so a
+better encoder can be dropped in later. The default one, ``LexicalHashingBackend``, is pure
+standard library: no network, no weights, no filesystem, no process spawn. A backend that cannot
+say that about itself registers with ``query_safe=False`` and a catalog sealed by it refuses to
+load at query time. ``DEFAULT_BACKENDS`` -- the allowlist the query-time load path falls back to --
+is sealed at construction, so that fallback cannot be widened at run time by an importer; a
+different inventory is a registry you build and pass, visibly, at the call site.
+
+**No IEEE-754 arithmetic anywhere.** Every step from text to vector is integer: sha256 bucketing,
+an integer signed accumulator, ``math.isqrt`` for the norm, and integer round-half-away-from-zero
+for the scaling. That is not fastidiousness. A vector that differs in its last bit between arm64
+and x86_64 would make retrieval host-dependent, and the digest authority refuses non-integer JSON
+numbers outright, so a real number could never be sealed even if one were computed. A test reads
+this module's own syntax tree and fails if a division, a square root, or the name of the
+real-number type appears in it.
+
+**One bounded seam carries provider-scale encoding.** ``encode_many`` is the only entry point a
+provider-scale publisher may use, and it is a bounded loop over the same scalar ``encode`` rather
+than a second encoder: at most 1000 items and 8 MiB of aggregate UTF-8 per invocation, one exact
+vector per input in input order, all-or-nothing, and a counter record the seam mints from work it
+measured itself. Nothing about it is faster than the scalar path, which is the honest position --
+no native batch was measured on either declared runtime tuple, so none is offered and none is
+claimed to be equivalent.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import math
+import re
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from mostlyright.data_harness.sources.catalog.contracts import QUANTIZATIONS
+from mostlyright.data_harness.sources.contracts import SourceContractError
+
+DIMENSIONS = 256
+SCALE = 10_000
+MAX_ENCODE_CHARS = 100_000
+MIN_SUBTOKEN_SOURCE = 5
+SUBTOKEN_LENGTH = 4
+
+# The bounds on one outer `encode_many` call. Both are hard: a publisher that has more work than
+# this issues more invocations, each one separately bounded and separately metered.
+MAX_BATCH_ITEMS = 1_000
+MAX_BATCH_UTF8_BYTES = 8 * 1024 * 1024
+
+#: The only encoding-adapter mode that exists. It is a fixed property of this seam rather than a
+#: caller-selectable option: every backend behind `encode_many` runs the same bounded scalar loop.
+BOUNDED_SCALAR_ADAPTER = "bounded_scalar_adapter"
+
+_TOKEN = re.compile(r"[a-z0-9]+")
+_IDENTIFIER = re.compile(r"^[a-z][a-z0-9]*(?:[._-][a-z0-9]+)*$")
+_SEMVER = re.compile(r"^[1-9][0-9]*\.[0-9]+\.[0-9]+$")
+
+
+def _identifier(value: Any, path: str) -> str:
+    if not isinstance(value, str) or not _IDENTIFIER.fullmatch(value):
+        raise SourceContractError(
+            "IDENTIFIER",
+            path,
+            "must be a lowercase dotted/dashed/underscored identifier",
+        )
+    return value
+
+
+def _semver(value: Any, path: str) -> str:
+    if not isinstance(value, str) or not _SEMVER.fullmatch(value):
+        raise SourceContractError("SEMVER", path, "must be a positive-major semantic version")
+    return value
+
+
+@dataclass(frozen=True)
+class BackendDescriptor:
+    """The immutable identity of one embedding backend.
+
+    ``query_safe`` is the query-time trust boundary. It is true only for a backend that is *pure*:
+    it opens no socket, reads no file, loads no weights, spawns no process, and calls no paid API.
+    A remote or weight-loading backend must register with ``query_safe=False``; a catalog sealed
+    by such a backend is then refused by the query-time load path rather than quietly pulling a
+    model onto the runtime path.
+    """
+
+    backend_id: str
+    backend_version: str
+    dimensions: int
+    quantization: str
+    query_safe: bool
+
+    def __post_init__(self) -> None:
+        _identifier(self.backend_id, "backend.backend_id")
+        _semver(self.backend_version, "backend.backend_version")
+        if type(self.dimensions) is not int or not 1 <= self.dimensions <= 4_096:
+            raise SourceContractError(
+                "INTEGER",
+                "backend.dimensions",
+                "must be an integer in [1, 4096]",
+            )
+        if self.quantization not in QUANTIZATIONS:
+            raise SourceContractError(
+                "ENUM",
+                "backend.quantization",
+                f"must be one of {sorted(QUANTIZATIONS)}",
+            )
+        if type(self.query_safe) is not bool:
+            raise SourceContractError("BOOLEAN", "backend.query_safe", "must be a boolean")
+
+    @property
+    def coordinate(self) -> str:
+        return f"{self.backend_id}@{self.backend_version}"
+
+
+@dataclass(frozen=True)
+class PinnedLocalBackendDescriptor(BackendDescriptor):
+    """Explicit v2 identity for a verified local-model sidecar.
+
+    This is deliberately not a reinterpretation of ``query_safe``. A local model reads governed
+    sidecar bytes and loads an optional runtime, so it remains false and cannot enter a v1 sealed
+    index. The v2 caller must select the resolver explicitly and receives the root, PIN, artifact,
+    and supported runtime-tuple coordinate it selected.
+    """
+
+    model_root: str
+    pin_sha256: str
+    coordinate_sha256: str
+    runtime_tuple: str
+    runtime_closure_sha256: str
+    artifact_sha256s: tuple[tuple[str, str], ...]
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.query_safe:
+            raise SourceContractError(
+                "NEURAL_QUERY_SAFE",
+                "backend.query_safe",
+                "a sidecar-loading backend cannot claim the v1 query-safe capability",
+            )
+        if not isinstance(self.model_root, str) or not self.model_root.startswith("/"):
+            raise SourceContractError(
+                "NEURAL_MODEL_ROOT_ABSOLUTE",
+                "backend.model_root",
+                "must be an absolute caller-supplied path",
+            )
+        for value, path in (
+            (self.pin_sha256, "backend.pin_sha256"),
+            (self.coordinate_sha256, "backend.coordinate_sha256"),
+            (self.runtime_closure_sha256, "backend.runtime_closure_sha256"),
+        ):
+            if not isinstance(value, str) or re.fullmatch(r"[0-9a-f]{64}", value) is None:
+                raise SourceContractError("SHA256", path, "must be a lowercase SHA-256 digest")
+        _identifier(self.runtime_tuple, "backend.runtime_tuple")
+        if (
+            not self.artifact_sha256s
+            or tuple(sorted(self.artifact_sha256s)) != self.artifact_sha256s
+        ):
+            raise SourceContractError(
+                "NEURAL_ARTIFACT_COORDINATE",
+                "backend.artifact_sha256s",
+                "must be a non-empty path-sorted artifact coordinate",
+            )
+        for path, digest in self.artifact_sha256s:
+            if not isinstance(path, str) or not path or "/" in path or "\\" in path:
+                raise SourceContractError(
+                    "NEURAL_ARTIFACT_COORDINATE",
+                    "backend.artifact_sha256s",
+                    "artifact paths must be single safe names",
+                )
+            if not isinstance(digest, str) or re.fullmatch(r"[0-9a-f]{64}", digest) is None:
+                raise SourceContractError(
+                    "NEURAL_ARTIFACT_COORDINATE",
+                    "backend.artifact_sha256s",
+                    "artifact digests must be lowercase SHA-256 values",
+                )
+
+    @property
+    def coordinate(self) -> str:
+        return (
+            f"{self.backend_id}@{self.backend_version}"
+            f"#{self.runtime_tuple}:{self.coordinate_sha256}:{self.runtime_closure_sha256}"
+        )
+
+
+@dataclass(frozen=True)
+class BatchEncodeResult:
+    """One outer bounded encoding invocation: its vectors, and the work it metered doing them.
+
+    The counters are the point. A provider-scale publisher has to state how much encoding a
+    generation cost, and a total it computes from its own intentions proves nothing -- it would
+    reconcile with itself whether or not the encoder ran. So this record is minted by
+    ``encode_many`` from work it observed: ``batch_invocations`` is always exactly one because this
+    record *is* one invocation, ``scalar_invocations`` is the number of times the underlying scalar
+    ``encode`` was actually called, and ``encoded_utf8_bytes`` is the aggregate the seam measured
+    before it encoded anything. A publisher spanning several invocations sums ordered records; it
+    never authors one.
+
+    ``adapter_mode`` is fixed. It exists so the publication evidence says out loud which
+    implementation produced the vectors, and there is exactly one: the bounded scalar adapter. No
+    native batching was measured on either declared runtime tuple, so no other mode may be recorded
+    here, and a record claiming one is refused rather than believed.
+    """
+
+    adapter_mode: str
+    backend_coordinate: str
+    vectors: tuple[tuple[int, ...], ...]
+    batch_invocations: int
+    scalar_invocations: int
+    encoded_utf8_bytes: int
+
+    def __post_init__(self) -> None:
+        if self.adapter_mode != BOUNDED_SCALAR_ADAPTER:
+            raise SourceContractError(
+                "EMBEDDING_BATCH_ADAPTER",
+                "embedding.batch.adapter_mode",
+                f"the only encoding adapter is {BOUNDED_SCALAR_ADAPTER!r}",
+            )
+        if not isinstance(self.backend_coordinate, str) or not self.backend_coordinate:
+            raise SourceContractError(
+                "EMBEDDING_BATCH_COUNTER",
+                "embedding.batch.backend_coordinate",
+                "the invocation must name the exact backend coordinate that produced it",
+            )
+        if type(self.vectors) is not tuple or not 1 <= len(self.vectors) <= MAX_BATCH_ITEMS:
+            raise SourceContractError(
+                "EMBEDDING_BATCH_COUNTER",
+                "embedding.batch.vectors",
+                f"one invocation carries 1 to {MAX_BATCH_ITEMS} vectors",
+            )
+        if type(self.batch_invocations) is not int or self.batch_invocations != 1:
+            raise SourceContractError(
+                "EMBEDDING_BATCH_COUNTER",
+                "embedding.batch.batch_invocations",
+                "one counter record accounts for exactly one outer batch invocation",
+            )
+        if type(self.scalar_invocations) is not int or self.scalar_invocations != len(self.vectors):
+            raise SourceContractError(
+                "EMBEDDING_BATCH_COUNTER",
+                "embedding.batch.scalar_invocations",
+                "the bounded scalar adapter invokes scalar encode exactly once per item",
+            )
+        if (
+            type(self.encoded_utf8_bytes) is not int
+            or not 0 <= self.encoded_utf8_bytes <= MAX_BATCH_UTF8_BYTES
+        ):
+            raise SourceContractError(
+                "EMBEDDING_BATCH_COUNTER",
+                "embedding.batch.encoded_utf8_bytes",
+                f"metered input bytes must be in [0, {MAX_BATCH_UTF8_BYTES}]",
+            )
+
+
+class EmbeddingBackend(Protocol):
+    """One way of turning text into a sealed integer vector."""
+
+    @property
+    def descriptor(self) -> BackendDescriptor:
+        """Return the immutable backend identity."""
+
+    def encode(self, text: str) -> tuple[int, ...]:
+        """Return this text's quantized vector, of exactly ``descriptor.dimensions`` integers."""
+
+    def encode_many(self, texts: tuple[str, ...] | list[str]) -> BatchEncodeResult:
+        """Return one exact vector per input, in order, with this invocation's own counters."""
+
+
+def encode_many_bounded(
+    backend: EmbeddingBackend, texts: tuple[str, ...] | list[str]
+) -> BatchEncodeResult:
+    """The one bounded provider-scale seam: a scalar loop, in order, with metered work.
+
+    Every backend's ``encode_many`` is this function. That is deliberate. A per-backend batch
+    implementation is where a batch stops being equal to its scalar equivalent -- padding to the
+    longest member of a batch, a different reduction order, a runtime kernel selected by shape --
+    and no cross-platform receipt proves that any of those still produce the sealed bytes.
+    A loop needs no such receipt: the vectors are the scalar vectors, because they *are* the scalar
+    vectors.
+
+    Two things are metered separately and neither is inferred from the other. Aggregate UTF-8 bytes
+    are counted first, over the whole input, before a single item is encoded, so an oversized batch
+    costs nothing to refuse. Underlying scalar invocations are then counted as they happen, so the
+    reported count is an observation rather than ``len(texts)`` restated.
+
+    The seam is all-or-nothing. A typed refusal from the scalar encoder -- a deadline, a lost
+    worker, an over-long item -- propagates unchanged and no partial batch is returned, because a
+    publisher that received nine of ten vectors would have no honest way to say so.
+    """
+
+    if type(texts) is not tuple and type(texts) is not list:
+        raise SourceContractError(
+            "EMBEDDING_BATCH_SEQUENCE",
+            "embedding.texts",
+            "must be a materialized tuple or list of texts, not a string, mapping, or iterator",
+        )
+    items = tuple(texts)
+    if not items:
+        raise SourceContractError(
+            "EMBEDDING_BATCH_EMPTY",
+            "embedding.texts",
+            "one invocation encodes at least one text",
+        )
+    if len(items) > MAX_BATCH_ITEMS:
+        raise SourceContractError(
+            "EMBEDDING_BATCH_ITEM_LIMIT",
+            "embedding.texts",
+            f"one invocation encodes at most {MAX_BATCH_ITEMS} texts",
+        )
+
+    encoded_utf8_bytes = 0
+    for position, text in enumerate(items):
+        if type(text) is not str:
+            raise SourceContractError(
+                "EMBEDDING_BATCH_ITEM_TYPE",
+                f"embedding.texts[{position}]",
+                "every item must be a string",
+            )
+        encoded_utf8_bytes += len(text.encode("utf-8"))
+        if encoded_utf8_bytes > MAX_BATCH_UTF8_BYTES:
+            raise SourceContractError(
+                "EMBEDDING_BATCH_BYTE_LIMIT",
+                "embedding.texts",
+                f"one invocation encodes at most {MAX_BATCH_UTF8_BYTES} aggregate UTF-8 bytes",
+            )
+
+    dimensions = backend.descriptor.dimensions
+    vectors: list[tuple[int, ...]] = []
+    first_seen: dict[str, tuple[int, ...]] = {}
+    scalar_invocations = 0
+    for position, text in enumerate(items):
+        vector = backend.encode(text)
+        scalar_invocations += 1
+        if (
+            type(vector) is not tuple
+            or len(vector) != dimensions
+            or any(type(value) is not int for value in vector)
+        ):
+            raise SourceContractError(
+                "EMBEDDING_BATCH_RESULT",
+                f"embedding.vectors[{position}]",
+                f"each scalar result must be exactly {dimensions} integers",
+            )
+        # A repeated text must receive its own identical vector. The encoder is deterministic, so
+        # a disagreement here is a result arriving against the wrong request -- a reordered or
+        # stale response -- which is the one mis-binding a serial loop cannot otherwise see.
+        if first_seen.setdefault(text, vector) != vector:
+            raise SourceContractError(
+                "EMBEDDING_BATCH_ORDER",
+                f"embedding.vectors[{position}]",
+                "a repeated text received a different vector, so results are not in input order",
+            )
+        vectors.append(vector)
+
+    if len(vectors) != len(items) or scalar_invocations != len(items):
+        raise SourceContractError(
+            "EMBEDDING_BATCH_RESULT",
+            "embedding.vectors",
+            "the invocation produced a partial result set",
+        )
+    return BatchEncodeResult(
+        adapter_mode=BOUNDED_SCALAR_ADAPTER,
+        backend_coordinate=backend.descriptor.coordinate,
+        vectors=tuple(vectors),
+        batch_invocations=1,
+        scalar_invocations=scalar_invocations,
+        encoded_utf8_bytes=encoded_utf8_bytes,
+    )
+
+
+def resolve_local_neural_backend(
+    model_root: Any,
+) -> EmbeddingBackend:
+    """Resolve the governed v2 neural backend from one explicit sidecar root.
+
+    The exact module name is the allowlisted boundary. A built-in dynamic import keeps this module
+    and all lexical/query imports free of the optional runtime. Resolution prepares the
+    authenticated private worker before returning, so each later ``encode`` call covers inference
+    only.
+    """
+
+    module = __import__(
+        "mostlyright.data_harness.sources.catalog.neural",
+        fromlist=("resolve_local_neural_backend",),
+    )
+    return module.resolve_local_neural_backend(model_root)
+
+
+class BackendRegistry:
+    """Allowlist of embedding backends, in the shape of the coordinator's adapter allowlist.
+
+    ``sealed`` closes the allowlist to further registration. It exists for one instance --
+    ``DEFAULT_BACKENDS``, the registry ``SealedEmbeddingIndex.load`` resolves against when a caller
+    names none. A process-global whose membership any
+    importer can widen at run time is not a control: widening it with a ``query_safe=True``
+    descriptor is all it would take to make ``load`` accept a catalog sealed by a backend that
+    opens a socket. A caller that genuinely wants a different inventory builds its own registry and
+    passes it, which is a visible act at the call site rather than an invisible one at import time.
+
+    Sealing is one-way and set at construction. There is no unseal.
+    """
+
+    def __init__(
+        self,
+        backends: tuple[EmbeddingBackend, ...] = (),
+        *,
+        sealed: bool = False,
+    ) -> None:
+        self._backends: dict[tuple[str, str], EmbeddingBackend] = {}
+        self._sealed = False
+        for backend in backends:
+            self.register(backend)
+        if type(sealed) is not bool:
+            raise SourceContractError("BOOLEAN", "registry.sealed", "must be a boolean")
+        self._sealed = sealed
+
+    @property
+    def sealed(self) -> bool:
+        """Whether this allowlist is closed to further registration."""
+
+        return self._sealed
+
+    def register(self, backend: EmbeddingBackend) -> None:
+        if self._sealed:
+            raise SourceContractError(
+                "EMBEDDING_BACKEND_REGISTRY_SEALED",
+                "registry",
+                "this backend allowlist is sealed and cannot be widened; build your own registry "
+                "and pass it explicitly instead",
+            )
+        descriptor = backend.descriptor
+        if not isinstance(descriptor, BackendDescriptor):
+            raise SourceContractError(
+                "EMBEDDING_BACKEND_DESCRIPTOR",
+                "backend.descriptor",
+                "backend descriptor must use the strict contract",
+            )
+        key = (descriptor.backend_id, descriptor.backend_version)
+        if key in self._backends:
+            raise SourceContractError(
+                "EMBEDDING_BACKEND_DUPLICATE",
+                "registry",
+                f"backend {descriptor.coordinate} is already registered",
+            )
+        self._backends[key] = backend
+
+    def resolve(self, backend_id: str, backend_version: str) -> EmbeddingBackend:
+        try:
+            return self._backends[(backend_id, backend_version)]
+        except KeyError:
+            raise SourceContractError(
+                "EMBEDDING_BACKEND_UNAVAILABLE",
+                "vector.backend_id",
+                f"backend {backend_id}@{backend_version} is not allowlisted",
+            ) from None
+
+    def inventory(self) -> tuple[BackendDescriptor, ...]:
+        return tuple(self._backends[key].descriptor for key in sorted(self._backends))
+
+
+class LexicalHashingBackend:
+    """The default local encoder: hashed unigrams plus character 4-grams, integer-quantized.
+
+    The 4-grams are a cheap stemming surrogate. They are what lets a query phrased as "produce"
+    reach a description that says "production" without a stemmer, a stopword list, or a model.
+    """
+
+    _DESCRIPTOR = BackendDescriptor(
+        backend_id="lexical_hashing",
+        backend_version="1.0.0",
+        dimensions=DIMENSIONS,
+        quantization="int32_fixed_point_10000",
+        query_safe=True,
+    )
+
+    @property
+    def descriptor(self) -> BackendDescriptor:
+        return self._DESCRIPTOR
+
+    def encode(self, text: str) -> tuple[int, ...]:
+        if not isinstance(text, str):
+            raise SourceContractError("TYPE", "embedding.text", "must be a string")
+        if len(text) > MAX_ENCODE_CHARS:
+            raise SourceContractError(
+                "EMBEDDING_INPUT_LIMIT",
+                "embedding.text",
+                f"encode input exceeds {MAX_ENCODE_CHARS} characters",
+            )
+        accumulator = [0] * DIMENSIONS
+        for token in _tokens(text):
+            digest = hashlib.sha256(token.encode("utf-8")).digest()
+            bucket = int.from_bytes(digest[:4], "big") % DIMENSIONS
+            accumulator[bucket] += 1 if digest[4] & 1 else -1
+        square_sum = sum(value * value for value in accumulator)
+        if square_sum == 0:
+            return (0,) * DIMENSIONS
+        # Exact integer square root. No IEEE-754 operation is involved anywhere on this path, so
+        # the vector is bit-identical on every host.
+        norm = math.isqrt(square_sum)
+        if norm == 0:
+            return (0,) * DIMENSIONS
+        scaled: list[int] = []
+        for value in accumulator:
+            sign = -1 if value < 0 else 1
+            # Integer round-half-away-from-zero of |value| * SCALE / norm.
+            scaled.append(sign * ((abs(value) * SCALE * 2 + norm) // (norm * 2)))
+        return tuple(scaled)
+
+    def encode_many(self, texts: tuple[str, ...] | list[str]) -> BatchEncodeResult:
+        return encode_many_bounded(self, texts)
+
+
+def _tokens(text: str) -> list[str]:
+    """Lowercased unigrams, plus character 4-grams of every unigram longer than five."""
+
+    tokens: list[str] = []
+    for match in _TOKEN.finditer(text.lower()):
+        token = match.group()
+        tokens.append(token)
+        if len(token) > MIN_SUBTOKEN_SOURCE:
+            for start in range(len(token) - SUBTOKEN_LENGTH + 1):
+                tokens.append(token[start : start + SUBTOKEN_LENGTH])
+    return tokens
+
+
+# Sealed. This is the allowlist `SealedEmbeddingIndex.load` uses when a caller names none.
+# See `BackendRegistry` for why it may not be widened in place.
+DEFAULT_BACKENDS = BackendRegistry((LexicalHashingBackend(),), sealed=True)