"""LlamaIndex VectorStore backed by turbovec's quantized index. Install with: ``pip install turbovec[llama-index]``. """ from __future__ import annotations import json import os from pathlib import Path from typing import Any, List, Optional, Sequence import numpy as np from ._dedup import DuplicatePolicy, resolve_duplicates from ._persist import check_persisted_handles from ._turbovec import IdMapIndex try: from llama_index.core.bridge.pydantic import PrivateAttr from llama_index.core.schema import ( BaseNode, NodeRelationship, RelatedNodeInfo, TextNode, ) from llama_index.core.vector_stores.types import ( BasePydanticVectorStore, FilterCondition, FilterOperator, MetadataFilter, MetadataFilters, VectorStoreQuery, VectorStoreQueryMode, VectorStoreQueryResult, ) from llama_index.core.vector_stores.utils import ( metadata_dict_to_node, node_to_metadata_dict, ) except ImportError as exc: raise ImportError( "llama-index-core is required to use turbovec.llama_index. " "Install with: pip install turbovec[llama-index]" ) from exc # Persistence layout: a single user-facing ``persist_path`` (the file the # framework asks us to write) is split into two files by extension — # ``{base}.tvim`` for the binary IdMapIndex and ``{base}.nodes.json`` for # the node side-car. Both live next to each other so the layout fits the # directory-of-namespaced-files pattern used by StorageContext. _INDEX_EXT = ".tvim" _STORE_EXT = ".nodes.json" # Filename template used by SimpleVectorStore for namespace lookup — # mirrored here so `from_persist_dir` works the same way. _NAMESPACE_SEP = "__" _DEFAULT_PERSIST_FNAME = "vector_store.json" _DEFAULT_VECTOR_STORE = "default" # Bump when the nodes.json shape changes; loader accepts the current # version plus any older versions whose missing fields we know how to # reconstruct (currently v1, written before full-node round-trip was # added — v1 entries are reconstructed as bare TextNodes with only # text + metadata + SOURCE relationship, matching the original # lossy behaviour rather than failing to load). _NODES_SCHEMA_VERSION = 2 _NODES_SCHEMA_COMPAT = (1, 2) def _split_persist_base(persist_path: str | Path) -> Path: """Strip the framework-provided extension off `persist_path` so the binary index and JSON side-car can sit next to each other under a shared base. We then append our own extensions in persist / load.""" p = Path(persist_path) # Use the path without its suffix so both .tvim and .nodes.json share # a base. If the input has no suffix (e.g. a bare folder-like name), # use it as-is. return p.with_suffix("") if p.suffix else p class TurboQuantVectorStore(BasePydanticVectorStore): """LlamaIndex VectorStore backed by a :class:`IdMapIndex`. Vectors are quantized to 2–4 bits per dimension. A side-car dictionary holds node text and metadata keyed by ``node_id``. Supports ``delete`` (by ``ref_doc_id``, removing every node with that ref) and ``delete_nodes`` (by ``node_id``) — both O(1) per node. """ stores_text: bool = True is_embedding_query: bool = True flat_metadata: bool = False _index: Any = PrivateAttr() _nodes: dict[str, dict[str, Any]] = PrivateAttr() _node_id_to_u64: dict[str, int] = PrivateAttr() _u64_to_node_id: dict[int, str] = PrivateAttr() _next_u64: int = PrivateAttr() def __init__(self, index: IdMapIndex | None = None, *, bit_width: int = 4, **kwargs: Any) -> None: """Construct the vector store. :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 callers use the no-arg construction pattern common to LlamaIndex's other vector stores (e.g. via ``StorageContext.from_defaults(vector_store=TurboQuantVectorStore())``). :param bit_width: Quantization width used when constructing the lazy index. Ignored if ``index`` is supplied. """ super().__init__(**kwargs) # IdMapIndex itself supports lazy construction now — no per-store # lazy wrapping needed. self._index = index if index is not None else IdMapIndex(bit_width=bit_width) self._nodes = {} self._node_id_to_u64 = {} self._u64_to_node_id = {} self._next_u64 = 0 def _issue_handle(self) -> int: self._next_u64 += 1 return self._next_u64 @classmethod def class_name(cls) -> str: return "TurboQuantVectorStore" @classmethod def from_params(cls, dim: int | None = None, bit_width: int = 4) -> "TurboQuantVectorStore": """Build a store with a known ``dim`` (eager) or lazy when ``dim`` is omitted.""" return cls(index=IdMapIndex(dim, bit_width)) @property def client(self) -> IdMapIndex: return self._index def add(self, nodes: list[BaseNode], **_: Any) -> list[str]: if not nodes: return [] # Reject intra-batch duplicates loudly. Letting them through would # leave the index with N vectors but only the last node_id mapped # back to one of them — the earlier handles become orphans that # `query` later resolves through the duplicate node_id, returning # the second node's payload attached to the first node's vector. # Caller's job to deduplicate before calling add. node_ids = [n.node_id for n in nodes] try: resolve_duplicates(node_ids, DuplicatePolicy.REJECT) except ValueError: seen: set[str] = set() dup = next(nid for nid in node_ids if nid in seen or seen.add(nid)) raise ValueError( f"duplicate node_id {dup!r} appears multiple times " "in the input batch; deduplicate before calling add()" ) from None embeddings = [node.get_embedding() for node in nodes] vectors = np.asarray(embeddings, dtype=np.float32) if vectors.ndim != 2: raise ValueError( f"expected 2D embedding batch, got {vectors.ndim}D" ) # IdMapIndex.add_with_ids handles eager (dim must match) and lazy # (locks dim on first add) — 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"node embedding dim {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 nodes], dtype=np.uint64) # Add first; if validation above or encoding here (e.g. non-finite # values) rejects the batch, it raises before any existing data is # touched. Only after the add succeeds do we remove the old entries # for colliding node_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-like: if a node_id is already present in the STORE, remove # the old entry so the new embedding wins. duplicates = [n.node_id for n in nodes if n.node_id in self._node_id_to_u64] for node_id in duplicates: self._remove_node_by_id(node_id) ids: list[str] = [] for node, handle in zip(nodes, handles): h = int(handle) nid = node.node_id self._node_id_to_u64[nid] = h self._u64_to_node_id[h] = nid # `metadata` and `ref_doc_id` are kept at top level for fast # filter / doc-id lookup (queries hit these on every hit; # parsing _node_content per hit would be wasteful). `node_dict` # is the framework's canonical metadata representation # (`_node_content` + `_node_type` + original metadata keys), # which `metadata_dict_to_node` reconstructs into a full # BaseNode — preserving relationships (PREVIOUS / NEXT / # PARENT / CHILD), excluded_*_metadata_keys, template fields, # start/end_char_idx and mimetype on retrieval. The narrow # `{text, metadata, ref_doc_id}` schema we used to keep lost # all of those silently. self._nodes[nid] = { "metadata": dict(node.metadata), "ref_doc_id": node.ref_doc_id, "node_dict": node_to_metadata_dict( node, remove_text=False, flat_metadata=False ), } ids.append(nid) return ids def delete(self, ref_doc_id: str, **_: Any) -> None: """Delete every node whose ``ref_doc_id`` matches.""" matching = [ nid for nid, data in self._nodes.items() if data.get("ref_doc_id") == ref_doc_id ] for nid in matching: self._remove_node_by_id(nid) def delete_nodes( self, node_ids: Optional[List[str]] = None, filters: Optional[MetadataFilters] = None, **_: Any, ) -> None: """Delete every node matching ``node_ids`` and/or ``filters``. Both constraints intersect when supplied. Missing node_ids are ignored. Matches the signature and semantics of ``SimpleVectorStore.delete_nodes``. """ if not node_ids and filters is None: return candidates = list(self._nodes.items()) if node_ids is not None: node_id_set = set(node_ids) candidates = [(nid, data) for nid, data in candidates if nid in node_id_set] if filters is not None: candidates = [ (nid, data) for nid, data in candidates if self._filters_match(data["metadata"], filters) ] for nid, _data in candidates: self._remove_node_by_id(nid) def clear(self) -> None: """Drop every node from the store and reset to a fresh lazy index. The new index keeps the same ``bit_width`` so subsequent adds commit a new ``dim`` lazily. """ bw = self._index.bit_width self._index = IdMapIndex(bit_width=bw) self._nodes = {} self._node_id_to_u64 = {} self._u64_to_node_id = {} self._next_u64 = 0 def get(self, text_id: str) -> List[float]: """LlamaIndex's protocol expects this to return the full-precision embedding for a given node id. turbovec discards full-precision embeddings after quantization, so we raise loudly with an explanation rather than return a lossy reconstruction or zeroes. """ raise NotImplementedError( "TurboQuantVectorStore.get(text_id) cannot return the original " "embedding because turbovec quantizes vectors to 2-4 bits per " "dimension and discards full precision after encoding. Keep a " "parallel docstore if you need the raw embedding." ) def get_nodes( self, node_ids: Optional[List[str]] = None, filters: Optional[MetadataFilters] = None, ) -> List[BaseNode]: """Return the nodes matching ``node_ids`` and/or ``filters``. Both constraints intersect when supplied; missing node_ids are silently skipped. Unlike ``SimpleVectorStore`` (which raises NotImplementedError here because it doesn't store nodes), turbovec keeps node text and metadata in a side-car so this can return populated ``TextNode`` objects directly. """ candidates = list(self._nodes.items()) if node_ids is not None: node_id_set = set(node_ids) candidates = [(nid, data) for nid, data in candidates if nid in node_id_set] if filters is not None: candidates = [ (nid, data) for nid, data in candidates if self._filters_match(data["metadata"], filters) ] return [self._reconstruct_node(nid, data) for nid, data in candidates] @staticmethod def _reconstruct_node(nid: str, data: dict[str, Any]) -> BaseNode: # v2 entries carry `node_dict` — round-trip via the framework's # own helper so we get the full BaseNode subclass back # (TextNode / IndexNode / ImageNode) with every field populated. if "node_dict" in data: return metadata_dict_to_node(data["node_dict"]) # v1 fallback: stores that were persisted before the full-node # round-trip landed only have {text, metadata, ref_doc_id}. # Reconstruct the minimum-fidelity TextNode they used to produce # so old on-disk stores keep loading without manual migration. node = TextNode( id_=nid, text=data["text"], metadata=dict(data["metadata"]), ) if data.get("ref_doc_id") is not None: node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo( node_id=data["ref_doc_id"] ) return node def _remove_node_by_id(self, node_id: str) -> bool: handle = self._node_id_to_u64.pop(node_id, None) if handle is None: return False self._u64_to_node_id.pop(handle, None) self._nodes.pop(node_id, None) self._index.remove(handle) return True def _resolve_allowed_handles( self, filters: MetadataFilters | None, node_ids: list[str] | None, doc_ids: list[str] | None, ) -> list[int]: """Resolve ``query.filters``, ``query.node_ids`` and ``query.doc_ids`` to the list of internal u64 handles that satisfy the filter. Empty list means no node matches. Semantics (matching the SimpleVectorStore reference where applicable): - ``node_ids``: filter by node_id (set membership). - ``doc_ids``: filter by ``ref_doc_id`` only (source document id). - ``filters``: apply metadata filters. All three intersect when more than one is supplied. """ candidates = self._nodes.items() if node_ids: node_id_set = set(node_ids) candidates = [(nid, data) for nid, data in candidates if nid in node_id_set] if doc_ids: doc_id_set = set(doc_ids) candidates = [ (nid, data) for nid, data in candidates if data.get("ref_doc_id") in doc_id_set ] if filters is None: return [self._node_id_to_u64[nid] for nid, _ in candidates] return [ self._node_id_to_u64[nid] for nid, data in candidates if self._filters_match(data["metadata"], filters) ] @classmethod def _filters_match( cls, metadata: dict[str, Any], filters: MetadataFilters ) -> bool: condition = getattr(filters, "condition", None) or FilterCondition.AND results: list[bool] = [] for f in filters.filters: if isinstance(f, MetadataFilters): results.append(cls._filters_match(metadata, f)) else: results.append(cls._single_filter_match(metadata, f)) if condition == FilterCondition.AND: return all(results) if results else True if condition == FilterCondition.OR: return any(results) if results else True if condition == FilterCondition.NOT: # Reference semantics (`build_metadata_filter_fn`, # `utils.py:187-189`): NOT matches when none of the inner # filters match. Empty inner list trivially satisfies NOT. return not any(results) raise NotImplementedError( f"filter condition {condition!r} not supported by TurboQuantVectorStore" ) @staticmethod def _single_filter_match(metadata: dict[str, Any], f: MetadataFilter) -> bool: # Semantics mirror SimpleVectorStore's _build_metadata_filter_fn # (llama_index/core/vector_stores/simple.py) so that filtered # results agree with the in-tree reference store. op = f.operator target = f.value value = metadata.get(f.key) # IS_EMPTY is the only operator that treats a missing key as a hit. if op == FilterOperator.IS_EMPTY: return value is None or value == "" or value == [] # Every other operator returns False when the key is absent — this # matches the reference implementation (notably NE returns False on # missing, not True). if value is None: return False if op == FilterOperator.EQ: return value == target if op == FilterOperator.NE: return value != target if op == FilterOperator.GT: return value > target if op == FilterOperator.LT: return value < target if op == FilterOperator.GTE: return value >= target if op == FilterOperator.LTE: return value <= target if op == FilterOperator.IN: return value in target if op == FilterOperator.NIN: return value not in target if op == FilterOperator.CONTAINS: return target in value if op == FilterOperator.TEXT_MATCH: # Reference (`utils.py:138-144`): case-SENSITIVE substring, # both sides must be strings. Previous turbovec impl # lowercased both sides — a silent semantic divergence that # caused our results to disagree with SimpleVectorStore on # mixed-case keys. if isinstance(target, str) and isinstance(value, str): return target in value raise TypeError( "Both metadata value and filter value must be strings " "for the TEXT_MATCH operator" ) if op == FilterOperator.TEXT_MATCH_INSENSITIVE: if isinstance(target, str) and isinstance(value, str): return target.lower() in value.lower() raise TypeError( "Both metadata value and filter value must be strings " "for the TEXT_MATCH_INSENSITIVE operator" ) if op == FilterOperator.ALL: # Reference (`utils.py:152-153`): every element of `target` # must be present in the metadata value (which is typically # a list — tag-set matching). return all(t in value for t in target) if op == FilterOperator.ANY: return any(t in value for t in target) raise NotImplementedError( f"filter operator {op!r} not supported by TurboQuantVectorStore" ) def query(self, query: VectorStoreQuery, **_: Any) -> VectorStoreQueryResult: # MMR / SVM / LINEAR_REGRESSION / HYBRID etc. all need access to # full-precision vectors (for pairwise diversity, learned scoring, # or sparse-dense fusion). turbovec discards full precision after # quantization, so any non-DEFAULT mode is unsupportable here. # Raise loudly instead of silently treating it as DEFAULT, which # the previous impl did and which let callers think they were # getting e.g. MMR diversity when they were not. if query.mode != VectorStoreQueryMode.DEFAULT: raise NotImplementedError( f"TurboQuantVectorStore does not support query mode " f"{query.mode!r}. Only VectorStoreQueryMode.DEFAULT is " "supported — MMR / SVM / hybrid modes need access to " "full-precision vectors which turbovec discards after " "quantization. Maintain a parallel store with full vectors " "if you need a non-default scoring mode." ) if query.query_embedding is None: raise ValueError( "TurboQuantVectorStore requires a pre-computed query_embedding " "(is_embedding_query=True)." ) qvec = np.asarray(query.query_embedding, dtype=np.float32) if qvec.ndim == 1: qvec = qvec[None, :] if not qvec.flags["C_CONTIGUOUS"]: qvec = np.ascontiguousarray(qvec) if len(self._index) == 0: return VectorStoreQueryResult(nodes=[], similarities=[], ids=[]) has_filters = ( query.filters is not None or bool(query.node_ids) or bool(query.doc_ids) ) if not has_filters: k = min(query.similarity_top_k, len(self._index)) scores, handles = self._index.search(qvec, k) else: allowed_handles = self._resolve_allowed_handles( query.filters, query.node_ids, query.doc_ids ) if not allowed_handles: return VectorStoreQueryResult(nodes=[], similarities=[], ids=[]) allowlist = np.asarray(allowed_handles, dtype=np.uint64) scores, handles = self._index.search( qvec, query.similarity_top_k, allowlist=allowlist ) result_nodes: list[TextNode] = [] similarities: list[float] = [] ids: list[str] = [] for score, handle in zip(scores[0], handles[0]): nid = self._u64_to_node_id[int(handle)] data = self._nodes[nid] result_nodes.append(self._reconstruct_node(nid, data)) similarities.append(float(score)) ids.append(nid) return VectorStoreQueryResult(nodes=result_nodes, similarities=similarities, ids=ids) # ---- Async overrides -------------------------------------------------- # # The base class provides default async impls that delegate to sync via # `return self.(...)`. We override them explicitly so the signature # is visible on the class and an autodoc tool / IDE doesn't make # callers chase the abstract base class for the documentation. async def async_add( self, nodes: Sequence[BaseNode], **kwargs: Any ) -> List[str]: return self.add(list(nodes), **kwargs) async def adelete(self, ref_doc_id: str, **kwargs: Any) -> None: self.delete(ref_doc_id, **kwargs) async def adelete_nodes( self, node_ids: Optional[List[str]] = None, filters: Optional[MetadataFilters] = None, **kwargs: Any, ) -> None: self.delete_nodes(node_ids=node_ids, filters=filters, **kwargs) async def aclear(self) -> None: self.clear() async def aquery( self, query: VectorStoreQuery, **kwargs: Any ) -> VectorStoreQueryResult: return self.query(query, **kwargs) async def aget_nodes( self, node_ids: Optional[List[str]] = None, filters: Optional[MetadataFilters] = None, ) -> List[BaseNode]: return self.get_nodes(node_ids=node_ids, filters=filters) # ---- Config serialization --------------------------------------------- def to_dict(self, **_: Any) -> dict[str, Any]: """Serialize the store's *configuration* (not its data) so a fresh instance can be reconstructed via ``from_dict``. Mirrors the contract of ``SimpleVectorStore.to_dict`` — config-only; node data round-trips through ``persist`` / ``from_persist_path``. """ return { "bit_width": self._index.bit_width, "dim": self._index.dim, # may be None (lazy uncommitted) } @classmethod def from_dict(cls, data: dict[str, Any], **_: Any) -> "TurboQuantVectorStore": """Construct an empty store from a config dict produced by ``to_dict``. To restore data, use ``from_persist_path``.""" dim = data.get("dim") bit_width = data.get("bit_width", 4) return cls(index=IdMapIndex(dim, bit_width)) def persist(self, persist_path: str, fs: Any = None) -> None: """Persist the store. ``persist_path`` is treated as a path *stem*: the binary index goes to ``{stem}.tvim`` and the node side-car to ``{stem}.nodes.json``. Any extension on ``persist_path`` (e.g. ``.json`` from a StorageContext default) is replaced. This matches the layout assumed by ``StorageContext.persist`` — which calls us with ``persist_path = {persist_dir}/{namespace}__vector_store.json`` — and lets multiple namespaced stores coexist in the same directory. Node metadata must be JSON-serializable (same constraint as ``SimpleVectorStore``). ``fs`` (fsspec) is not yet supported; pass a local path. """ if fs is not None: raise NotImplementedError( "fsspec filesystems are not supported yet; pass a local path." ) base = _split_persist_base(persist_path) base.parent.mkdir(parents=True, exist_ok=True) self._index.write(str(base.with_suffix(_INDEX_EXT))) payload = { "schema_version": _NODES_SCHEMA_VERSION, "nodes": self._nodes, # JSON object keys must be strings; round-trip int keys via # an explicit list of [node_id, handle] pairs to preserve # type fidelity. "node_id_to_u64": list(self._node_id_to_u64.items()), "next_u64": self._next_u64, } with open(base.with_suffix(_STORE_EXT), "w") as f: json.dump(payload, f) @classmethod def from_persist_path( cls, persist_path: str, fs: Any = None, ) -> "TurboQuantVectorStore": """Load a previously-persisted store. ``persist_path`` is the same path that was passed to :meth:`persist` (extension is ignored; ``{stem}.tvim`` and ``{stem}.nodes.json`` are read). Safe to call on any path — the side-car is plain JSON, never pickle, so there's no deserialization-of-code risk. """ if fs is not None: raise NotImplementedError( "fsspec filesystems are not supported yet; pass a local path." ) base = _split_persist_base(persist_path) index = IdMapIndex.load(str(base.with_suffix(_INDEX_EXT))) with open(base.with_suffix(_STORE_EXT)) as f: state = json.load(f) version = state.get("schema_version", 0) if version not in _NODES_SCHEMA_COMPAT: raise ValueError( f"{_STORE_EXT.lstrip('.')} has schema version {version}; " f"this turbovec accepts versions {list(_NODES_SCHEMA_COMPAT)}" ) store = cls(index=index) # v1 entries lack `node_dict` and reconstruct as narrow TextNodes; # v2 entries carry it and reconstruct with full BaseNode fidelity. # `_reconstruct_node` dispatches on shape, so we just load the # dict as-is. store._nodes = state["nodes"] # Reconstruct {node_id: int handle} from the list-of-pairs form. store._node_id_to_u64 = {nid: int(h) for nid, h in state["node_id_to_u64"]} store._u64_to_node_id = {h: nid for nid, h in store._node_id_to_u64.items()} store._next_u64 = int(state["next_u64"]) check_persisted_handles(index, store._u64_to_node_id.keys(), what="node") return store @classmethod def from_persist_dir( cls, persist_dir: str, namespace: str = _DEFAULT_VECTOR_STORE, fs: Any = None, ) -> "TurboQuantVectorStore": """Load a store from a ``StorageContext``-style persist directory. Builds the namespaced filename ``{persist_dir}/{namespace}__vector_store.json`` and forwards to :meth:`from_persist_path`. The ``.json`` suffix is conventional — our actual on-disk files use ``.tvim`` and ``.nodes.json`` extensions derived from the same stem. """ persist_fname = f"{namespace}{_NAMESPACE_SEP}{_DEFAULT_PERSIST_FNAME}" persist_path = os.path.join(persist_dir, persist_fname) return cls.from_persist_path(persist_path, fs=fs) __all__ = ["TurboQuantVectorStore"]