Spaces:
Running
Running
File size: 11,295 Bytes
5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 5959a78 eb89f47 | 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 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | """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
|