"""turbovec backend: the dynamic-corpus specialist.
turbovec (Rust + PyO3, ships wheels on PyPI incl. Apple Silicon) is the engine
the router reaches for when the corpus **changes constantly**: it supports O(1)
``add_with_ids`` and ``remove(id)`` with no index rebuild, and its TurboQuant
2-4 bit quantisation gives ~16x compression while keeping recall above
FAISS-PQ. That combination (mutable *and* compact *and* fast on Apple Silicon)
is what earns it the "frequent updates" branch of the policy, ahead of the
graph indexes whose deletes rot the structure.
This is the extracted, packaged form of the brute-force → turbovec routing that
already ships inside the ``roitelet`` prototype's ``core/personal.py``.
Consumes: ``turbovec`` (optional, ``pip install 'ann-router[turbovec]'``).
Produces: :class:`TurboVecIndex`.
Author: Warith Harchaoui <warith.harchaoui@deraison.ai>
"""
from __future__ import annotations
import numpy as np
from ..base import ANNIndex, BackendUnavailable, Capabilities
from ..spec import MetricName
def _require():
"""Import turbovec or raise an actionable :class:`BackendUnavailable`."""
try:
import turbovec # local import so `import ann_router` never needs the Rust wheel
except ImportError as exc: # pragma: no cover - exercised only when absent
raise BackendUnavailable(
"turbovec not installed. Run: pip install 'ann-router[turbovec]'"
) from exc
return turbovec
[docs]
class TurboVecIndex(ANNIndex):
"""turbovec ``IdMapIndex``: mutable, quantised, id-native.
turbovec is id-native (``add_with_ids`` / ``remove(id)``) and returns
``(distances, ids)`` from ``search`` — this adapter flips that to the
package's ``(ids, distances)`` order. Vectors are normalised for cosine so
the quantiser's inner product matches cosine similarity.
Parameters
----------
dim : int
Embedding dimensionality.
metric : {"cosine", "l2", "ip"}, optional
Distance metric. Defaults to ``"cosine"``. turbovec is inner-product /
cosine oriented; L2 is approximated on normalised vectors.
bit_width : int, optional
TurboQuant bit width (2 or 4). Defaults to 4 — the recall/size sweet
spot measured in the roitelet study.
Examples
--------
>>> TurboVecIndex.capabilities().supports_remove
True
"""
def __init__(self, dim: int, metric: MetricName = "cosine", **kwargs: object) -> None:
super().__init__(dim, metric, **kwargs)
self._bit_width = int(kwargs.get("bit_width", 4))
# Cosine == inner product on unit vectors; normalise on the way in.
self._normalise = metric in ("cosine", "l2")
[docs]
@classmethod
def capabilities(cls) -> Capabilities:
"""Return the turbovec capability descriptor (fully mutable)."""
return Capabilities(
name="turbovec",
supports_add=True,
supports_remove=True, # O(1) remove(id), the reason it exists
supports_filter=False, # allowlist exists but is not metadata filtering
persistent=False,
needs_gpu=False,
approximate=True,
metrics=("cosine", "ip", "l2"),
pip_extra="turbovec",
)
[docs]
@classmethod
def is_available(cls) -> bool:
"""Return ``True`` if turbovec is importable.
Examples
--------
>>> isinstance(TurboVecIndex.is_available(), bool)
True
"""
try:
import turbovec # noqa: F401
except ImportError:
return False
return True
def _prep(self, vectors: np.ndarray) -> np.ndarray:
"""Coerce to float32 and L2-normalise when the metric needs it.
Parameters
----------
vectors : numpy.ndarray
Any 2-D array of vectors.
Returns
-------
numpy.ndarray
Contiguous float32, L2-normalised per row for cosine/l2 metrics.
"""
arr = self._as_f32(vectors)
if self._normalise:
norms = np.linalg.norm(arr, axis=1, keepdims=True)
norms[norms == 0] = 1.0
arr = arr / norms
return arr
[docs]
def build(self, vectors: np.ndarray, ids: np.ndarray | None = None) -> TurboVecIndex:
"""Create the index and insert the initial corpus.
Parameters
----------
vectors : numpy.ndarray
Shape ``(n, dim)``.
ids : numpy.ndarray, optional
Shape ``(n,)``; defaults to ``range(n)``.
Returns
-------
TurboVecIndex
``self``.
"""
turbovec = _require()
self._index = turbovec.IdMapIndex(dim=self.dim, bit_width=self._bit_width)
arr = self._prep(vectors)
labels = np.arange(arr.shape[0]) if ids is None else np.asarray(ids)
self._index.add_with_ids(arr, labels.astype(np.uint64))
# Seed the high-water mark from the ids just inserted, so a later bare
# add() (no explicit ids) continues past them instead of restarting at
# 0 and colliding with the ids build() just assigned.
self._next_id = int(labels.max()) + 1 if labels.size else 0
return self
[docs]
def add(self, vectors: np.ndarray) -> None:
"""Append vectors with the next contiguous ids.
Parameters
----------
vectors : numpy.ndarray
Shape ``(m, dim)``.
Raises
------
RuntimeError
If called on an index that was populated via :meth:`load` with no
subsequent :meth:`build`/:meth:`add_with_ids` call in this
process — the high-water mark can't be recovered from disk (see
:meth:`load`'s docstring), so guessing ``0`` here would silently
collide with ids already on disk instead of failing loudly. Call
:meth:`add_with_ids` with explicit ids instead.
"""
# turbovec has no count getter we rely on, so track the high-water mark
# from the corpus we have inserted so far via a private counter. A
# freshly constructed, never-built index has no attribute at all
# (safe to start at 0); `load()` sets it to the None *sentinel*
# (unsafe to guess) -- these are deliberately not the same case.
if self._index is not None and getattr(self, "_next_id", 0) is None:
raise RuntimeError(
"TurboVecIndex.add() cannot infer the next id after load(): "
"the id high-water mark is not recoverable from the on-disk "
"index. Call add_with_ids(vectors, ids) with explicit ids "
"instead."
)
start = getattr(self, "_next_id", None) or 0
self.add_with_ids(vectors, np.arange(start, start + len(vectors)))
[docs]
def add_with_ids(self, vectors: np.ndarray, ids: np.ndarray) -> None:
"""Append vectors with explicit ids (O(1), no rebuild).
Parameters
----------
vectors : numpy.ndarray
Shape ``(m, dim)``.
ids : numpy.ndarray
Shape ``(m,)`` integer ids.
"""
if self._index is None:
turbovec = _require()
self._index = turbovec.IdMapIndex(dim=self.dim, bit_width=self._bit_width)
arr = self._prep(vectors)
labels = np.asarray(ids, dtype=np.uint64)
self._index.add_with_ids(arr, labels) # type: ignore[union-attr]
# Remember the next free contiguous id for a subsequent bare add().
self._next_id = int(labels.max()) + 1 if labels.size else getattr(self, "_next_id", 0)
[docs]
def remove(self, ids: np.ndarray) -> None:
"""Delete vectors by id — O(1) each, no structural degradation.
Parameters
----------
ids : numpy.ndarray
Shape ``(m,)`` integer ids to drop.
"""
for i in np.asarray(ids, dtype=np.uint64).tolist():
self._index.remove(int(i)) # type: ignore[union-attr]
[docs]
def search(self, queries: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
"""Return approximate top-``k`` neighbours per query.
Parameters
----------
queries : numpy.ndarray
Shape ``(q, dim)``.
k : int
Neighbours per query.
Returns
-------
ids : numpy.ndarray
Shape ``(q, k)`` neighbour ids.
distances : numpy.ndarray
Shape ``(q, k)`` distances under the index metric.
"""
arr = self._prep(queries)
# turbovec returns (distances, ids); the package contract is (ids, dist),
# so we swap the two halves of the tuple here.
distances, ids = self._index.search(arr, k) # type: ignore[union-attr]
ids = np.asarray(ids, dtype=np.int64)
distances = np.asarray(distances, dtype=np.float32)
# A corpus smaller than k gives turbovec fewer than k columns; pad back
# to the (q, k) contract every backend promises (see ANNIndex.search).
return self._pad(ids, k), self._pad(distances, k, fill=np.inf)
[docs]
def save(self, path: str) -> None:
"""Persist via turbovec's native ``write``.
Parameters
----------
path : str
Destination file path.
"""
self._index.write(path) # type: ignore[union-attr]
[docs]
def load(self, path: str) -> TurboVecIndex:
"""Load an index written by :meth:`save`.
Parameters
----------
path : str
Source path produced by :meth:`save`.
Returns
-------
TurboVecIndex
``self``, populated from disk.
Notes
-----
turbovec's native ``IdMapIndex`` exposes no way to enumerate or count
the ids it holds, so the id high-water mark :meth:`add` relies on
cannot be recovered here. ``_next_id`` is set to ``None`` — a
sentinel :meth:`add` checks for and refuses to guess past, raising a
clear error rather than silently restarting id assignment at 0 and
colliding with ids already in the loaded index. Use
:meth:`add_with_ids` with explicit ids after loading.
"""
turbovec = _require()
self._index = turbovec.IdMapIndex.load(path)
self._next_id = None
return self