Spaces:
Sleeping
Sleeping
File size: 5,676 Bytes
7d3b88b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """Vector store — semantic similarity index for memory retrieval.
Provides fast nearest-neighbour search over embedding vectors.
Backend strategy:
1. FAISS (preferred) — fast, production-grade
2. NumPy cosine similarity (fallback) — zero extra deps
"""
from __future__ import annotations
import logging
import threading
import numpy as np
logger = logging.getLogger(__name__)
# ── Backend detection ──────────────────────────────────────────────────────
try:
import faiss # type: ignore[import-untyped]
_HAS_FAISS = True
except ImportError:
_HAS_FAISS = False
# ── VectorStore ────────────────────────────────────────────────────────────
class VectorStore:
"""Thread-safe vector index with add / search / remove support.
Uses FAISS when available; falls back to brute-force NumPy cosine
similarity for environments where FAISS cannot be installed.
"""
def __init__(self, dimension: int = 384) -> None:
self._dimension = dimension
self._lock = threading.Lock()
if _HAS_FAISS:
base = faiss.IndexFlatIP(dimension)
self._index: faiss.IndexIDMap = faiss.IndexIDMap(base)
self._backend = "faiss"
else:
self._vectors: dict[int, np.ndarray] = {}
self._backend = "numpy"
logger.info(
"[VECTOR] Initialised (dim=%d, backend=%s)",
dimension, self._backend,
)
# ── Properties ─────────────────────────────────────────────────────
@property
def backend(self) -> str:
"""Return the active backend name."""
return self._backend
@property
def size(self) -> int:
"""Current number of indexed vectors."""
with self._lock:
if self._backend == "faiss":
return self._index.ntotal
return len(self._vectors)
# ── Public API ─────────────────────────────────────────────────────
def add(self, embedding: list[float], entry_id: int) -> None:
"""Index an embedding under the given *entry_id*."""
vec = self._to_unit_vec(embedding)
with self._lock:
if self._backend == "faiss":
ids = np.array([entry_id], dtype=np.int64)
self._index.add_with_ids(vec, ids)
else:
self._vectors[entry_id] = vec.flatten()
logger.debug("[VECTOR] Added entry_id=%d (total=%d)", entry_id, self.size)
def remove(self, entry_id: int) -> None:
"""Remove an embedding by *entry_id*."""
with self._lock:
if self._backend == "faiss":
ids = np.array([entry_id], dtype=np.int64)
try:
self._index.remove_ids(ids)
except Exception as exc:
logger.warning(
"[VECTOR] FAISS remove failed for id=%d: %s",
entry_id, exc,
)
else:
self._vectors.pop(entry_id, None)
logger.debug("[VECTOR] Removed entry_id=%d", entry_id)
def search(
self,
query_embedding: list[float],
top_k: int = 5,
) -> list[tuple[int, float]]:
"""Return the *top_k* most similar entries.
Returns:
List of ``(entry_id, similarity_score)`` sorted descending.
"""
vec = self._to_unit_vec(query_embedding)
with self._lock:
if self._backend == "faiss":
return self._search_faiss(vec, top_k)
return self._search_numpy(vec, top_k)
def clear(self) -> None:
"""Drop all indexed vectors."""
with self._lock:
if self._backend == "faiss":
base = faiss.IndexFlatIP(self._dimension)
self._index = faiss.IndexIDMap(base)
else:
self._vectors.clear()
logger.info("[VECTOR] Index cleared")
# ── Internals ──────────────────────────────────────────────────────
def _to_unit_vec(self, embedding: list[float]) -> np.ndarray:
"""Convert to float32 row vector and L2-normalise."""
vec = np.array(embedding, dtype=np.float32).reshape(1, -1)
norm = np.linalg.norm(vec)
if norm > 0:
vec = vec / norm
return vec
def _search_faiss(
self, query_vec: np.ndarray, top_k: int,
) -> list[tuple[int, float]]:
if self._index.ntotal == 0:
return []
k = min(top_k, self._index.ntotal)
scores, ids = self._index.search(query_vec, k)
return [
(int(eid), float(score))
for score, eid in zip(scores[0], ids[0])
if eid >= 0
]
def _search_numpy(
self, query_vec: np.ndarray, top_k: int,
) -> list[tuple[int, float]]:
if not self._vectors:
return []
q = query_vec.flatten()
scored = [
(eid, float(np.dot(q, vec)))
for eid, vec in self._vectors.items()
]
scored.sort(key=lambda x: x[1], reverse=True)
return scored[:top_k]
|