wehub-resource-sync
0d29a02500
CI / Rust (ubuntu-latest) (push) Failing after 1s
CI / Python (ubuntu-latest) (push) Failing after 0s
CI / Rust (macos-14) (push) Has been cancelled
CI / Rust (windows-latest) (push) Has been cancelled
CI / Python (macos-14) (push) Has been cancelled
CI / Python (windows-latest) (push) Has been cancelled
679 lines
28 KiB
Python
679 lines
28 KiB
Python
"""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.<sync>(...)`. 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"]
|