isotrieve 0.2.1__py3-none-any.whl

isotrieve/__init__.py

@@ -0,0 +1,22 @@
+"""Isotrieve — Embedding migration without re-embedding.
+
+Learn a linear (or optional shallow non-linear) mapping between source and
+target embedding spaces from a small calibration sample, then transform stored
+vectors in place instead of re-embedding an entire corpus.
+
+This package does **not** claim algorithmic novelty. It productizes known
+linear-mapping techniques with a quality gate, providers, store adapters, and
+a reproducible benchmark harness. See README "Prior Art & Research Basis".
+"""
+
+from isotrieve.mapping.base import Mapping, ValidationReport
+from isotrieve.mapping.linear import OrthogonalProcrustesMapping, RidgeMapping
+
+__version__ = "0.2.1"
+__all__ = [
+    "Mapping",
+    "RidgeMapping",
+    "OrthogonalProcrustesMapping",
+    "ValidationReport",
+    "__version__",
+]

isotrieve/adapters/__init__.py

@@ -0,0 +1,6 @@
+"""Vector database adapters for Isotrieve.
+
+Provides thin wrappers that apply Isotrieve mappings to popular vector databases.
+Core mapping/gate/recalibration logic stays in isotrieve.adapters.base; DB-specific
+code is in per-DB modules.
+"""

isotrieve/adapters/base.py

@@ -0,0 +1,127 @@
+"""Core adapter abstractions.
+
+``EmbeddingAdapter`` wraps any embedding model.
+``VectorStoreAdapter`` wraps any vector DB with serve + offline modes.
+``MigrationReport`` captures what happened during offline migration.
+
+All mapping/gate/recalibration logic is imported from isotrieve core;
+DB-specific code stays thin.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+import numpy as np
+
+from isotrieve.mapping.base import Mapping, l2_normalize
+
+
+@dataclass
+class MigrationReport:
+    """Result of an offline migration pass."""
+
+    rows_processed: int = 0
+    elapsed_seconds: float = 0.0
+    sampled_recall_at_10: float | None = None
+    mapping_checksum: str = ""
+    source_collection: str = ""
+    target_collection: str = ""
+    errors: list[str] = field(default_factory=list)
+    idempotent: bool = True  # False if double-migration detected
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "rows_processed": self.rows_processed,
+            "elapsed_seconds": self.elapsed_seconds,
+            "sampled_recall_at_10": self.sampled_recall_at_10,
+            "mapping_checksum": self.mapping_checksum,
+            "source_collection": self.source_collection,
+            "target_collection": self.target_collection,
+            "errors": self.errors,
+            "idempotent": self.idempotent,
+        }
+
+
+class EmbeddingAdapter(ABC):
+    """Wraps any embedding model; source of truth for model identity."""
+
+    model_id: str
+    dim: int
+
+    @abstractmethod
+    def embed(self, texts: list[str]) -> np.ndarray:
+        """Embed a list of texts. Returns (n, dim) array."""
+
+
+class VectorStoreAdapter(ABC):
+    """Thin per-DB shell. All mapping logic lives in isotrieve core.
+
+    Parameters
+    ----------
+    mapping:
+        A fitted Isotrieve Mapping (with inverse + optional recalibrator).
+    mode:
+        ``"serve"`` = transform queries on-the-fly, corpus untouched.
+        ``"migrated"`` = corpus already transformed, no query mapping needed.
+    """
+
+    def __init__(
+        self, mapping: Mapping, mode: Literal["serve", "migrated"] = "serve"
+    ) -> None:
+        self._mapping = mapping
+        self._mode = mode
+
+    @property
+    def has_recalibrator(self) -> bool:
+        return self._mapping.has_recalibrator
+
+    def recalibrate_scores(self, scores: np.ndarray) -> np.ndarray:
+        """Apply score recalibration (WS-A) to raw similarity scores."""
+        return self._mapping.recalibrate_scores(scores)
+
+    def preflight(self, sample_vectors: np.ndarray, target_vectors: np.ndarray) -> Any:
+        """Run quality gate on a sample before serve/migrate.
+
+        Returns a GateReport. Caller should check verdict before proceeding.
+        """
+        from isotrieve.quality.gate import QualityGate
+
+        gate = QualityGate()
+        return gate.evaluate(self._mapping, sample_vectors, target_vectors)
+
+    def _map_query(self, vec: np.ndarray) -> np.ndarray:
+        """Map a single query vector from new-model space to legacy space."""
+        return l2_normalize(self._mapping.inverse_transform(vec.reshape(1, -1)).ravel())
+
+    def _map_queries(self, vecs: np.ndarray) -> np.ndarray:
+        """Batch map query vectors."""
+        return l2_normalize(self._mapping.inverse_transform(vecs))
+
+    @abstractmethod
+    def query(
+        self,
+        query_vectors: np.ndarray,
+        k: int = 10,
+        **kwargs: Any,
+    ) -> list[list[dict[str, Any]]]:
+        """Serve-mode query: map queries, search, return results.
+
+        Returns a list (per query) of lists of result dicts with keys
+        ``id``, ``score``, ``metadata``.
+        """
+
+    @abstractmethod
+    def migrate(
+        self,
+        batch_size: int = 1000,
+        dry_run: bool = False,
+        new_collection: str | None = None,
+    ) -> MigrationReport:
+        """Offline migration: transform corpus vectors and write to new collection.
+
+        Never modifies the source collection. If ``new_collection`` is None,
+        appends ``_migrated`` to the source name.
+        """

isotrieve/adapters/chroma.py

@@ -0,0 +1,281 @@
+"""ChromaDB adapter for Isotrieve.
+
+Two modes:
+- **Serve mode**: ``IsotrieveChromaFunction`` — a Chroma ``EmbeddingFunction`` that
+  transparently maps new-model queries into legacy space before searching.
+- **Offline migration**: ``migrate_collection()`` — transforms stored vectors
+  into the new model's space and writes to a new collection.
+
+Requires: ``pip install isotrieve[chroma]`` or ``pip install chromadb``.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import time
+from typing import Any
+
+import numpy as np
+
+from isotrieve.adapters.base import MigrationReport
+from isotrieve.mapping.base import Mapping, l2_normalize
+
+
+def _require_chroma():
+    try:
+        import chromadb
+        from chromadb import Documents, EmbeddingFunction, Embeddings
+
+        return chromadb, Documents, EmbeddingFunction, Embeddings
+    except ImportError:
+        raise ImportError(
+            "ChromaDB adapter requires chromadb. Install with: pip install chromadb"
+        )
+
+
+class IsotrieveChromaFunction:
+    """Chroma ``EmbeddingFunction`` that applies an Isotrieve mapping.
+
+    In serve mode, ``__call__`` embeds with the new model then maps into
+    legacy space, so the legacy index is searched without modification.
+
+    Usage::
+
+        from isotrieve.adapters.chroma import IsotrieveChromaFunction
+        from isotrieve.mapping.base import Mapping
+
+        mapping = Mapping.load("ada002_to_te3.isotrieve")
+        ef = IsotrieveChromaFunction(mapping, new_model_embedder=my_embed_fn)
+        col = client.get_collection("docs", embedding_function=ef)
+        results = col.query(query_texts=["..."], n_results=10)
+    """
+
+    def __init__(
+        self,
+        mapping: Mapping,
+        new_model_embedder: Any = None,
+        *,
+        recalibrate: bool = True,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        mapping:
+            Fitted Isotrieve mapping with inverse direction.
+        new_model_embedder:
+            Callable that takes ``list[str]`` and returns ``np.ndarray``
+            of shape ``(n, d_new)``. If None, the EF embeds with the
+            new model (caller must provide this).
+        recalibrate:
+            If True and the mapping has a recalibrator, post-process
+            query scores.
+        """
+        self._mapping = mapping
+        self._embedder = new_model_embedder
+        self._recalibrate = recalibrate
+
+    def __call__(self, input: list[str]) -> list[list[float]]:
+        """Embed texts into legacy space (for querying legacy index)."""
+        if self._embedder is None:
+            raise RuntimeError(
+                "No new_model_embedder provided. Pass a callable that "
+                "embeds text with the new model."
+            )
+        # Embed with new model
+        new_vecs = np.asarray(self._embedder(input), dtype=np.float64)
+        # Map to legacy space
+        legacy_vecs = l2_normalize(self._mapping.inverse_transform(new_vecs))
+        return legacy_vecs.tolist()
+
+    def embed_query(self, input: str) -> list[float]:
+        """Embed a single query into legacy space."""
+        return self.__call__([input])[0]
+
+    @property
+    def has_recalibrator(self) -> bool:
+        return self._mapping.has_recalibrator
+
+    def default_space(self) -> str:
+        return "cosine"
+
+    def name(self) -> str:
+        return "isotrieve_chroma"
+
+    def get_config(self) -> dict[str, Any]:
+        return {
+            "mapping_type": self._mapping.mapping_type,
+            "d_src": self._mapping._d_src,
+            "d_tgt": self._mapping._d_tgt,
+        }
+
+
+def _hash_mapping(mapping: Mapping) -> str:
+    """Short checksum of the mapping matrix for metadata."""
+    if mapping._W is None:
+        return "unfitted"
+    h = hashlib.sha256(mapping._W.tobytes()).hexdigest()[:12]
+    return h
+
+
+def migrate_collection(
+    client: Any,
+    collection_name: str,
+    mapping: Mapping,
+    *,
+    new_collection: str | None = None,
+    batch_size: int = 1000,
+    dry_run: bool = False,
+    new_model_embedder: Any = None,
+) -> MigrationReport:
+    """Migrate a Chroma collection from legacy to new model space.
+
+    Reads all vectors from ``collection_name``, applies the forward
+    mapping (legacy → new), and writes to ``new_collection``.
+    Source collection is never modified.
+
+    Parameters
+    ----------
+    client:
+        ``chromadb.Client`` instance.
+    collection_name:
+        Source collection (legacy embeddings).
+    mapping:
+        Fitted Isotrieve mapping (forward direction: legacy → new).
+    new_collection:
+        Target name. Default: ``{collection_name}_migrated``.
+    batch_size:
+        Rows per read/write batch.
+    dry_run:
+        If True, read and transform but don't write.
+    new_model_embedder:
+        If provided, re-embed texts with new model instead of mapping
+        (for comparison / ground-truth baseline).
+
+    Returns
+    -------
+    MigrationReport with rows processed, timing, and sampled recall.
+    """
+    chromadb, _, _, _ = _require_chroma()
+
+    t0 = time.perf_counter()
+    target_name = new_collection or f"{collection_name}_migrated"
+    report = MigrationReport(
+        source_collection=collection_name,
+        target_collection=target_name,
+        mapping_checksum=_hash_mapping(mapping),
+    )
+
+    # Open source collection
+    src = client.get_collection(collection_name)
+    total = src.count()
+
+    if dry_run:
+        # Just read a batch and report dimensions
+        sample = src.get(
+            limit=min(batch_size, total),
+            include=["embeddings", "metadatas", "documents"],
+        )
+        if sample["embeddings"]:
+            dims = len(sample["embeddings"][0])
+            report.rows_processed = len(sample["embeddings"])
+            report.elapsed_seconds = time.perf_counter() - t0
+            report.errors.append(
+                f"DRY RUN: would migrate {total} vectors of dim {dims}"
+            )
+        return report
+
+    # Check for double-migration: does the source already have isotrieve metadata?
+    sample_meta = src.get(limit=1, include=["metadatas"])
+    if sample_meta["metadatas"] and "isotrieve_mapping_id" in (
+        sample_meta["metadatas"][0] or {}
+    ):
+        report.idempotent = False
+        report.errors.append(
+            "Source collection already has isotrieve_mapping_id metadata. "
+            "This may be a double migration. Proceeding anyway."
+        )
+
+    # Create target collection with correct dimensions
+    # Read first batch to determine dims
+    first_batch = src.get(
+        limit=batch_size, include=["embeddings", "metadatas", "documents"]
+    )
+    if not first_batch["embeddings"]:
+        report.errors.append("Source collection is empty")
+        return report
+
+    # Create target
+    try:
+        client.delete_collection(target_name)
+    except Exception:
+        pass
+    target = client.create_collection(
+        name=target_name,
+        metadata={"hnsw:space": "cosine"},
+    )
+
+    # Process in batches
+    offset = 0
+    rows_written = 0
+    sampled_correct = 0
+    sampled_total = 0
+
+    while offset < total:
+        batch = src.get(
+            offset=offset,
+            limit=batch_size,
+            include=["embeddings", "metadatas", "documents", "ids"],
+        )
+        if not batch["ids"]:
+            break
+
+        ids = batch["ids"]
+        metadatas = batch["metadatas"] or [{} for _ in ids]
+        documents = batch["documents"] or [None for _ in ids]
+        embeddings = np.array(batch["embeddings"], dtype=np.float64)
+
+        # Transform: legacy → new space
+        mapped = mapping.transform(embeddings)
+        mapped_list = mapped.tolist()
+
+        # Add Isotrieve metadata to each row
+        for i, meta in enumerate(metadatas):
+            if meta is None:
+                meta = {}
+                metadatas[i] = meta
+            meta["isotrieve_mapping_id"] = _hash_mapping(mapping)
+            meta["isotrieve_format_version"] = 1
+            meta["isotrieve_source_collection"] = collection_name
+
+        # Write to target
+        target.add(
+            ids=ids,
+            embeddings=mapped_list,
+            metadatas=metadatas,
+            documents=documents,
+        )
+        rows_written += len(ids)
+
+        # Sample recall check (every 10th batch)
+        if rows_written % (batch_size * 10) == 0 or rows_written >= total:
+            sample_n = min(50, len(ids))
+            src_sample = embeddings[:sample_n]
+            tgt_sample = mapped[:sample_n]
+            # Recall@10: for each mapped vector, how many of its top-10
+            # nearest in mapped space match its top-10 in source space
+            src_sims = l2_normalize(src_sample) @ l2_normalize(src_sample).T
+            tgt_sims = l2_normalize(tgt_sample) @ l2_normalize(tgt_sample).T
+            for i in range(sample_n):
+                src_top10 = set(np.argsort(-src_sims[i])[:10])
+                tgt_top10 = set(np.argsort(-tgt_sims[i])[:10])
+                sampled_correct += len(src_top10 & tgt_top10)
+                sampled_total += 10
+
+        offset += batch_size
+
+    report.rows_processed = rows_written
+    report.elapsed_seconds = time.perf_counter() - t0
+    if sampled_total > 0:
+        report.sampled_recall_at_10 = sampled_correct / sampled_total
+
+    return report

isotrieve/adapters/langchain.py

@@ -0,0 +1,117 @@
+"""LangChain embeddings adapter for Isotrieve.
+
+``IsotrieveEmbeddings`` is a drop-in ``langchain_core.embeddings.Embeddings``
+shim that transparently applies an Isotrieve mapping. No index migration
+needed — query vectors are mapped on-the-fly.
+
+Requires: ``pip install isotrieve[langchain]`` or ``pip install langchain-core``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from isotrieve.mapping.base import Mapping, l2_normalize
+
+
+def _require_langchain():
+    try:
+        from langchain_core.embeddings import Embeddings
+
+        return Embeddings
+    except ImportError:
+        raise ImportError(
+            "LangChain adapter requires langchain-core. "
+            "Install with: pip install langchain-core"
+        )
+
+
+class IsotrieveEmbeddings:
+    """Drop-in ``Embeddings`` that applies an Isotrieve mapping.
+
+    Usage::
+
+        from langchain_chroma import Chroma
+        from isotrieve.adapters.langchain import IsotrieveEmbeddings
+        from isotrieve.mapping.base import Mapping
+
+        mapping = Mapping.load("ada002_to_te3.isotrieve")
+        base_embed = OpenAIEmbeddings(model="text-embedding-3-small")
+        ae = IsotrieveEmbeddings(mapping, base_embed)
+
+        # Works with any LangChain vector store
+        db = Chroma.from_documents(docs, embedding=ae)
+        results = db.similarity_search("query", k=10)
+    """
+
+    def __init__(
+        self,
+        mapping: Mapping,
+        base_embeddings: Any,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        mapping:
+            Fitted Isotrieve mapping. Forward direction should be
+            ``legacy → new`` so that ``mapping.inverse_transform()``
+            maps new-model vectors back to legacy space.
+        base_embeddings:
+            A ``langchain_core.embeddings.Embeddings`` instance for the
+            new model (e.g. ``OpenAIEmbeddings``).
+        """
+        _require_langchain()  # side effect: raises ImportError if langchain missing
+        # Validate that base_embeddings satisfies the interface
+        if not hasattr(base_embeddings, "embed_documents"):
+            raise TypeError("base_embeddings must implement embed_documents()")
+        self._mapping = mapping
+        self._base = base_embeddings
+
+    def _map_vectors(self, vecs: np.ndarray) -> np.ndarray:
+        """Map new-model vectors to legacy space."""
+        return l2_normalize(self._mapping.inverse_transform(vecs))
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed documents into legacy space (for legacy index)."""
+        new_vecs = np.asarray(self._base.embed_documents(texts), dtype=np.float64)
+        legacy_vecs = self._map_vectors(new_vecs)
+        return legacy_vecs.tolist()
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed a single query into legacy space."""
+        new_vec = np.asarray(self._base.embed_query(text), dtype=np.float64)
+        legacy_vec = self._map_vectors(new_vec.reshape(1, -1)).ravel()
+        return legacy_vec.tolist()
+
+    # Async variants — delegate to base + map
+    async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
+        if hasattr(self._base, "aembed_documents"):
+            new_vecs = np.asarray(
+                await self._base.aembed_documents(texts), dtype=np.float64
+            )
+        else:
+            new_vecs = np.asarray(self._base.embed_documents(texts), dtype=np.float64)
+        legacy_vecs = self._map_vectors(new_vecs)
+        return legacy_vecs.tolist()
+
+    async def aembed_query(self, text: str) -> list[float]:
+        if hasattr(self._base, "aembed_query"):
+            new_vec = np.asarray(await self._base.aembed_query(text), dtype=np.float64)
+        else:
+            new_vec = np.asarray(self._base.embed_query(text), dtype=np.float64)
+        legacy_vec = self._map_vectors(new_vec.reshape(1, -1)).ravel()
+        return legacy_vec.tolist()
+
+    @property
+    def has_recalibrator(self) -> bool:
+        return self._mapping.has_recalibrator
+
+    def __repr__(self) -> str:
+        base_cls = type(self._base).__name__
+        return (
+            f"IsotrieveEmbeddings(base={base_cls}, "
+            f"mapping={self._mapping.mapping_type}, "
+            f"d_src={self._mapping._d_src}, d_tgt={self._mapping._d_tgt})"
+        )