core/vector_store.py

# 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