Add vrom_hub/hnsw.py

vrom_hub/hnsw.py

@@ -0,0 +1,348 @@
+"""
+Pure-Python HNSW index builder.
+
+Produces index.json files that are 100% compatible with the Rust/WASM
+VectorDB.load(json) method. Mirrors the exact serde serialization format
+of the vecdb-wasm crate.
+
+Key invariants:
+- Node ID = index in the nodes array
+- neighbors[i] = connections at layer i (length = node.max_layer + 1)
+- Layer 0 uses m_max0 (= 2*m) max neighbors; higher layers use m
+- metric variants are PascalCase: "Cosine", "Euclidean", "DotProduct"
+- metadata is a JSON string (not an object)
+- entry_point is null for empty indexes
+"""
+
+from __future__ import annotations
+
+import heapq
+import json
+import logging
+import math
+import random
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class HnswConfig:
+    """HNSW algorithm parameters — mirrors Rust HnswConfig exactly."""
+    m: int = 16
+    m_max0: int = 32
+    ef_construction: int = 128
+    ef_search: int = 40
+    level_multiplier: float = 0.0  # computed from m if 0
+    metric: str = "Cosine"
+
+    def __post_init__(self):
+        if self.m_max0 == 0:
+            self.m_max0 = 2 * self.m
+        if self.level_multiplier == 0.0:
+            self.level_multiplier = 1.0 / math.log(self.m)
+
+    def to_dict(self) -> dict:
+        return {
+            "m": self.m,
+            "m_max0": self.m_max0,
+            "ef_construction": self.ef_construction,
+            "ef_search": self.ef_search,
+            "level_multiplier": self.level_multiplier,
+            "metric": self.metric,
+        }
+
+
+@dataclass
+class HnswNode:
+    """A single node in the HNSW graph — mirrors Rust HnswNode exactly."""
+    vector: list[float]          # len = dim
+    neighbors: list[list[int]]   # neighbors[layer] = [node_ids...]
+    max_layer: int               # highest layer this node is in
+    metadata: Optional[str]      # JSON string or None
+
+    def to_dict(self) -> dict:
+        return {
+            "vector": self.vector,
+            "neighbors": self.neighbors,
+            "max_layer": self.max_layer,
+            "metadata": self.metadata,
+        }
+
+
+class HnswIndex:
+    """
+    Top-level HNSW index — serializes to the exact JSON format
+    that VectorDB.load() expects.
+    """
+
+    def __init__(self, config: HnswConfig, dim: int):
+        self.config = config
+        self.dim = dim
+        self.nodes: list[HnswNode] = []
+        self.entry_point: Optional[int] = None
+        self.max_layer: int = 0
+
+    def to_dict(self) -> dict:
+        return {
+            "config": self.config.to_dict(),
+            "nodes": [n.to_dict() for n in self.nodes],
+            "entry_point": self.entry_point,
+            "max_layer": self.max_layer,
+            "dim": self.dim,
+        }
+
+    def to_json(self, indent: int | None = None) -> str:
+        return json.dumps(self.to_dict(), indent=indent)
+
+
+def _cosine_distance(a: np.ndarray, b: np.ndarray) -> float:
+    """1 - cosine_similarity. For pre-normalized vectors, this is 1 - dot(a,b)."""
+    dot = float(np.dot(a, b))
+    return 1.0 - dot
+
+
+def _euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
+    return float(np.linalg.norm(a - b))
+
+
+def _dot_product_distance(a: np.ndarray, b: np.ndarray) -> float:
+    """Negative dot product (lower = more similar)."""
+    return -float(np.dot(a, b))
+
+
+def _get_distance_fn(metric: str):
+    if metric == "Cosine":
+        return _cosine_distance
+    elif metric == "Euclidean":
+        return _euclidean_distance
+    elif metric == "DotProduct":
+        return _dot_product_distance
+    else:
+        raise ValueError(f"Unknown metric: {metric}")
+
+
+class HnswBuilder:
+    """
+    Builds an HNSW index from vectors + metadata.
+
+    This is a faithful Python implementation of the HNSW insertion algorithm
+    that produces a graph loadable by the Rust VectorDB.load() method.
+    """
+
+    def __init__(self, config: HnswConfig | None = None, dim: int = 384):
+        self.config = config or HnswConfig()
+        self.dim = dim
+        self.index = HnswIndex(self.config, dim)
+        self._vectors: list[np.ndarray] = []  # numpy cache for fast distance
+        self._dist_fn = _get_distance_fn(self.config.metric)
+        self._rng = random.Random(42)
+
+    def _random_layer(self) -> int:
+        """Sample a random layer using the HNSW level multiplier."""
+        r = self._rng.random()
+        return int(-math.log(r) * self.config.level_multiplier)
+
+    def _distance(self, a_id: int, b_vec: np.ndarray) -> float:
+        """Distance between stored node a_id and query vector b_vec."""
+        return self._dist_fn(self._vectors[a_id], b_vec)
+
+    def _search_layer(
+        self, query: np.ndarray, entry_id: int, ef: int, layer: int
+    ) -> list[tuple[float, int]]:
+        """
+        Search a single layer of the HNSW graph.
+        Returns up to ef nearest (distance, node_id) pairs.
+        """
+        entry_dist = self._distance(entry_id, query)
+        candidates = [(entry_dist, entry_id)]  # min-heap
+        results = [(-entry_dist, entry_id)]     # max-heap (negative for max)
+        visited = {entry_id}
+
+        while candidates:
+            c_dist, c_id = heapq.heappop(candidates)
+
+            # Furthest in results
+            f_dist = -results[0][0]
+            if c_dist > f_dist:
+                break
+
+            # Explore neighbors at this layer
+            node = self.index.nodes[c_id]
+            if layer < len(node.neighbors):
+                for neighbor_id in node.neighbors[layer]:
+                    if neighbor_id in visited:
+                        continue
+                    visited.add(neighbor_id)
+
+                    n_dist = self._distance(neighbor_id, query)
+                    f_dist = -results[0][0]
+
+                    if n_dist < f_dist or len(results) < ef:
+                        heapq.heappush(candidates, (n_dist, neighbor_id))
+                        heapq.heappush(results, (-n_dist, neighbor_id))
+                        if len(results) > ef:
+                            heapq.heappop(results)
+
+        # Convert results (stored as negative distances)
+        return [(abs(d), nid) for d, nid in results]
+
+    def _select_neighbors_simple(
+        self, candidates: list[tuple[float, int]], m: int
+    ) -> list[int]:
+        """Select the M nearest neighbors from candidates."""
+        candidates.sort(key=lambda x: x[0])
+        return [nid for _, nid in candidates[:m]]
+
+    def add(self, vector: np.ndarray, metadata: str | None = None) -> int:
+        """
+        Insert a vector into the HNSW index.
+
+        Args:
+            vector: numpy array of shape (dim,)
+            metadata: Optional JSON string metadata
+
+        Returns:
+            The node ID (= index in nodes array)
+        """
+        assert vector.shape == (self.dim,), f"Expected ({self.dim},), got {vector.shape}"
+        node_id = len(self.index.nodes)
+        node_layer = self._random_layer()
+
+        # Create node with empty neighbor lists
+        node = HnswNode(
+            vector=vector.tolist(),
+            neighbors=[[] for _ in range(node_layer + 1)],
+            max_layer=node_layer,
+            metadata=metadata,
+        )
+        self.index.nodes.append(node)
+        self._vectors.append(vector.copy())
+
+        # First node — set as entry point
+        if self.index.entry_point is None:
+            self.index.entry_point = node_id
+            self.index.max_layer = node_layer
+            return node_id
+
+        # Traverse from top to the node's layer, greedily
+        ep_id = self.index.entry_point
+        current_layer = self.index.max_layer
+
+        # Phase 1: Greedy descent from top to node_layer + 1
+        while current_layer > node_layer:
+            results = self._search_layer(vector, ep_id, ef=1, layer=current_layer)
+            if results:
+                ep_id = min(results, key=lambda x: x[0])[1]
+            current_layer -= 1
+
+        # Phase 2: Search and connect at each layer from min(node_layer, max_layer) down to 0
+        for lc in range(min(node_layer, self.index.max_layer), -1, -1):
+            results = self._search_layer(
+                vector, ep_id, ef=self.config.ef_construction, layer=lc
+            )
+
+            # Select neighbors
+            m_for_layer = self.config.m_max0 if lc == 0 else self.config.m
+            neighbors = self._select_neighbors_simple(results, m_for_layer)
+
+            # Set this node's neighbors at layer lc
+            node.neighbors[lc] = neighbors
+
+            # Add bidirectional connections
+            for neighbor_id in neighbors:
+                neighbor_node = self.index.nodes[neighbor_id]
+                # Ensure neighbor has enough layers
+                while len(neighbor_node.neighbors) <= lc:
+                    neighbor_node.neighbors.append([])
+
+                neighbor_node.neighbors[lc].append(node_id)
+
+                # Shrink if over capacity
+                max_conn = self.config.m_max0 if lc == 0 else self.config.m
+                if len(neighbor_node.neighbors[lc]) > max_conn:
+                    # Keep only the closest
+                    n_vec = self._vectors[neighbor_id]
+                    scored = [
+                        (self._dist_fn(self._vectors[nid], n_vec), nid)
+                        for nid in neighbor_node.neighbors[lc]
+                    ]
+                    neighbor_node.neighbors[lc] = self._select_neighbors_simple(
+                        scored, max_conn
+                    )
+
+            # Update entry point for next layer
+            if results:
+                ep_id = min(results, key=lambda x: x[0])[1]
+
+        # Update global entry point if new node is higher
+        if node_layer > self.index.max_layer:
+            self.index.entry_point = node_id
+            self.index.max_layer = node_layer
+
+        return node_id
+
+    def build(
+        self,
+        vectors: np.ndarray,
+        metadatas: list[str | None] | None = None,
+    ) -> HnswIndex:
+        """
+        Build an entire HNSW index from a batch of vectors.
+
+        Args:
+            vectors: np.ndarray of shape (n, dim)
+            metadatas: Optional list of JSON string metadata, one per vector
+
+        Returns:
+            The built HnswIndex
+        """
+        n = vectors.shape[0]
+        if metadatas is None:
+            metadatas = [None] * n
+
+        assert len(metadatas) == n, f"Mismatch: {n} vectors, {len(metadatas)} metadatas"
+        assert vectors.shape[1] == self.dim, f"Expected dim={self.dim}, got {vectors.shape[1]}"
+
+        logger.info(f"Building HNSW index: {n} vectors, dim={self.dim}, m={self.config.m}")
+
+        for i in range(n):
+            self.add(vectors[i], metadatas[i])
+            if (i + 1) % 100 == 0 or i == n - 1:
+                logger.info(f"  Indexed {i + 1}/{n} vectors (max_layer={self.index.max_layer})")
+
+        logger.info(
+            f"HNSW index built: {n} nodes, max_layer={self.index.max_layer}, "
+            f"entry_point={self.index.entry_point}"
+        )
+        return self.index
+
+    def search(self, query: np.ndarray, k: int = 5) -> list[tuple[float, int]]:
+        """
+        Search the index for the k nearest neighbors.
+
+        Returns:
+            List of (distance, node_id) tuples, sorted by distance (ascending).
+        """
+        if self.index.entry_point is None:
+            return []
+
+        ep_id = self.index.entry_point
+        current_layer = self.index.max_layer
+
+        # Greedy descent to layer 0
+        while current_layer > 0:
+            results = self._search_layer(query, ep_id, ef=1, layer=current_layer)
+            if results:
+                ep_id = min(results, key=lambda x: x[0])[1]
+            current_layer -= 1
+
+        # Search layer 0 with ef_search
+        results = self._search_layer(
+            query, ep_id, ef=max(k, self.config.ef_search), layer=0
+        )
+        results.sort(key=lambda x: x[0])
+        return results[:k]