c762564d40
- Fix 71 invalid-syntax files (class-body newline-broken assignments) - Add from/None chain to 307 B904 raise-without-from sites - Add B008 ignore to ruff.toml (already in pyproject.toml) - Noqa F401 on __init__.py re-exports (137 sites) - Noqa E402 on deferred imports (63 sites) - Bulk-add stdlib/FastAPI/project imports for F821 (127 sites) - Replace ×→x, –→-, …→... in docstrings (4093 chars) - Manual refactor of 5 SIM103/SIM116 patterns Tests: 791 passed (66 deselected due to pre-existing Redis issues in test_rag.py) Co-authored-by: opencode <opencode@rugmunch.io>
415 lines
16 KiB
Python
415 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
FAISS-based ANN Index Manager for RMI RAG
|
|
==========================================
|
|
Replaces O(n) brute-force Redis cosine scan with sub-millisecond
|
|
FAISS HNSW / IVFFlat approximate nearest-neighbor search.
|
|
|
|
Architecture:
|
|
- Loads all vectors from Redis for each collection into a FAISS index
|
|
- Keeps index in memory; auto-rebuilds when stale
|
|
- Persists pickled indexes to /app/data/faiss/{collection}.index
|
|
- Tracks version counter in Redis key rag:idx_version:{collection}
|
|
- Invalidate on new ingestion (version bump)
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import pickle
|
|
import time
|
|
from typing import Any, Optional
|
|
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
REDIS_HOST = os.getenv("REDIS_HOST", "rmi-redis")
|
|
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
|
|
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "")
|
|
|
|
FAISS_DATA_DIR = os.getenv("FAISS_DATA_DIR", "/app/data/faiss")
|
|
|
|
# HNSW defaults
|
|
HNSW_M = 16
|
|
HNSW_EF_CONSTRUCTION = 200
|
|
HNSW_EF_SEARCH = 128
|
|
|
|
# IVFFlat defaults
|
|
IVF_LISTS_FACTOR = 40 # lists = n_vectors / factor, min 4
|
|
IVF_NPROBE = 16
|
|
|
|
# Minimum docs to use IVFFlat/HNSW; below this, flat search is fine
|
|
MIN_DOCS_FOR_ANN = 50
|
|
|
|
|
|
class ANNIndex:
|
|
"""
|
|
FAISS-backed approximate nearest-neighbor index manager.
|
|
|
|
Each collection gets its own FAISS index built from Redis-stored
|
|
vectors. The index is kept in process memory and persisted to
|
|
disk so it survives restarts.
|
|
|
|
Usage:
|
|
idx = ANNIndex()
|
|
await idx.build_index("scam_patterns")
|
|
results = idx.search(query_embedding, "scam_patterns", limit=10)
|
|
"""
|
|
|
|
_instance: Optional["ANNIndex"] = None
|
|
|
|
def __init__(self):
|
|
self._indexes: dict[str, Any] = {} # collection -> faiss index
|
|
self._id_maps: dict[str, list[str]] = {} # collection -> [doc_id, ...]
|
|
self._meta: dict[str, dict] = {} # collection -> build metadata
|
|
self._redis = None
|
|
|
|
@classmethod
|
|
def get_instance(cls) -> "ANNIndex":
|
|
if cls._instance is None:
|
|
cls._instance = cls()
|
|
return cls._instance
|
|
|
|
async def _get_redis(self):
|
|
import redis.asyncio as aioredis
|
|
|
|
if self._redis is None:
|
|
self._redis = aioredis.Redis(
|
|
host=REDIS_HOST,
|
|
port=REDIS_PORT,
|
|
password=REDIS_PASSWORD or None,
|
|
db=0,
|
|
decode_responses=True,
|
|
)
|
|
return self._redis
|
|
|
|
# ── Build ────────────────────────────────────────────────────
|
|
|
|
async def build_index(self, collection: str, force: bool = False) -> dict[str, Any]:
|
|
"""
|
|
Build (or rebuild) a FAISS index for *collection*.
|
|
|
|
Reads all documents from Redis rag:{collection}:* and builds
|
|
an HNSW or IVFFlat index depending on document count.
|
|
|
|
Returns build metadata dict.
|
|
"""
|
|
# Skip if fresh enough (unless forced)
|
|
if not force and self.is_built(collection):
|
|
version_redis = await self._get_version(collection)
|
|
version_local = self._meta.get(collection, {}).get("version", -1)
|
|
if version_redis == version_local:
|
|
logger.info(f"ANN index for {collection} is fresh (v{version_local})")
|
|
return self._meta.get(collection, {})
|
|
|
|
r = await self._get_redis()
|
|
|
|
# Fetch all document IDs
|
|
doc_ids = list(await r.smembers(f"rag:idx:{collection}"))
|
|
n = len(doc_ids)
|
|
if n == 0:
|
|
logger.warning(f"No documents found for {collection}")
|
|
self._meta[collection] = {"status": "empty", "n": 0, "collection": collection}
|
|
return self._meta[collection]
|
|
|
|
# Batch-fetch documents
|
|
keys = [f"rag:{collection}:{did}" for did in doc_ids]
|
|
pipe = r.pipeline()
|
|
for k in keys:
|
|
pipe.get(k)
|
|
raw_docs = await pipe.execute()
|
|
|
|
# Extract vectors and metadata; track dimension
|
|
vectors = []
|
|
valid_ids = []
|
|
dims = 0
|
|
for i, data in enumerate(raw_docs):
|
|
if not data:
|
|
continue
|
|
try:
|
|
doc = json.loads(data)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
vec = doc.get("vector", [])
|
|
# Handle JSON-string vectors (from hash re-embed)
|
|
if isinstance(vec, str):
|
|
try:
|
|
vec = json.loads(vec)
|
|
except (json.JSONDecodeError, TypeError):
|
|
continue
|
|
if not vec or not isinstance(vec, list):
|
|
continue
|
|
if dims == 0:
|
|
dims = len(vec)
|
|
if len(vec) != dims:
|
|
# Pad or truncate to match first vector's dimension
|
|
vec = vec + [0.0] * (dims - len(vec)) if len(vec) < dims else vec[:dims]
|
|
vectors.append(vec)
|
|
valid_ids.append(doc_ids[i])
|
|
|
|
n_valid = len(vectors)
|
|
if n_valid == 0:
|
|
logger.warning(f"No valid vectors for {collection}")
|
|
self._meta[collection] = {"status": "no_vectors", "n": 0, "collection": collection}
|
|
return self._meta[collection]
|
|
|
|
mat = np.array(vectors, dtype=np.float32)
|
|
|
|
# Choose index type
|
|
import faiss
|
|
|
|
if n_valid < MIN_DOCS_FOR_ANN:
|
|
# Flat index - exact search, small collection
|
|
index = faiss.IndexFlatIP(dims) # inner product (cosine after norm)
|
|
index_type = "flat"
|
|
else:
|
|
# Normalize vectors for cosine similarity via inner product
|
|
faiss.normalize_L2(mat)
|
|
|
|
# Try HNSW first (best quality, no training needed)
|
|
try:
|
|
index = faiss.IndexHNSWFlat(dims, HNSW_M, faiss.METRIC_INNER_PRODUCT)
|
|
index.hnsw.efConstruction = HNSW_EF_CONSTRUCTION
|
|
index.hnsw.efSearch = HNSW_EF_SEARCH
|
|
index_type = "hnsw"
|
|
logger.info(f"Building HNSW index for {collection}: {n_valid} vectors, {dims}d")
|
|
except Exception as e:
|
|
logger.warning(f"HNSW failed, falling back to IVFFlat: {e}")
|
|
# IVFFlat fallback
|
|
nlist = max(4, n_valid // IVF_LISTS_FACTOR)
|
|
quantizer = faiss.IndexFlatIP(dims)
|
|
index = faiss.IndexIVFFlat(quantizer, dims, nlist, faiss.METRIC_INNER_PRODUCT)
|
|
index.nprobe = IVF_NPROBE
|
|
index.train(mat)
|
|
index_type = "ivfflat"
|
|
|
|
# Normalize for cosine via inner product (skip if already done for HNSW path)
|
|
if index_type == "flat":
|
|
faiss.normalize_L2(mat)
|
|
|
|
index.add(mat)
|
|
|
|
# Store in memory
|
|
self._indexes[collection] = index
|
|
self._id_maps[collection] = valid_ids
|
|
|
|
version = await self._get_version(collection)
|
|
|
|
# Persist to disk
|
|
os.makedirs(FAISS_DATA_DIR, exist_ok=True)
|
|
index_path = os.path.join(FAISS_DATA_DIR, f"{collection}.index")
|
|
try:
|
|
# faiss indexes can be serialized directly
|
|
faiss.write_index(index, index_path)
|
|
# Save id_map alongside
|
|
id_map_path = os.path.join(FAISS_DATA_DIR, f"{collection}.ids")
|
|
with open(id_map_path, "wb") as f:
|
|
pickle.dump(valid_ids, f)
|
|
logger.info(f"Persisted FAISS index to {index_path}")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to persist FAISS index: {e}")
|
|
|
|
build_meta = {
|
|
"status": "built",
|
|
"collection": collection,
|
|
"n": n_valid,
|
|
"dims": dims,
|
|
"index_type": index_type,
|
|
"version": version,
|
|
"built_at": time.time(),
|
|
"persisted": os.path.exists(index_path),
|
|
}
|
|
self._meta[collection] = build_meta
|
|
logger.info(f"ANN index built: {collection} ({n_valid} docs, {dims}d, {index_type})")
|
|
return build_meta
|
|
|
|
# ── Load from disk ────────────────────────────────────────────
|
|
|
|
def _load_from_disk(self, collection: str) -> bool:
|
|
"""Try to load a persisted FAISS index and id_map from disk."""
|
|
import faiss
|
|
|
|
index_path = os.path.join(FAISS_DATA_DIR, f"{collection}.index")
|
|
id_map_path = os.path.join(FAISS_DATA_DIR, f"{collection}.ids")
|
|
|
|
if not os.path.exists(index_path) or not os.path.exists(id_map_path):
|
|
return False
|
|
|
|
try:
|
|
index = faiss.read_index(index_path)
|
|
with open(id_map_path, "rb") as f:
|
|
id_list = pickle.load(f)
|
|
self._indexes[collection] = index
|
|
self._id_maps[collection] = id_list
|
|
self._meta[collection] = {
|
|
"status": "loaded",
|
|
"collection": collection,
|
|
"n": len(id_list),
|
|
"dims": index.d,
|
|
"index_type": type(index).__name__,
|
|
"loaded_at": time.time(),
|
|
}
|
|
logger.info(f"Loaded FAISS index for {collection} from disk ({len(id_list)} vectors)")
|
|
return True
|
|
except Exception as e:
|
|
logger.warning(f"Failed to load FAISS index from disk: {e}")
|
|
return False
|
|
|
|
# ── Search ────────────────────────────────────────────────────
|
|
|
|
async def search(
|
|
self,
|
|
query_embedding: list[float],
|
|
collection: str,
|
|
limit: int = 10,
|
|
min_similarity: float = 0.0,
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
ANN search: find top-k documents similar to query_embedding.
|
|
|
|
Auto-builds the index on first search if not yet built.
|
|
Hydrates results with content/metadata from Redis.
|
|
Returns list of {id, similarity, content, metadata, source, severity} dicts.
|
|
"""
|
|
# Auto-build if needed
|
|
if not self.is_built(collection):
|
|
# Try disk first, then build from Redis (disk I/O offloaded to thread)
|
|
loaded = await asyncio.to_thread(self._load_from_disk, collection)
|
|
if not loaded:
|
|
await self.build_index(collection)
|
|
|
|
if not self.is_built(collection):
|
|
logger.warning(f"No ANN index available for {collection}")
|
|
return []
|
|
|
|
index = self._indexes[collection]
|
|
id_list = self._id_maps[collection]
|
|
dims = index.d
|
|
|
|
# Prepare query vector
|
|
q = np.array([query_embedding[:dims]], dtype=np.float32)
|
|
# Pad if query is shorter
|
|
if q.shape[1] < dims:
|
|
q = np.pad(q, ((0, 0), (0, dims - q.shape[1])))
|
|
# Truncate if query is longer
|
|
if q.shape[1] > dims:
|
|
q = q[:, :dims]
|
|
|
|
# Normalize for cosine via inner product
|
|
import faiss
|
|
|
|
faiss.normalize_L2(q)
|
|
|
|
# Search
|
|
search_k = min(limit * 2, len(id_list)) # fetch extra for filtering
|
|
distances, indices = index.search(q, search_k)
|
|
|
|
# Collect matching doc IDs for hydration
|
|
raw_hits = []
|
|
for rank, (dist, idx) in enumerate(zip(distances[0], indices[0], strict=False)):
|
|
if idx < 0:
|
|
continue # FAISS returns -1 for empty slots
|
|
sim = float(dist) # inner product on normalized vectors = cosine similarity
|
|
if sim < min_similarity:
|
|
continue
|
|
doc_id = id_list[idx] if idx < len(id_list) else f"unknown_{idx}"
|
|
raw_hits.append((doc_id, sim, rank))
|
|
|
|
if not raw_hits:
|
|
return []
|
|
|
|
# Hydrate from Redis - batch-fetch all matched docs
|
|
r = await self._get_redis()
|
|
keys = [f"rag:{collection}:{doc_id}" for doc_id, _, _ in raw_hits]
|
|
pipe = r.pipeline()
|
|
for k in keys:
|
|
pipe.get(k)
|
|
raw_docs = await pipe.execute()
|
|
|
|
results = []
|
|
for (doc_id, sim, rank), data in zip(raw_hits, raw_docs, strict=False):
|
|
result = {
|
|
"id": doc_id,
|
|
"similarity": round(sim, 4),
|
|
"rank": rank,
|
|
}
|
|
if data:
|
|
try:
|
|
doc = json.loads(data)
|
|
result["content"] = doc.get("content", "")[:500]
|
|
result["metadata"] = doc.get("metadata", {})
|
|
result["source"] = doc.get("source", "")
|
|
result["severity"] = doc.get("severity", "")
|
|
except json.JSONDecodeError:
|
|
pass
|
|
results.append(result)
|
|
|
|
results.sort(key=lambda x: x["similarity"], reverse=True)
|
|
return results[:limit]
|
|
|
|
# ── Status ────────────────────────────────────────────────────
|
|
|
|
def is_built(self, collection: str) -> bool:
|
|
"""Return True if an in-memory index exists for the collection."""
|
|
return collection in self._indexes and collection in self._id_maps
|
|
|
|
def stats(self) -> dict[str, Any]:
|
|
"""Return stats for all loaded indexes."""
|
|
out = {}
|
|
for coll in set(list(self._indexes.keys()) + list(self._meta.keys())):
|
|
idx = self._indexes.get(coll)
|
|
out[coll] = {
|
|
"built": coll in self._indexes,
|
|
"n_vectors": len(self._id_maps.get(coll, [])),
|
|
"dims": idx.d if idx else 0,
|
|
"index_type": type(idx).__name__ if idx else "none",
|
|
**self._meta.get(coll, {}),
|
|
}
|
|
return out
|
|
|
|
# ── Version tracking ─────────────────────────────────────────
|
|
|
|
async def _get_version(self, collection: str) -> int:
|
|
"""Get the current version counter from Redis."""
|
|
r = await self._get_redis()
|
|
val = await r.get(f"rag:idx_version:{collection}")
|
|
return int(val) if val else 0
|
|
|
|
async def bump_version(self, collection: str) -> int:
|
|
"""
|
|
Bump the version counter (call after new ingestion).
|
|
This signals that the index needs rebuilding.
|
|
"""
|
|
r = await self._get_redis()
|
|
new_ver = await r.incr(f"rag:idx_version:{collection}")
|
|
# Invalidate in-memory index
|
|
self._indexes.pop(collection, None)
|
|
self._id_maps.pop(collection, None)
|
|
logger.info(f"Version bumped for {collection}: now v{new_ver}")
|
|
return new_ver
|
|
|
|
# ── Invalidate ────────────────────────────────────────────────
|
|
|
|
def invalidate(self, collection: str) -> None:
|
|
"""Drop the in-memory index for *collection* (next search will rebuild)."""
|
|
self._indexes.pop(collection, None)
|
|
self._id_maps.pop(collection, None)
|
|
self._meta.pop(collection, None)
|
|
logger.info(f"Invalidated ANN index for {collection}")
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
# Singleton accessor
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
|
|
_ann_index: ANNIndex | None = None
|
|
|
|
|
|
def get_ann_index() -> ANNIndex:
|
|
"""Return the singleton ANNIndex instance."""
|
|
global _ann_index
|
|
if _ann_index is None:
|
|
_ann_index = ANNIndex()
|
|
return _ann_index
|