ltcai 11.7.0 → 12.0.0

package/latticeai/core/vector_index/hnsw.py

@@ -0,0 +1,383 @@
+"""Approximate nearest-neighbour index with incremental append.
+
+``hnswlib`` is optional (``pip install hnswlib``; the published
+``ltcai[hnsw]`` extra was retired in 11.6.0 when the write path moved to
+Rust). When the compiled module is missing the import is caught here and
+reported as a reason string; search then returns no ANN hits rather than
+pretending.
+
+Policy, also stated on :meth:`HnswIndex.add_items`:
+
+* **Append** new ids onto the live graph (``resize_index`` + ``add_items``).
+* **Full rebuild** on provider/dim/space change, on any deletion when the
+  graph was loaded from a sidecar without source vectors, or when deletes
+  exceed :data:`DELETE_REBUILD_RATIO` of the current size.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
+
+HNSW_BACKEND = "hnsw"
+HNSWLIB_MODULE = "hnswlib"
+HNSW_SUFFIX = ".hnsw"
+HNSW_META_SUFFIX = ".hnsw.meta.json"
+#: Rebuild instead of tombstoning once this fraction of the graph would go.
+DELETE_REBUILD_RATIO = 0.10
+
+IndexItem = Tuple[str, Sequence[float], Mapping[str, Any]]
+ScoredId = Tuple[str, float]
+
+
+def load_hnswlib() -> Tuple[Optional[Any], Optional[str]]:
+    """Import ``hnswlib`` or explain why it is unavailable (never raises)."""
+    try:
+        import hnswlib  # noqa: PLC0415 — guarded, optional, import-on-demand
+    except Exception as exc:  # noqa: BLE001 — any import failure is "no ANN"
+        return None, (
+            f"{HNSWLIB_MODULE} is not available ({exc}); install it with "
+            "`pip install hnswlib` to use the approximate index"
+        )
+    return hnswlib, None
+
+
+def hnswlib_available() -> bool:
+    """True when the optional ANN engine can actually be imported."""
+    module, _ = load_hnswlib()
+    return module is not None
+
+
+def sidecar_paths(db_path: Any) -> Tuple[Path, Path]:
+    """``(index file, meta file)`` for a brain database path."""
+    base = Path(db_path)
+    return (
+        base.with_suffix(HNSW_SUFFIX),
+        base.with_suffix(HNSW_META_SUFFIX),
+    )
+
+
+class HnswIndex:
+    """Small-M HNSW graph over the brain's embeddings."""
+
+    def __init__(
+        self,
+        *,
+        dim: int,
+        space: str = "cosine",
+        ef_construction: int = 400,
+        m: int = 32,
+        ef_search: int = 400,
+        model_id: str = "",
+    ) -> None:
+        self._dim = int(dim)
+        self._space = str(space)
+        self._ef_construction = int(ef_construction)
+        self._m = int(m)
+        self._ef_search = int(ef_search)
+        self._model_id = str(model_id)
+        self._module, self._detail = load_hnswlib()
+        self._vectors: Dict[str, List[float]] = {}
+        self._metadata: Dict[str, Dict[str, Any]] = {}
+        self._labels: List[str] = []
+        self._ann: Any = None
+        self._dirty = True
+        self._loaded = False
+
+    @property
+    def backend(self) -> str:
+        return HNSW_BACKEND
+
+    @property
+    def approx(self) -> bool:
+        return True
+
+    @property
+    def exhaustive(self) -> bool:
+        return False
+
+    @property
+    def available(self) -> bool:
+        return self._module is not None
+
+    @property
+    def unavailable_detail(self) -> Optional[str]:
+        return self._detail
+
+    @property
+    def loaded_from_sidecar(self) -> bool:
+        """True when the graph came off disk instead of being rebuilt."""
+        return self._loaded
+
+    @property
+    def dim(self) -> int:
+        return self._dim
+
+    @property
+    def model_id(self) -> str:
+        return self._model_id
+
+    def _identity_matches(self, dim: int, space: str, model_id: str) -> bool:
+        return (
+            int(dim) == self._dim
+            and str(space) == self._space
+            and str(model_id) == self._model_id
+        )
+
+    def add(
+        self,
+        id: str,
+        vector: Sequence[float],
+        metadata: Optional[Mapping[str, Any]] = None,
+    ) -> None:
+        if self._loaded:
+            raise RuntimeError(
+                "this HnswIndex was loaded from a sidecar and holds no source "
+                "vectors; rebuild it from vector_embeddings before mutating"
+            )
+        key = str(id)
+        self._vectors[key] = [float(value) for value in vector]
+        self._metadata[key] = dict(metadata or {})
+        self._dirty = True
+
+    def remove(self, id: str) -> None:
+        if self._loaded:
+            raise RuntimeError(
+                "this HnswIndex was loaded from a sidecar and holds no source "
+                "vectors; rebuild it from vector_embeddings before mutating"
+            )
+        key = str(id)
+        self._vectors.pop(key, None)
+        self._metadata.pop(key, None)
+        self._dirty = True
+
+    def rebuild(self, items: Iterable[IndexItem]) -> None:
+        """Replace the held set and mark the graph for a full rebuild."""
+        self._vectors.clear()
+        self._metadata.clear()
+        self._labels.clear()
+        self._ann = None
+        self._loaded = False
+        for item_id, vector, metadata in items:
+            self.add(item_id, vector, metadata)
+        self._dirty = True
+
+    def add_items(
+        self,
+        items: Iterable[IndexItem],
+        *,
+        dim: Optional[int] = None,
+        space: Optional[str] = None,
+        model_id: Optional[str] = None,
+    ) -> str:
+        """Append new vectors, or rebuild when the policy says so.
+
+        Returns ``"append"`` or ``"rebuild"`` so a bench can time the two
+        doors without guessing.
+
+        Full rebuild when the embedder identity changed, when any currently
+        indexed id is missing from the incoming set (a deletion), when
+        deletes would exceed :data:`DELETE_REBUILD_RATIO`, or when this
+        instance was loaded from a sidecar and cannot append (no source
+        vectors to keep the label map honest).
+        """
+        incoming: Dict[str, IndexItem] = {}
+        for item_id, vector, metadata in items:
+            incoming[str(item_id)] = (str(item_id), vector, metadata)
+
+        want_dim = self._dim if dim is None else int(dim)
+        want_space = self._space if space is None else str(space)
+        want_model = self._model_id if model_id is None else str(model_id)
+        if not self._identity_matches(want_dim, want_space, want_model):
+            self._dim = want_dim
+            self._space = want_space
+            self._model_id = want_model
+            self.rebuild(incoming.values())
+            return "rebuild"
+
+        if self._loaded:
+            self.rebuild(incoming.values())
+            return "rebuild"
+
+        current = set(self._vectors)
+        incoming_ids = set(incoming)
+        deleted = current - incoming_ids
+        if deleted:
+            ratio = len(deleted) / max(1, len(current))
+            if ratio >= DELETE_REBUILD_RATIO or not current:
+                self.rebuild(incoming.values())
+                return "rebuild"
+            # Sparse deletes: drop the gone ids and rebuild so we never
+            # leave a tombstoned hole the sidecar cannot describe.
+            self.rebuild(incoming.values())
+            return "rebuild"
+
+        new_ids = [item_id for item_id in incoming if item_id not in current]
+        if not new_ids and not self._dirty and self._ann is not None:
+            return "append"
+        for item_id in new_ids:
+            _, vector, metadata = incoming[item_id]
+            self.add(item_id, vector, metadata)
+        if self._ann is not None and new_ids and self._module is not None:
+            self._append_live(new_ids)
+            return "append"
+        self._dirty = True
+        return "append" if new_ids or self._ann is not None else "rebuild"
+
+    def _new_graph(self, module: Any, capacity: int) -> Any:
+        graph = module.Index(space=self._space, dim=self._dim)
+        graph.init_index(
+            max_elements=max(1, capacity),
+            ef_construction=self._ef_construction,
+            M=self._m,
+        )
+        return graph
+
+    def _append_live(self, new_ids: Sequence[str]) -> None:
+        """``hnswlib`` ``resize_index`` + ``add_items`` for ``new_ids`` only."""
+        module = self._module
+        graph = self._ann
+        if module is None or graph is None:
+            self._dirty = True
+            return
+        start = len(self._labels)
+        needed = start + len(new_ids)
+        try:
+            current_max = int(graph.get_max_elements())
+        except Exception:  # noqa: BLE001 — some builds expose no getter
+            current_max = start
+        if needed > current_max:
+            graph.resize_index(max(needed, current_max * 2 if current_max else needed))
+        graph.add_items(
+            [self._vectors[key] for key in new_ids],
+            list(range(start, start + len(new_ids))),
+        )
+        self._labels.extend(new_ids)
+        self._dirty = False
+
+    def _ensure_graph(self) -> Any:
+        """Build the graph if the held vectors changed (None when disabled)."""
+        module = self._module
+        if module is None:
+            return None
+        if self._ann is not None and not self._dirty:
+            return self._ann
+        labels = list(self._vectors)
+        graph = self._new_graph(module, max(len(labels), 1))
+        if labels:
+            graph.add_items(
+                [self._vectors[key] for key in labels],
+                list(range(len(labels))),
+            )
+        self._ann = graph
+        self._labels = labels
+        self._dirty = False
+        return graph
+
+    def search(
+        self,
+        query: Sequence[float],
+        top_k: int,
+        filter: Optional[Mapping[str, Any]] = None,
+    ) -> List[ScoredId]:
+        graph = self._ensure_graph()
+        if graph is None or not self._labels:
+            return []
+        floor = float("-inf")
+        if filter and filter.get("min_score") is not None:
+            try:
+                floor = float(filter["min_score"])
+            except (TypeError, ValueError):
+                floor = float("-inf")
+        wanted = max(1, min(int(top_k), len(self._labels)))
+        graph.set_ef(max(self._ef_search, wanted))
+        labels, distances = graph.knn_query([list(query)], k=wanted)
+        scored: List[ScoredId] = []
+        for label, distance in zip(labels[0], distances[0], strict=True):
+            score = 1.0 - float(distance)
+            if score < floor:
+                continue
+            scored.append((self._labels[int(label)], score))
+        return scored
+
+    def stats(self) -> Dict[str, Any]:
+        return {
+            "backend": self.backend,
+            "size": len(self._labels) if self._loaded else len(self._vectors),
+            "dim": self._dim,
+            "approx": True,
+            "exhaustive": False,
+            "available": self.available,
+            "detail": self._detail,
+            "model_id": self._model_id,
+        }
+
+    def save(self, db_path: Any, *, fingerprint: str) -> bool:
+        """Write the graph + label map beside the brain database."""
+        graph = self._ensure_graph()
+        if graph is None or not self._labels:
+            return False
+        index_path, meta_path = sidecar_paths(db_path)
+        try:
+            index_path.parent.mkdir(parents=True, exist_ok=True)
+            graph.save_index(str(index_path))
+            meta_path.write_text(
+                json.dumps(
+                    {
+                        "fingerprint": str(fingerprint),
+                        "dim": self._dim,
+                        "space": self._space,
+                        "model_id": self._model_id,
+                        "labels": self._labels,
+                    },
+                    ensure_ascii=False,
+                ),
+                encoding="utf-8",
+            )
+        except Exception:  # noqa: BLE001 — persistence is best-effort
+            return False
+        return True
+
+    def load(self, db_path: Any, *, fingerprint: str) -> bool:
+        """Adopt a sidecar graph when it provably matches ``fingerprint``."""
+        module = self._module
+        if module is None:
+            return False
+        index_path, meta_path = sidecar_paths(db_path)
+        try:
+            meta = json.loads(meta_path.read_text(encoding="utf-8"))
+            matches = str(meta["fingerprint"]) == str(fingerprint)
+            matches = matches and int(meta["dim"]) == self._dim
+            if meta.get("model_id") not in (None, "", self._model_id):
+                matches = False
+            labels = [str(label) for label in meta["labels"]]
+        except Exception:  # noqa: BLE001 — absent/corrupt sidecar = rebuild
+            return False
+        if not (matches and labels):
+            return False
+        try:
+            graph = module.Index(space=self._space, dim=self._dim)
+            graph.load_index(str(index_path), max_elements=len(labels))
+        except Exception:  # noqa: BLE001 — corrupt binary = rebuild
+            return False
+        self._ann = graph
+        self._labels = labels
+        self._vectors.clear()
+        self._metadata.clear()
+        self._dirty = False
+        self._loaded = True
+        return True
+
+
+__all__ = [
+    "DELETE_REBUILD_RATIO",
+    "HNSW_BACKEND",
+    "HNSW_META_SUFFIX",
+    "HNSW_SUFFIX",
+    "HNSWLIB_MODULE",
+    "HnswIndex",
+    "hnswlib_available",
+    "load_hnswlib",
+    "sidecar_paths",
+]

package/latticeai/core/vector_index/sidecar.py

@@ -0,0 +1,329 @@
+"""Load, refresh, and query the HNSW sidecar next to a brain database.
+
+The live graph is cached per ``(db_path, model_id, dim)``. A missing sidecar
+is ``index: "none"``. When the store has vectors the cache (or a rebuild from
+``vector_embeddings``) can serve, the reply is ``index: "hnsw"``. Staleness is
+sidecar size vs ``COUNT(*)`` for that identity — reported honestly, and
+refreshed through :meth:`HnswIndex.add_items` so an ingest append does not
+rebuild the whole graph.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+import struct
+import threading
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+from latticeai.core.config import default_data_dir
+
+from .hnsw import HnswIndex, hnswlib_available, sidecar_paths
+
+GRAPH_DB_NAME = "knowledge_graph.sqlite"
+VECTOR_QUERY_K_CAP = 200
+
+_CACHE: Dict[Tuple[str, str, int], HnswIndex] = {}
+_LOCK = threading.Lock()
+
+
+def resolve_graph_db(workspace: Optional[str] = None) -> Path:
+    """Brain sqlite path: explicit file, a data dir, or ``LATTICEAI_DATA_DIR``."""
+    if workspace:
+        raw = Path(str(workspace)).expanduser()
+        if raw.is_file():
+            return raw
+        if raw.suffix.lower() == ".sqlite":
+            return raw
+        return raw / GRAPH_DB_NAME
+    return default_data_dir() / GRAPH_DB_NAME
+
+
+def decode_f32le(blob: bytes, dim: int) -> List[float]:
+    """Little-endian f32 payload — the same layout Rust ``encode`` writes."""
+    if not blob:
+        return []
+    count = dim if dim > 0 else len(blob) // 4
+    if len(blob) != count * 4:
+        count = len(blob) // 4
+    if count <= 0:
+        return []
+    return list(struct.unpack(f"<{count}f", blob[: count * 4]))
+
+
+def sidecar_fingerprint(model_id: str, dim: int, size: int) -> str:
+    """Identity the ``.hnsw.meta.json`` file is keyed on."""
+    return f"{model_id}|{int(dim)}|{int(size)}"
+
+
+def store_vector_count(conn: sqlite3.Connection, model_id: str, dim: int) -> int:
+    row = conn.execute(
+        "SELECT COUNT(*) FROM vector_embeddings "
+        "WHERE embedding_model=? AND embedding_dim=?",
+        (model_id, int(dim)),
+    ).fetchone()
+    return int(row[0]) if row else 0
+
+
+def load_store_items(
+    conn: sqlite3.Connection, model_id: str, dim: int
+) -> List[Tuple[str, List[float], Dict[str, Any]]]:
+    items: List[Tuple[str, List[float], Dict[str, Any]]] = []
+    for item_id, blob in conn.execute(
+        "SELECT item_id, embedding FROM vector_embeddings "
+        "WHERE embedding_model=? AND embedding_dim=? ORDER BY indexed_at ASC, item_id ASC",
+        (model_id, int(dim)),
+    ):
+        vector = decode_f32le(bytes(blob or b""), int(dim))
+        if len(vector) != int(dim):
+            continue
+        items.append((str(item_id), vector, {}))
+    return items
+
+
+def sidecar_meta_size(db_path: Path) -> Optional[int]:
+    """Label count from the sidecar meta file, or ``None`` when it is absent."""
+    _, meta_path = sidecar_paths(db_path)
+    try:
+        import json
+
+        meta = json.loads(meta_path.read_text(encoding="utf-8"))
+        labels = meta.get("labels") or []
+        return len(labels)
+    except Exception:  # noqa: BLE001 — absent/corrupt meta is "no sidecar"
+        return None
+
+
+def _open_readonly(db_path: Path) -> Optional[sqlite3.Connection]:
+    if not db_path.is_file():
+        return None
+    try:
+        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
+    except sqlite3.Error:
+        try:
+            conn = sqlite3.connect(str(db_path))
+        except sqlite3.Error:
+            return None
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def _none_reply(
+    *,
+    size: int = 0,
+    store_size: int = 0,
+    stale: bool = False,
+    detail: Optional[str] = None,
+) -> Dict[str, Any]:
+    return {
+        "ids": [],
+        "scores": [],
+        "index": "none",
+        "size": int(size),
+        "store_size": int(store_size),
+        "stale": bool(stale),
+        "detail": detail,
+        "refreshed": None,
+    }
+
+
+def _hnsw_reply(
+    ids: Sequence[str],
+    scores: Sequence[float],
+    *,
+    size: int,
+    store_size: int,
+    stale: bool,
+    detail: Optional[str],
+    refreshed: Optional[str],
+) -> Dict[str, Any]:
+    return {
+        "ids": [str(item) for item in ids],
+        "scores": [float(score) for score in scores],
+        "index": "hnsw",
+        "size": int(size),
+        "store_size": int(store_size),
+        "stale": bool(stale),
+        "detail": detail,
+        "refreshed": refreshed,
+    }
+
+
+def _refresh_index(
+    index: HnswIndex,
+    items: Sequence[Tuple[str, List[float], Dict[str, Any]]],
+    model_id: str,
+    dim: int,
+) -> str:
+    return index.add_items(items, dim=dim, model_id=model_id)
+
+
+def query_sidecar(
+    *,
+    workspace: Optional[str],
+    embedding_model: str,
+    embedding_dim: int,
+    vector: Sequence[float],
+    k: int,
+    db_path: Optional[Path] = None,
+) -> Dict[str, Any]:
+    """Serve one ANN query from the sidecar (or say the index is absent)."""
+    dim = int(embedding_dim)
+    model_id = str(embedding_model or "")
+    wanted = max(1, min(int(k), VECTOR_QUERY_K_CAP))
+    path = Path(db_path) if db_path is not None else resolve_graph_db(workspace)
+    if not hnswlib_available():
+        return _none_reply(detail="hnswlib is not available in this worker")
+    if dim <= 0 or len(vector) != dim:
+        return _none_reply(
+            detail=(
+                f"query vector width {len(vector)} does not match embedding_dim {dim}"
+                if dim > 0
+                else "embedding_dim must be a positive integer"
+            )
+        )
+    if not model_id:
+        return _none_reply(detail="embedding_model is required")
+
+    conn = _open_readonly(path)
+    store_size = 0
+    items: List[Tuple[str, List[float], Dict[str, Any]]] = []
+    if conn is not None:
+        try:
+            store_size = store_vector_count(conn, model_id, dim)
+            if store_size > 0:
+                items = load_store_items(conn, model_id, dim)
+                store_size = len(items)
+        finally:
+            conn.close()
+
+    key = (str(path), model_id, dim)
+    with _LOCK:
+        index = _CACHE.get(key)
+        refreshed: Optional[str] = None
+        if index is None:
+            index = HnswIndex(dim=dim, model_id=model_id)
+            loaded = index.load(path, fingerprint=sidecar_fingerprint(model_id, dim, store_size))
+            if not loaded:
+                if not items:
+                    return _none_reply(
+                        store_size=store_size,
+                        detail="no HNSW sidecar for this embedding identity",
+                    )
+                refreshed = _refresh_index(index, items, model_id, dim)
+                index.save(path, fingerprint=sidecar_fingerprint(model_id, dim, len(items)))
+            _CACHE[key] = index
+
+        cached_size = int(index.stats().get("size") or 0)
+        stale = cached_size != store_size
+        if stale:
+            if not items and store_size > 0:
+                conn = _open_readonly(path)
+                if conn is not None:
+                    try:
+                        items = load_store_items(conn, model_id, dim)
+                    finally:
+                        conn.close()
+            if items:
+                refreshed = _refresh_index(index, items, model_id, dim)
+                index.save(path, fingerprint=sidecar_fingerprint(model_id, dim, len(items)))
+                cached_size = int(index.stats().get("size") or 0)
+                stale = cached_size != len(items)
+            elif store_size == 0:
+                _CACHE.pop(key, None)
+                return _none_reply(
+                    store_size=0,
+                    detail="vector store is empty for this embedding identity",
+                )
+
+        if cached_size <= 0:
+            return _none_reply(
+                store_size=store_size,
+                stale=stale,
+                detail="HNSW sidecar has no labels for this identity",
+            )
+
+        index._ef_search = max(index._ef_search, wanted * 2, 200)
+        hits = index.search(list(vector), top_k=min(wanted, cached_size))
+        ids = [item_id for item_id, _ in hits]
+        scores = [score for _, score in hits]
+        detail = None
+        if stale:
+            detail = (
+                f"sidecar size {cached_size} != store count {store_size} "
+                "for this embedding identity"
+            )
+        return _hnsw_reply(
+            ids,
+            scores,
+            size=cached_size,
+            store_size=store_size,
+            stale=stale,
+            detail=detail,
+            refreshed=refreshed,
+        )
+
+
+def sidecar_freshness(
+    *,
+    workspace: Optional[str] = None,
+    embedding_model: Optional[str] = None,
+    embedding_dim: Optional[int] = None,
+    db_path: Optional[Path] = None,
+) -> Dict[str, Any]:
+    """Cheap on-disk vs store comparison — no ANN query, no rebuild."""
+    path = Path(db_path) if db_path is not None else resolve_graph_db(workspace)
+    meta_size = sidecar_meta_size(path)
+    store_size = 0
+    conn = _open_readonly(path)
+    if conn is not None:
+        try:
+            if embedding_model and embedding_dim:
+                store_size = store_vector_count(conn, embedding_model, int(embedding_dim))
+            else:
+                row = conn.execute("SELECT COUNT(*) FROM vector_embeddings").fetchone()
+                store_size = int(row[0]) if row else 0
+        except sqlite3.Error:
+            store_size = 0
+        finally:
+            conn.close()
+    if meta_size is None:
+        return {
+            "index": "none",
+            "size": 0,
+            "store_size": store_size,
+            "stale": store_size > 0,
+            "detail": "no HNSW sidecar next to the brain database",
+        }
+    stale = meta_size != store_size
+    return {
+        "index": "hnsw",
+        "size": meta_size,
+        "store_size": store_size,
+        "stale": stale,
+        "detail": (
+            f"sidecar size {meta_size} != store count {store_size}"
+            if stale
+            else "sidecar matches the vector store"
+        ),
+    }
+
+
+def reset_sidecar_cache() -> None:
+    """Drop the process cache — tests only."""
+    with _LOCK:
+        _CACHE.clear()
+
+
+__all__ = [
+    "GRAPH_DB_NAME",
+    "VECTOR_QUERY_K_CAP",
+    "decode_f32le",
+    "query_sidecar",
+    "reset_sidecar_cache",
+    "resolve_graph_db",
+    "sidecar_fingerprint",
+    "sidecar_freshness",
+    "sidecar_meta_size",
+    "store_vector_count",
+]