Spaces:
Running
Running
| """Persistent vector storage: ChromaDB (text + metadata) with FAISS or Chroma similarity search.""" | |
| from typing import Dict, List, Optional | |
| import time | |
| import numpy as np | |
| import chromadb | |
| from chromadb.config import Settings as ChromaSettings | |
| from backend.core import settings, get_logger | |
| from backend.services.embeddings import get_embedding_service | |
| from backend.services.pdf_processor import ProcessedDocument | |
| logger = get_logger(__name__) | |
| class VectorStoreService: | |
| # Class-level defaults keep partially constructed instances (e.g. in tests) on the Chroma path. | |
| backend = "chroma" | |
| _faiss = None | |
| def __init__(self): | |
| self.client = chromadb.PersistentClient(path=settings.vectorstore_path, | |
| settings=ChromaSettings(anonymized_telemetry=False)) | |
| self.embedding_service = None | |
| # embedding_function=None: we always pass our own vectors. Without this, Chroma | |
| # would silently download a second (ONNX) embedding model when none is supplied. | |
| self.collection = self.client.get_or_create_collection( | |
| name="documents", metadata={"hnsw:space": "cosine"}, embedding_function=None) | |
| self.registry = self.client.get_or_create_collection( | |
| name="document_registry", metadata={"description": "Document metadata"}, | |
| embedding_function=None) | |
| self._documents_cache = (0.0, []) | |
| if settings.vector_backend == "faiss": | |
| try: | |
| from backend.services.faiss_index import FaissIndex | |
| self._faiss = FaissIndex() | |
| self.backend = "faiss" | |
| except ImportError: | |
| logger.warning("VECTOR_BACKEND=faiss but faiss-cpu is not installed; " | |
| "falling back to ChromaDB search. Fix: pip install faiss-cpu") | |
| logger.info(f"Vector search backend: {self.backend}") | |
| # ── FAISS <-> Chroma synchronisation ───────────────────────────────────── | |
| def _ensure_faiss(self) -> None: | |
| """Chroma is the source of truth: rebuild FAISS whenever the vector counts differ.""" | |
| if self._faiss is not None and self._faiss.ntotal != self.collection.count(): | |
| self._rebuild_faiss() | |
| def _rebuild_faiss(self, page_size: int = 2000) -> None: | |
| started = time.perf_counter() | |
| self._faiss.clear() | |
| offset = 0 | |
| while True: | |
| page = self.collection.get(limit=page_size, offset=offset, | |
| include=["embeddings", "metadatas"]) | |
| ids = page["ids"] | |
| if len(ids) == 0: | |
| break | |
| metas = page["metadatas"] | |
| self._faiss.add(ids, np.asarray(page["embeddings"], dtype="float32"), | |
| [(m or {}).get("doc_id", "") for m in metas], | |
| [(m or {}).get("owner_id", "") for m in metas]) | |
| offset += len(ids) | |
| logger.info(f"FAISS index rebuilt from ChromaDB | vectors={self._faiss.ntotal} " | |
| f"| {(time.perf_counter() - started) * 1000:.0f} ms") | |
| def index_document(self, processed_doc: ProcessedDocument, owner_id: str = "") -> bool: | |
| if not processed_doc.chunks: | |
| return False | |
| chunks = processed_doc.chunks | |
| try: | |
| if self.embedding_service is None: | |
| self.embedding_service = get_embedding_service() | |
| vectors = self.embedding_service.embed_texts([c.text for c in chunks]) | |
| if self._faiss is not None: | |
| self._ensure_faiss() # bring FAISS level with Chroma before adding to both | |
| self.collection.upsert(ids=[c.chunk_id for c in chunks], | |
| embeddings=vectors.tolist(), documents=[c.text for c in chunks], | |
| metadatas=[{"doc_id": c.doc_id, "doc_name": c.doc_name, | |
| "page_number": c.page_number, "chunk_index": c.chunk_index, | |
| "char_count": c.char_count, "owner_id": owner_id} for c in chunks]) | |
| if self._faiss is not None: | |
| self._faiss.add([c.chunk_id for c in chunks], vectors, | |
| [c.doc_id for c in chunks], [owner_id] * len(chunks)) | |
| # Registry is the commit marker; write it only after every vector succeeds. | |
| # The registry only stores metadata; a constant 1-d vector satisfies Chroma. | |
| self.registry.upsert(ids=[processed_doc.doc_id], embeddings=[[0.0]], | |
| documents=[processed_doc.filename], | |
| metadatas=[{"doc_id": processed_doc.doc_id, "filename": processed_doc.filename, | |
| "total_pages": processed_doc.total_pages, "total_chunks": len(chunks), | |
| "file_size_kb": processed_doc.file_size_kb, "owner_id": owner_id}]) | |
| self.invalidate_documents_cache() | |
| return True | |
| except Exception: | |
| self.collection.delete(where={"doc_id": processed_doc.doc_id}) | |
| if self._faiss is not None: | |
| self._faiss.remove_doc(processed_doc.doc_id) | |
| raise | |
| def list_documents(self, owner_id: Optional[str] = None) -> List[Dict]: | |
| try: | |
| cached_at, cached_docs = self._documents_cache | |
| if owner_id is None and time.monotonic() - cached_at < 5: | |
| return list(cached_docs) | |
| result = self.registry.get(include=["metadatas"]) | |
| docs = [] | |
| for doc_id, meta in zip(result.get("ids", []), result.get("metadatas", [])): | |
| meta = meta or {} | |
| if owner_id is not None and meta.get("owner_id", "") != owner_id: | |
| continue | |
| docs.append({k: meta.get(k, default) for k, default in { | |
| "doc_id": doc_id, "filename": "", "total_pages": 0, | |
| "total_chunks": 0, "file_size_kb": 0, "owner_id": ""}.items()}) | |
| if owner_id is None: | |
| self._documents_cache = (time.monotonic(), docs) | |
| return docs | |
| except Exception: | |
| # Do not hide storage failures: returning [] would look like "no documents". | |
| logger.exception("Failed to list documents") | |
| raise | |
| def search(self, query: str, doc_ids: Optional[List[str]] = None, | |
| top_k: int = None, owner_id: Optional[str] = None) -> List[Dict]: | |
| top_k = top_k or settings.top_k_results | |
| docs = self.list_documents(owner_id) | |
| selected = [d["doc_id"] for d in docs if doc_ids is None or d["doc_id"] in doc_ids] | |
| count = self.collection.count() | |
| if not selected or top_k <= 0 or count == 0: | |
| return [] | |
| if self.embedding_service is None: | |
| self.embedding_service = get_embedding_service() | |
| n_candidates = min(max(top_k, 20 if settings.hybrid_search or settings.enable_rerank else top_k), count) | |
| query_vector = self.embedding_service.embed_query(query) | |
| if self._faiss is not None: | |
| rows = self._faiss_candidates(query_vector, n_candidates, selected, owner_id) | |
| else: | |
| rows = self._chroma_candidates(query_vector, n_candidates, selected, owner_id) | |
| if settings.hybrid_search: | |
| from rank_bm25 import BM25Okapi | |
| corpus = [r["text"].lower().split() for r in rows] | |
| if corpus: | |
| scores = BM25Okapi(corpus).get_scores(query.lower().split()) | |
| lexical = sorted(range(len(rows)), key=lambda i: scores[i], reverse=True) | |
| vec = sorted(range(len(rows)), key=lambda i: rows[i]["similarity"], reverse=True) | |
| rrf = {} | |
| for ranking in (vec, lexical): | |
| for rank, ix in enumerate(ranking, 1): rrf[ix] = rrf.get(ix, 0) + 1/(60+rank) | |
| rows = [rows[i] for i in sorted(rrf, key=rrf.get, reverse=True)] | |
| if settings.enable_rerank and rows: | |
| from sentence_transformers import CrossEncoder | |
| if not hasattr(self, "reranker"): | |
| self.reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2") | |
| scores = self.reranker.predict([(query, row["text"]) for row in rows]) | |
| rows = [row for _, row in sorted(zip(scores, rows), reverse=True, key=lambda p: p[0])] | |
| return rows[:top_k] | |
| def _chroma_candidates(self, query_vector, n: int, selected: List[str], | |
| owner_id: Optional[str]) -> List[Dict]: | |
| filters = {"doc_id": {"$in": selected}} | |
| if owner_id is not None: | |
| filters = {"$and": [filters, {"owner_id": owner_id}]} | |
| result = self.collection.query(query_embeddings=[query_vector.tolist()], n_results=n, | |
| where=filters, include=["documents", "metadatas", "distances"]) | |
| rows = [] | |
| for text, meta, distance in zip(result["documents"][0], result["metadatas"][0], result["distances"][0]): | |
| if owner_id is not None and meta.get("owner_id", "") != owner_id: | |
| continue | |
| rows.append({"text": text, "metadata": meta, | |
| "similarity": round(max(0.0, min(1.0, 1.0 - float(distance))), 4), | |
| "doc_name": meta.get("doc_name", ""), "page_number": meta.get("page_number", 0)}) | |
| return rows | |
| def _faiss_candidates(self, query_vector, n: int, selected: List[str], | |
| owner_id: Optional[str]) -> List[Dict]: | |
| """FAISS finds the nearest chunk ids; Chroma supplies their text and metadata.""" | |
| self._ensure_faiss() | |
| hits = self._faiss.search(query_vector, n, doc_ids=set(selected), owner_id=owner_id) | |
| if not hits: | |
| return [] | |
| got = self.collection.get(ids=[chunk_id for chunk_id, _ in hits], include=["documents", "metadatas"]) | |
| by_id = dict(zip(got["ids"], zip(got["documents"], got["metadatas"]))) | |
| rows = [] | |
| for chunk_id, score in hits: | |
| if chunk_id not in by_id: # removed between the FAISS search and the Chroma lookup | |
| continue | |
| text, meta = by_id[chunk_id] | |
| rows.append({"text": text, "metadata": meta, | |
| "similarity": round(max(0.0, min(1.0, score)), 4), | |
| "doc_name": meta.get("doc_name", ""), "page_number": meta.get("page_number", 0)}) | |
| return rows | |
| def delete_document(self, doc_id: str, owner_id: Optional[str] = None) -> bool: | |
| docs = self.list_documents(owner_id) | |
| if not any(d["doc_id"] == doc_id for d in docs): return False | |
| self.collection.delete(where={"doc_id": doc_id}) | |
| if self._faiss is not None: | |
| self._faiss.remove_doc(doc_id) | |
| self.registry.delete(ids=[doc_id]) | |
| self.invalidate_documents_cache() | |
| return True | |
| def invalidate_documents_cache(self) -> None: | |
| self._documents_cache = (0.0, []) | |
| def document_exists(self, doc_id: str, owner_id: Optional[str] = None) -> bool: | |
| return any(d["doc_id"] == doc_id for d in self.list_documents(owner_id)) | |
| _vector_store_instance = None | |
| def get_vector_store() -> VectorStoreService: | |
| global _vector_store_instance | |
| if _vector_store_instance is None: _vector_store_instance = VectorStoreService() | |
| return _vector_store_instance | |