"""LangChain VectorStore backed by turbovec's quantized index. Install with: ``pip install turbovec[langchain]``. The public surface mirrors langchain_core's in-tree ``InMemoryVectorStore`` so this store can be swapped in wherever the in-memory store is used. """ from __future__ import annotations import json import uuid from pathlib import Path from typing import Any, Callable, Iterable, Sequence import numpy as np from ._dedup import DuplicatePolicy, resolve_duplicates from ._persist import check_persisted_handles from ._turbovec import IdMapIndex try: from langchain_core.documents import Document from langchain_core.embeddings import Embeddings from langchain_core.vectorstores import VectorStore except ImportError as exc: raise ImportError( "langchain-core is required to use turbovec.langchain. " "Install with: pip install turbovec[langchain]" ) from exc _INDEX_FILENAME = "index.tvim" _STORE_FILENAME = "docstore.json" # Bump when the docstore.json shape changes; loader refuses to deserialize # unknown versions. _DOCSTORE_SCHEMA_VERSION = 1 class TurboQuantVectorStore(VectorStore): """LangChain VectorStore backed by a :class:`IdMapIndex`. Vectors are quantized to 2–4 bits per dimension. A side-car dictionary holds the original text and metadata keyed by document id. Deletion is supported in O(1) per id via the underlying :class:`IdMapIndex`. """ def __init__( self, embedding: Embeddings, index: IdMapIndex | None = None, *, bit_width: int = 4, docs: dict[str, tuple[str, dict[str, Any]]] | None = None, str_to_u64: dict[str, int] | None = None, next_u64: int = 0, ) -> None: """ :param embedding: LangChain ``Embeddings`` instance used to encode documents and queries. :param index: Optional pre-built :class:`IdMapIndex`. When omitted, a lazy ``IdMapIndex`` is created — it commits to a dim on the first add and lets us match the no-arg constructor pattern of langchain_core's ``InMemoryVectorStore``. :param bit_width: Quantization width (2 or 4) used when the index is created from scratch. Ignored if ``index`` is supplied. """ self._embedding = embedding # IdMapIndex itself supports lazy construction now — no per-store # lazy wrapping needed. When `index` is None we create a lazy # IdMapIndex(dim=None, bit_width) and let it handle the rest. self._index = index if index is not None else IdMapIndex(bit_width=bit_width) self._docs: dict[str, tuple[str, dict[str, Any]]] = docs if docs is not None else {} self._str_to_u64: dict[str, int] = str_to_u64 if str_to_u64 is not None else {} # Reverse map (u64 handle → str id) kept in sync so search results # can translate handles back to LangChain document ids. self._u64_to_str: dict[int, str] = { handle: sid for sid, handle in self._str_to_u64.items() } self._next_u64: int = next_u64 def _issue_handle(self) -> int: self._next_u64 += 1 return self._next_u64 @property def embeddings(self) -> Embeddings: return self._embedding # ---- Relevance score normalization -------------------------------- def _select_relevance_score_fn(self) -> Callable[[float], float]: # turbovec returns the raw inner product of unit-normalized vectors — # ideally cosine similarity in [-1, 1]. Quantization noise can # push that very slightly outside the bounds, so clamp after # mapping to LangChain's [0, 1] relevance scale via (sim + 1) / 2. return lambda sim: max(0.0, min(1.0, (sim + 1.0) / 2.0)) # ---- Write path --------------------------------------------------- def add_texts( self, texts: Iterable[str], metadatas: list[dict] | None = None, ids: list[str] | None = None, **_: Any, ) -> list[str]: texts_list = list(texts) if not texts_list: return [] if metadatas is None: metadatas = [{} for _ in texts_list] if ids is None: ids = [str(uuid.uuid4()) for _ in texts_list] if len(metadatas) != len(texts_list) or len(ids) != len(texts_list): raise ValueError("texts, metadatas, and ids must all have the same length") vectors = np.asarray(self._embedding.embed_documents(texts_list), dtype=np.float32) return self._store_texts_and_vectors(texts_list, vectors, metadatas, ids) async def aadd_texts( self, texts: Iterable[str], metadatas: list[dict] | None = None, ids: list[str] | None = None, **_: Any, ) -> list[str]: texts_list = list(texts) if not texts_list: return [] if metadatas is None: metadatas = [{} for _ in texts_list] if ids is None: ids = [str(uuid.uuid4()) for _ in texts_list] if len(metadatas) != len(texts_list) or len(ids) != len(texts_list): raise ValueError("texts, metadatas, and ids must all have the same length") vectors = np.asarray( await self._embedding.aembed_documents(texts_list), dtype=np.float32 ) return self._store_texts_and_vectors(texts_list, vectors, metadatas, ids) def add_documents( self, documents: list[Document], ids: list[str] | None = None, **kwargs: Any, ) -> list[str]: # Override the base class default which drops the entire `ids` array # if any Document has a None id. The reference InMemoryVectorStore # falls back per-document so partial ids are honoured. texts = [doc.page_content for doc in documents] metadatas = [doc.metadata for doc in documents] if ids is None: ids = [doc.id or str(uuid.uuid4()) for doc in documents] return self.add_texts(texts=texts, metadatas=metadatas, ids=ids, **kwargs) async def aadd_documents( self, documents: list[Document], ids: list[str] | None = None, **kwargs: Any, ) -> list[str]: texts = [doc.page_content for doc in documents] metadatas = [doc.metadata for doc in documents] if ids is None: ids = [doc.id or str(uuid.uuid4()) for doc in documents] return await self.aadd_texts( texts=texts, metadatas=metadatas, ids=ids, **kwargs ) def _store_texts_and_vectors( self, texts_list: list[str], vectors: np.ndarray, metadatas: list[dict], ids: list[str], ) -> list[str]: if vectors.ndim != 2: raise ValueError(f"expected 2D embedding batch, got {vectors.ndim}D") # Dedup intra-batch duplicate ids, keeping the last occurrence — # matches InMemoryVectorStore, whose dict store silently overwrites # on a repeated id. Without this every row is added to the index but # _str_to_u64 keeps only the last handle per id, orphaning the # earlier vectors. The returned id list still mirrors the input # (one entry per input text), as the reference does. result_ids = ids keep = resolve_duplicates(ids, DuplicatePolicy.KEEP_LAST) if len(keep) != len(ids): ids = [ids[i] for i in keep] texts_list = [texts_list[i] for i in keep] metadatas = [metadatas[i] for i in keep] vectors = vectors[keep] # Validate before mutating any existing data. IdMapIndex.add_with_ids # handles both eager (dim must match) and lazy (locks dim on first # call) cases. Pre-check the eager case so we surface a clean # ValueError rather than a Rust panic. existing_dim = self._index.dim if existing_dim is not None and vectors.shape[1] != existing_dim: raise ValueError( f"embedding dimension {vectors.shape[1]} does not match index dim {existing_dim}" ) if not vectors.flags["C_CONTIGUOUS"]: vectors = np.ascontiguousarray(vectors) handles = np.array( [self._issue_handle() for _ in texts_list], dtype=np.uint64 ) # Add first; if encoding rejects the batch (e.g. non-finite values) # this raises before any existing data is touched. Only once the add # has succeeded do we remove the old vectors for colliding ids, so a # failed upsert never destroys existing data (issue #89). Handles are # freshly issued, so the old and new vectors coexist until the delete. self._index.add_with_ids(vectors, handles) # Upsert: any id that already existed is removed so the re-added # vector wins. Matches LangChain user expectation that `add_texts` # with an existing id updates in place. duplicates = [i for i in ids if i in self._str_to_u64] if duplicates: self.delete(duplicates) for id_, text, meta, handle in zip(ids, texts_list, metadatas, handles): h = int(handle) self._str_to_u64[id_] = h self._u64_to_str[h] = id_ self._docs[id_] = (text, dict(meta)) return result_ids # ---- Read path (similarity search) -------------------------------- def similarity_search( self, query: str, k: int = 4, filter: dict[str, Any] | Callable[[Document], bool] | None = None, **_: Any, ) -> list[Document]: return [ doc for doc, _score in self.similarity_search_with_score(query, k=k, filter=filter) ] async def asimilarity_search( self, query: str, k: int = 4, filter: dict[str, Any] | Callable[[Document], bool] | None = None, **_: Any, ) -> list[Document]: return [ doc for doc, _score in await self.asimilarity_search_with_score( query, k=k, filter=filter ) ] def similarity_search_with_score( self, query: str, k: int = 4, filter: dict[str, Any] | Callable[[Document], bool] | None = None, **_: Any, ) -> list[tuple[Document, float]]: qvec = np.asarray(self._embedding.embed_query(query), dtype=np.float32) return self._search_vector(qvec, k, filter=filter) async def asimilarity_search_with_score( self, query: str, k: int = 4, filter: dict[str, Any] | Callable[[Document], bool] | None = None, **_: Any, ) -> list[tuple[Document, float]]: qvec = np.asarray( await self._embedding.aembed_query(query), dtype=np.float32 ) return self._search_vector(qvec, k, filter=filter) def similarity_search_by_vector( self, embedding: list[float], k: int = 4, filter: dict[str, Any] | Callable[[Document], bool] | None = None, **_: Any, ) -> list[Document]: qvec = np.asarray(embedding, dtype=np.float32) return [doc for doc, _score in self._search_vector(qvec, k, filter=filter)] async def asimilarity_search_by_vector( self, embedding: list[float], k: int = 4, filter: dict[str, Any] | Callable[[Document], bool] | None = None, **_: Any, ) -> list[Document]: # The search itself is sync (no embedding step). Delegate. return self.similarity_search_by_vector(embedding, k=k, filter=filter) def _search_vector( self, qvec: np.ndarray, k: int, filter: dict[str, Any] | Callable[[Document], bool] | None = None, ) -> list[tuple[Document, float]]: if qvec.ndim == 1: qvec = qvec[None, :] if not qvec.flags["C_CONTIGUOUS"]: qvec = np.ascontiguousarray(qvec) # IdMapIndex handles the lazy-uncommitted case internally (returns # empty search results). A len-zero check covers both that and # the eager-but-empty case. if len(self._index) == 0: return [] if filter is None: search_k = min(k, len(self._index)) scores, handles = self._index.search(qvec, search_k) else: predicate = self._compile_filter(filter) allowed_handles = [ self._str_to_u64[sid] for sid, (text, meta) in self._docs.items() if predicate(Document(id=sid, page_content=text, metadata=dict(meta))) ] if not allowed_handles: return [] allowlist = np.asarray(allowed_handles, dtype=np.uint64) scores, handles = self._index.search(qvec, k, allowlist=allowlist) results: list[tuple[Document, float]] = [] for score, handle in zip(scores[0], handles[0]): sid = self._u64_to_str[int(handle)] text, meta = self._docs[sid] results.append( (Document(id=sid, page_content=text, metadata=dict(meta)), float(score)) ) return results @staticmethod def _compile_filter( filter: dict[str, Any] | Callable[[Document], bool], ) -> Callable[[Document], bool]: # Match the in-tree InMemoryVectorStore convention: callable filters # receive a Document, not a metadata dict # (langchain_core/vectorstores/in_memory.py). if callable(filter): return filter if isinstance(filter, dict): items = list(filter.items()) return lambda doc: all(doc.metadata.get(k) == v for k, v in items) raise TypeError( "filter must be a dict of metadata key/value pairs or a callable " f"taking a Document, got {type(filter).__name__}" ) # ---- Max marginal relevance --------------------------------------- # # MMR requires the full-precision vector of every candidate to compute # pairwise diversity scores. turbovec discards full vectors after # quantization (that's the point), so we can't faithfully implement # MMR. Raise loudly with a useful message rather than silently fall # back to the base class's bare NotImplementedError. _MMR_MSG = ( "TurboQuantVectorStore does not support max-marginal-relevance " "search because the underlying quantized index discards " "full-precision vectors after compression. MMR requires the " "original embedding for every candidate to compute pairwise " "diversity. Use `similarity_search` / `similarity_search_with_score` " "instead, or maintain a parallel store with full-precision " "embeddings if you need MMR specifically." ) def max_marginal_relevance_search( self, query: str, k: int = 4, fetch_k: int = 20, lambda_mult: float = 0.5, **kwargs: Any, ) -> list[Document]: raise NotImplementedError(self._MMR_MSG) def max_marginal_relevance_search_by_vector( self, embedding: list[float], k: int = 4, fetch_k: int = 20, lambda_mult: float = 0.5, *, filter: Callable[[Document], bool] | None = None, **kwargs: Any, ) -> list[Document]: raise NotImplementedError(self._MMR_MSG) async def amax_marginal_relevance_search( self, query: str, k: int = 4, fetch_k: int = 20, lambda_mult: float = 0.5, **kwargs: Any, ) -> list[Document]: raise NotImplementedError(self._MMR_MSG) # ---- Get / delete ------------------------------------------------- def get_by_ids(self, ids: Sequence[str], /) -> list[Document]: """Return Documents for the given ids. Missing ids are silently skipped (matches the InMemoryVectorStore reference).""" out: list[Document] = [] for sid in ids: if sid in self._docs: text, meta = self._docs[sid] out.append(Document(id=sid, page_content=text, metadata=dict(meta))) return out async def aget_by_ids(self, ids: Sequence[str], /) -> list[Document]: return self.get_by_ids(ids) def delete(self, ids: list[str] | None = None, **_: Any) -> None: """Remove documents by id. Missing ids are silently skipped — matches the InMemoryVectorStore reference (which also accepts ``ids=None`` as a no-op).""" if not ids: return for sid in ids: handle = self._str_to_u64.pop(sid, None) if handle is None: continue self._u64_to_str.pop(handle, None) self._docs.pop(sid, None) self._index.remove(handle) async def adelete(self, ids: list[str] | None = None, **_: Any) -> None: self.delete(ids) # ---- Construction helpers ----------------------------------------- @classmethod def from_texts( cls, texts: list[str], embedding: Embeddings, metadatas: list[dict] | None = None, *, bit_width: int = 4, ids: list[str] | None = None, **_: Any, ) -> "TurboQuantVectorStore": # The underlying index is created lazily on the first `add_texts` # call, picking up `dim` from the first batch of embeddings — same # no-`dim` ergonomics as InMemoryVectorStore. store = cls(embedding=embedding, bit_width=bit_width) if texts: store.add_texts(texts, metadatas=metadatas, ids=ids) return store @classmethod async def afrom_texts( cls, texts: list[str], embedding: Embeddings, metadatas: list[dict] | None = None, *, bit_width: int = 4, ids: list[str] | None = None, **_: Any, ) -> "TurboQuantVectorStore": store = cls(embedding=embedding, bit_width=bit_width) if texts: await store.aadd_texts(texts, metadatas=metadatas, ids=ids) return store # ---- Persistence -------------------------------------------------- # # Method names match the InMemoryVectorStore reference (`dump`/`load`), # but the on-disk layout is a folder containing the binary index file # plus a JSON side-car (we can't embed the binary Rust index in a # single JSON file the way the reference does with its raw-vector # store). def dump(self, folder_path: str | Path) -> None: """Persist the quantized index plus the side-car to disk. ``folder_path`` is a directory; turbovec writes ``index.tvim`` and ``docstore.json`` inside it. Document metadata must be JSON-serializable (same constraint as ``InMemoryVectorStore``). A lazy uncommitted index encodes its state via the index file's own ``dim = 0`` sentinel; no special-case handling needed here. """ folder = Path(folder_path) folder.mkdir(parents=True, exist_ok=True) self._index.write(str(folder / _INDEX_FILENAME)) # `_docs` stores tuples `(text, metadata)` — JSON would drop the # tuple-ness on round-trip, so serialize each entry as an explicit # `{"text": ..., "metadata": ...}` dict. docs_payload = { sid: {"text": text, "metadata": meta} for sid, (text, meta) in self._docs.items() } payload = { "schema_version": _DOCSTORE_SCHEMA_VERSION, "docs": docs_payload, "str_to_u64": self._str_to_u64, "next_u64": self._next_u64, # Pull bit_width off the live index — same value whether # the index was constructed eagerly or lazily. "bit_width": self._index.bit_width, } with open(folder / _STORE_FILENAME, "w") as f: json.dump(payload, f) @classmethod def load( cls, folder_path: str | Path, embedding: Embeddings, ) -> "TurboQuantVectorStore": """Reload a store from a folder previously written by :meth:`dump`. Safe to call on any path — the side-car is plain JSON, never pickle, so there's no deserialization-of-code risk.""" folder = Path(folder_path) with open(folder / _STORE_FILENAME) as f: state = json.load(f) version = state.get("schema_version", 0) if version != _DOCSTORE_SCHEMA_VERSION: raise ValueError( f"docstore.json has schema version {version}; " f"this turbovec expects version {_DOCSTORE_SCHEMA_VERSION}" ) # IdMapIndex.load handles the dim=0 (lazy-uncommitted) sentinel # internally and reconstructs the index in the right state. index = IdMapIndex.load(str(folder / _INDEX_FILENAME)) # Rehydrate `_docs` from the explicit `{"text", "metadata"}` form # back into the internal tuple representation. docs = { sid: (entry["text"], entry["metadata"]) for sid, entry in state["docs"].items() } # JSON object keys are strings; the str_to_u64 values are already # ints in the payload, just need to confirm. str_to_u64 = {sid: int(h) for sid, h in state["str_to_u64"].items()} check_persisted_handles(index, str_to_u64.values(), what="document") return cls( embedding=embedding, index=index, bit_width=state.get("bit_width", 4), docs=docs, str_to_u64=str_to_u64, next_u64=int(state["next_u64"]), ) __all__ = ["TurboQuantVectorStore"]