# vector_store.py — HF Spaces variant # Uses ChromaDB in-memory client only. # - Persistent Chroma (disk) is removed: HF Spaces filesystem is ephemeral # and resets on every restart, so persistence provides no benefit. # - Pinecone is removed from the default UI but the class is kept as an # optional import so power users can re-enable it with their own API key. # # Self-healing capabilities: # 1. Retrieval quality scoring — real L2 distances from Chroma (not fake 0.0) # 2. Relevance gate — filters out chunks whose distance exceeds the # threshold before they reach the LLM # 3. Fallback / retry retrieval — if too few chunks pass the gate, the search # is retried with a relaxed threshold and/or # plain similarity (no MMR) to maximise recall import os import logging import chromadb from langchain_chroma import Chroma from langchain.schema import Document from dotenv import load_dotenv load_dotenv() logger = logging.getLogger(__name__) # Clear any shared system cache from previous runs (belt-and-suspenders) try: chromadb.api.client.SharedSystemClient.clear_system_cache() except Exception: pass # ── Relevance gate thresholds ──────────────────────────────────────────────── # Chroma returns L2 distances: 0.0 = perfect match, higher = worse. # all-MiniLM-L6-v2 embeddings live in a unit-normalised space, so practical # L2 distances fall in [0, 2]. Empirically: # < 0.5 → strong match # 0.5–1.0 → acceptable # > 1.0 → weak / off-topic RELEVANCE_THRESHOLD: float = float(os.getenv("RELEVANCE_THRESHOLD", "1.0")) RELAXED_THRESHOLD: float = float(os.getenv("RELAXED_THRESHOLD", "1.4")) # Minimum number of chunks that must pass the gate before we skip the retry. MIN_PASSING_CHUNKS: int = int(os.getenv("MIN_PASSING_CHUNKS", "1")) class VectorStoreManager: @staticmethod def create_store(store_type: str, documents: list, embedding_function, **kwargs): if store_type == "Chroma": return ChromaStore(documents, embedding_function) else: raise ValueError( f"Unsupported store type for HF Spaces: '{store_type}'. " "Only 'Chroma' (in-memory) is supported in this deployment." ) class ChromaStore: def __init__(self, documents: list, embedding_function): # Use a brand-new ephemeral client every time so there is # zero chance of inheriting chunks from a previous session. # chromadb.EphemeralClient() is the modern in-memory API; # fall back to chromadb.Client() for older chromadb versions. try: _client = chromadb.EphemeralClient() except AttributeError: _client = chromadb.Client() # Delete the collection if it somehow already exists on this client try: _client.delete_collection("docquest") except Exception: pass self.store = Chroma( client=_client, collection_name="docquest", embedding_function=embedding_function, ) self.store.add_documents(documents) self._embedding_function = embedding_function # ── Internal helpers ───────────────────────────────────────────────────── def _collection_size(self) -> int: """Return the number of documents currently in the collection.""" try: collection = self.store._collection if hasattr(collection, "count"): return collection.count() except Exception: pass try: result = self.store.get() if isinstance(result, dict) and "ids" in result: return len(result["ids"]) except Exception: pass return 0 def _dedupe(self, doc_score_pairs: list) -> list: """Remove duplicate chunks by content fingerprint.""" seen = set() unique = [] for doc, score in doc_score_pairs: fingerprint = ( doc.page_content.strip()[:200], doc.metadata.get("filename"), doc.metadata.get("chunk_id"), ) if fingerprint not in seen: seen.add(fingerprint) unique.append((doc, score)) return unique def _embed_query(self, query: str) -> list: """ Embed *query* using whatever embedding object was passed at construction. LangChain embedding wrappers expose two calling conventions: • embed_query(str) → list[float] (standard interface) • embed_documents([str]) → list[list[float]] (batch interface) We try embed_query first, then fall back to embed_documents. """ ef = self._embedding_function if hasattr(ef, "embed_query"): return ef.embed_query(query) if hasattr(ef, "embed_documents"): return ef.embed_documents([query])[0] # Last resort: the object itself might be callable return ef(query) def _apply_gate( self, doc_score_pairs: list, threshold: float ) -> list: """ Return only pairs whose L2 distance is *below* threshold. Lower distance = more similar, so we KEEP scores < threshold. """ passing = [(doc, score) for doc, score in doc_score_pairs if score < threshold] blocked = len(doc_score_pairs) - len(passing) if blocked: logger.debug( "Relevance gate blocked %d/%d chunks (threshold=%.3f)", blocked, len(doc_score_pairs), threshold, ) return passing # ── MMR with real scores ───────────────────────────────────────────────── def _mmr_with_scores(self, query: str, k: int, fetch_k: int) -> list: """ Run MMR directly against the underlying Chroma collection and return (Document, L2_distance) tuples. WHY NOT retriever.invoke()? LangChain's MMR *retriever* wraps max_marginal_relevance_search(), which returns plain Documents with no distance info. We bypass the retriever and call max_marginal_relevance_search_with_score_by_vector() directly so we get real distances from Chroma — not the fake 0.0 patch from before. """ query_embedding = self._embed_query(query) # This method returns List[Tuple[Document, float]] where float is the # L2 distance between the query vector and each retrieved chunk. results = self.store.max_marginal_relevance_search_with_score_by_vector( embedding=query_embedding, k=k, fetch_k=fetch_k, ) return results # already List[(Document, float)] # ── Public search interface ────────────────────────────────────────────── def search(self, query: str, k: int = 4) -> list: """ Retrieve top-k most relevant unique chunks with self-healing logic. Pipeline -------- 1. MMR search → real (Document, L2_distance) pairs 2. Relevance gate (threshold = RELEVANCE_THRESHOLD) → if enough chunks pass → dedupe & return 3. Retry with relaxed threshold → if still not enough → plain similarity search as final fallback 4. Return whatever we have (never empty if the store has documents) All returned scores are genuine L2 distances in [0, 2], where: 0.0 = perfect match, 2.0 = maximally dissimilar. app.py's _score_to_pct() converts these correctly via (1 - d/2) * 100. """ size = self._collection_size() actual_k = max(1, min(k, size)) if size > 0 else k fetch_k = min(actual_k * 3, size) if size > 0 else actual_k * 3 # ── Step 1: MMR with real scores ───────────────────────────────────── mmr_results = [] try: mmr_results = self._mmr_with_scores(query, k=actual_k, fetch_k=fetch_k) except Exception as exc: logger.warning("MMR search failed (%s); will use similarity fallback.", exc) # ── Step 2: Apply strict relevance gate ────────────────────────────── if mmr_results: gated = self._apply_gate(mmr_results, threshold=RELEVANCE_THRESHOLD) if len(gated) >= MIN_PASSING_CHUNKS: logger.debug( "MMR + strict gate: %d/%d chunks passed", len(gated), len(mmr_results) ) return self._dedupe(gated)[:actual_k] # ── Step 3a: Gate passed too few — retry MMR with relaxed threshold logger.debug( "Only %d chunk(s) passed strict gate; retrying with relaxed threshold %.3f", len(gated), RELAXED_THRESHOLD, ) relaxed = self._apply_gate(mmr_results, threshold=RELAXED_THRESHOLD) if len(relaxed) >= MIN_PASSING_CHUNKS: return self._dedupe(relaxed)[:actual_k] # ── Step 3b: MMR failed or gate still blocking — plain similarity search logger.debug("Falling back to plain similarity_search_with_score.") try: sim_results = self.store.similarity_search_with_score(query, k=actual_k) sim_gated = self._apply_gate(sim_results, threshold=RELAXED_THRESHOLD) if sim_gated: return self._dedupe(sim_gated)[:actual_k] # Nothing passed even the relaxed gate — return raw similarity results # so the UI can at least show something and let the user judge. logger.warning( "All %d similarity chunks failed relaxed gate; returning ungated results.", len(sim_results), ) return self._dedupe(sim_results)[:actual_k] except Exception as exc: raise RuntimeError(f"Vector store search failed: {exc}") from exc