Update rag_engine.py

rag_engine.py

@@ -1,19 +1,21 @@
 """
-rag_engine.py — Multimodal RAG Engine with Multi-File Support & Conversation Memory
+rag_engine.py — Multimodal RAG Engine with Multi-File Support, Reranking & Memory
 Supports: PDF, TXT, DOCX, CSV, XLSX, Images (JPG/PNG/WEBP)
 Features: Up to 5 simultaneous files, per-file removal, additive indexing
 Memory: sliding window of last 6 exchanges
 
-KEY CHANGES (v4 — Multi-File):
-  1. Additive indexing — new uploads ADD to the vectorstore, not replace it.
-  2. Per-file chunk tracking — each file's chunk IDs are stored for clean removal.
-  3. remove_file(filename) — deletes a specific file's chunks from the vectorstore.
-  4. MAX_FILES = 5 — enforced in ingest_file().
-  5. _generate() is multi-file aware — system prompt lists all loaded files.
-  6. query() scales retrieval k with number of files for cross-doc coverage.
-  7. Memory is NOT cleared on upload (user may be chatting about multiple docs).
-
-Keeps all v3 image fixes: OCR, color analysis, BLIP raw bytes, VLM descriptions.
+KEY CHANGES (v5 — Cross-Encoder Reranking):
+  1. Cross-encoder reranker (ms-marco-MiniLM-L-6-v2) scores every retrieved
+     chunk for true semantic relevance to the query — not just embedding distance.
+  2. Over-fetches 12+ candidates from the vectorstore, then reranks to pick
+     the top-k most relevant chunks for the LLM context.
+  3. Graceful fallback — if the reranker fails to load, uses original order.
+
+Previous features preserved:
+  - Additive indexing, per-file removal, MAX_FILES=5
+  - Multi-file aware generation, cross-doc coverage
+  - OCR, color analysis, BLIP raw bytes, VLM descriptions for images
+  - Conversation memory (6-exchange sliding window)
 """
 
 import os
@@ -44,9 +46,11 @@ logger = logging.getLogger(__name__)
 
 # ── Constants ────────────────────────────────────────────────────────────────
 EMBED_MODEL     = "all-MiniLM-L6-v2"
+RERANK_MODEL    = "cross-encoder/ms-marco-MiniLM-L-6-v2"  # ~80MB, CPU-friendly
 CHUNK_SIZE      = 600
 CHUNK_OVERLAP   = 100
-TOP_K           = 4
+TOP_K           = 4       # final chunks sent to LLM after reranking
+RERANK_FETCH_K  = 12      # over-fetch this many candidates for reranking
 COLLECTION_NAME = "docmind_multimodal"
 HF_API_URL      = "https://router.huggingface.co/v1/chat/completions"
 MEMORY_WINDOW   = 6   # number of past Q&A pairs to keep
@@ -119,6 +123,7 @@ def _classify_color(r: int, g: int, b: int) -> str:
 class RAGEngine:
     def __init__(self):
         self._embeddings:  Optional[HuggingFaceEmbeddings] = None
+        self._reranker = None  # lazy-loaded cross-encoder
         self._vectorstore: Optional[Chroma] = None
         self._splitter = RecursiveCharacterTextSplitter(
             chunk_size=CHUNK_SIZE,
@@ -140,6 +145,58 @@ class RAGEngine:
             )
         return self._embeddings
 
+    @property
+    def reranker(self):
+        """Lazy-load the cross-encoder reranker (~80MB, CPU-friendly)."""
+        if self._reranker is None:
+            try:
+                from sentence_transformers import CrossEncoder
+                logger.info(f"Loading reranker model: {RERANK_MODEL}...")
+                self._reranker = CrossEncoder(RERANK_MODEL, max_length=512)
+                logger.info("Reranker loaded successfully.")
+            except Exception as e:
+                logger.warning(f"Failed to load reranker: {e}. Will skip reranking.")
+                self._reranker = False  # sentinel: don't retry
+        return self._reranker if self._reranker is not False else None
+
+    def _rerank_documents(self, question: str, docs: List[Document], top_k: int) -> List[Document]:
+        """Score and reorder documents using the cross-encoder reranker."""
+        if not docs:
+            return docs
+
+        ranker = self.reranker
+        if ranker is None:
+            # Reranker unavailable — fall back to original order
+            logger.info("Reranker not available, using original retrieval order.")
+            return docs[:top_k]
+
+        # Build query-document pairs for the cross-encoder
+        pairs = [(question, doc.page_content) for doc in docs]
+
+        try:
+            scores = ranker.predict(pairs)
+
+            # Pair each doc with its rerank score
+            scored = list(zip(docs, scores))
+            scored.sort(key=lambda x: x[1], reverse=True)
+
+            reranked = [doc for doc, score in scored[:top_k]]
+
+            # Log the reranking effect
+            original_sources = [d.metadata.get("source", "?")[:30] for d in docs[:top_k]]
+            reranked_sources = [d.metadata.get("source", "?")[:30] for d in reranked]
+            top_scores = [f"{s:.3f}" for _, s in scored[:top_k]]
+            logger.info(
+                f"Reranked {len(docs)} candidates → top {top_k}. "
+                f"Scores: {top_scores}. "
+                f"Before: {original_sources}, After: {reranked_sources}"
+            )
+
+            return reranked
+        except Exception as e:
+            logger.warning(f"Reranking failed: {e}. Using original order.")
+            return docs[:top_k]
+
     # ── Memory ───────────────────────────────────────────────────────────────
 
     def clear_memory(self):
@@ -805,15 +862,25 @@ class RAGEngine:
         sources = []
 
         try:
-            # Scale retrieval k with number of files for cross-doc coverage
-            k       = min(TOP_K + len(self._documents) - 1, 6)
-            fetch_k = k * 3
+            # ── Step 1: Over-fetch candidates ────────────────────────────────
+            # Retrieve more candidates than needed so the reranker can pick
+            # the truly relevant ones. Scale with number of loaded files.
+            num_files = len(self._documents)
+            fetch_k   = max(RERANK_FETCH_K, RERANK_FETCH_K + (num_files - 1) * 2)
+            initial_k = fetch_k  # MMR will return this many diverse candidates
 
             retriever = self._vectorstore.as_retriever(
                 search_type="mmr",
-                search_kwargs={"k": k, "fetch_k": fetch_k},
+                search_kwargs={"k": initial_k, "fetch_k": fetch_k * 2},
             )
-            docs = retriever.invoke(question)
+            candidate_docs = retriever.invoke(question)
+
+            # ── Step 2: Rerank with cross-encoder ────────────────────────────
+            # The cross-encoder scores each (query, chunk) pair for true
+            # semantic relevance — much more accurate than embedding distance.
+            final_k = min(TOP_K + num_files - 1, 6)
+            docs = self._rerank_documents(question, candidate_docs, top_k=final_k)
+
             context = "\n\n---\n\n".join(
                 f"[Chunk {i+1} | source: {d.metadata.get('source', '?')} | type: {d.metadata.get('type','text')}]\n{d.page_content}"
                 for i, d in enumerate(docs)