Spaces:
Sleeping
Sleeping
| """ | |
| backend/database/vector_store.py | |
| =================================== | |
| High-level vector store abstraction used by the RAG pipeline. | |
| Stores both vectors and associated metadata (chunk dicts). | |
| """ | |
| import json | |
| import os | |
| from typing import Any, Dict, List, Optional | |
| import numpy as np | |
| from backend.database.faiss_db import FAISSDatabase | |
| from backend.utils.helper import ensure_dir, save_json, load_json | |
| from backend.utils.logger import get_logger | |
| logger = get_logger(__name__) | |
| class VectorStore: | |
| """ | |
| Combines a FAISS index with a metadata store. | |
| Provides add / search / clear operations. | |
| """ | |
| def __init__(self, index_path: str = None, meta_path: str = None): | |
| self.index_path = index_path or "data/embeddings/faiss.index" | |
| self.meta_path = meta_path or "data/embeddings/metadata.json" | |
| self._db: Optional[FAISSDatabase] = None | |
| self._meta: List[Dict] = [] | |
| # ── Indexing ────────────────────────────────────────────── | |
| def add(self, vectors: np.ndarray, metadata: List[Dict]) -> None: | |
| """ | |
| Add vectors and their metadata to the store. | |
| Args: | |
| vectors: np.ndarray of shape (N, D) | |
| metadata: List of N dicts (chunk info, text, timestamps…) | |
| """ | |
| assert len(vectors) == len(metadata), "Vector/metadata length mismatch" | |
| dim = vectors.shape[1] | |
| self._db = FAISSDatabase(dim=dim, index_path=self.index_path) | |
| self._db.build(vectors) | |
| self._meta = metadata | |
| logger.info(f"VectorStore: {len(metadata)} items indexed") | |
| def save(self) -> None: | |
| """Persist index and metadata to disk.""" | |
| if self._db: | |
| self._db.save() | |
| save_json(self._meta, self.meta_path) | |
| logger.info("VectorStore saved to disk") | |
| def load(self) -> bool: | |
| """Load index and metadata from disk.""" | |
| if not os.path.exists(self.index_path): | |
| return False | |
| # Peek at metadata to get dimension | |
| if os.path.exists(self.meta_path): | |
| self._meta = load_json(self.meta_path) | |
| # Detect dim from a dummy load | |
| import faiss | |
| idx = faiss.read_index(self.index_path) | |
| dim = idx.d | |
| self._db = FAISSDatabase(dim=dim, index_path=self.index_path) | |
| self._db.load() | |
| logger.info(f"VectorStore loaded: {len(self._meta)} items") | |
| return True | |
| # ── Search ──────────────────────────────────────────────── | |
| def search(self, query_vector: np.ndarray, top_k: int = 5) -> List[Dict]: | |
| """ | |
| Search the store and return metadata for top-k matches. | |
| Returns: | |
| List of metadata dicts with an added 'score' field. | |
| """ | |
| if not self._db or not self._db.is_ready: | |
| logger.warning("VectorStore not ready. Call add() or load() first.") | |
| return [] | |
| distances, indices = self._db.search(query_vector, top_k) | |
| results = [] | |
| for dist, idx in zip(distances, indices): | |
| if 0 <= idx < len(self._meta): | |
| item = dict(self._meta[idx]) | |
| item["score"] = float(dist) | |
| results.append(item) | |
| return results | |
| # ── Utilities ───────────────────────────────────────────── | |
| def clear(self) -> None: | |
| """Reset the store.""" | |
| self._db = None | |
| self._meta = [] | |
| logger.info("VectorStore cleared") | |
| def size(self) -> int: | |
| return self._db.size if self._db else 0 | |
| def is_empty(self) -> bool: | |
| return self.size == 0 | |