Update src/qa.py

src/qa.py

@@ -89,24 +89,48 @@ REASONING_PROMPT = (
 )
 
 # ==========================================================
-# 5️⃣ Retrieval — FAISS + Cosine Re-Rank + Neighbor Fill
+# 5️⃣ Retrieval — FAISS + Re-rank + Neighbor Fill (Auto-Healing)
 # ==========================================================
+from vectorstore import build_faiss_index
+
 def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5,
-                    min_similarity: float = 0.6, candidate_multiplier: int = 3):
-    """Select top chunks via FAISS, rerank by cosine similarity, fill gaps with neighbors."""
+                    min_similarity: float = 0.6, candidate_multiplier: int = 3,
+                    embeddings: list = None):
+    """
+    Re-rank and optionally fill with neighbors for context continuity.
+    Auto-detects and rebuilds FAISS index if dimension mismatch occurs.
+    """
+
     if not index or not chunks:
+        print("⚠️ No FAISS index or chunks provided — returning empty result.")
         return []
 
     try:
+        # Encode query embedding
         q_emb = _query_model.encode(
-            [f"query: {query.strip()}"], convert_to_numpy=True, normalize_embeddings=True
+            [f"query: {query.strip()}"],
+            convert_to_numpy=True,
+            normalize_embeddings=True
         )[0]
 
-        # 1️⃣ Initial FAISS search
-        distances, indices = index.search(np.array([q_emb]).astype("float32"), top_k * candidate_multiplier)
-        candidate_indices = list(dict.fromkeys(indices[0]))  # dedup, preserve order
-
-        # 2️⃣ Compute true cosine similarity for rerank
+        # ✅ Sanity check: dimension match between query and FAISS index
+        if hasattr(index, "d") and q_emb.shape[0] != index.d:
+            print(f"⚠️ FAISS index dimension mismatch: index={index.d}, query={q_emb.shape[0]}")
+            if embeddings:
+                print("🔄 Rebuilding FAISS index to match embedding dimensions...")
+                index = build_faiss_index(embeddings)
+                print("✅ FAISS index successfully rebuilt.")
+            else:
+                print("❌ No embeddings available to rebuild FAISS index.")
+                return []
+
+        # Step 1️⃣ — Initial FAISS retrieval
+        num_candidates = max(top_k * candidate_multiplier, top_k + 2)
+        distances, indices = index.search(np.array([q_emb]).astype("float32"), num_candidates)
+        candidate_indices = [int(i) for i in indices[0] if i >= 0]
+        candidate_indices = list(dict.fromkeys(candidate_indices))  # de-dupe
+
+        # Step 2️⃣ — Re-rank by cosine similarity
         doc_embs = _query_model.encode(
             [f"passage: {chunks[i]}" for i in candidate_indices],
             convert_to_numpy=True,
@@ -115,28 +139,31 @@ def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5,
         sims = cosine_similarity([q_emb], doc_embs)[0]
         ranked = sorted(zip(candidate_indices, sims), key=lambda x: x[1], reverse=True)
 
-        # 3️⃣ Keep only chunks meeting threshold
-        filtered = [idx for idx, sim in ranked if sim >= min_similarity][:top_k]
+        # Step 3️⃣ — Filter by similarity threshold
+        filtered = [idx for idx, sim in ranked if sim >= min_similarity]
+        if len(filtered) > top_k:
+            filtered = filtered[:top_k]
 
-        # 4️⃣ Neighbor fill if not enough
+        # Step 4️⃣ — Neighbor fill (if not enough)
         if len(filtered) < top_k:
             expanded = set(filtered)
             for idx in filtered:
-                for nb in [idx - 1, idx + 1]:
-                    if 0 <= nb < len(chunks):
-                        expanded.add(nb)
+                for neighbor in [idx - 1, idx + 1]:
+                    if 0 <= neighbor < len(chunks):
+                        expanded.add(neighbor)
                         if len(expanded) >= top_k:
                             break
                 if len(expanded) >= top_k:
                     break
             filtered = sorted(expanded)[:top_k]
 
+        # Step 5️⃣ — Build final chunk list
         final_chunks = [chunks[i] for i in filtered]
-        print(f"✅ Retrieved {len(final_chunks)} chunks (semantic + neighbor fill)")
+        print(f"✅ Retrieved {len(final_chunks)} chunks (semantic + neighbor fill).")
         return final_chunks
 
     except Exception as e:
-        print(f"⚠️ Retrieval error: {e}")
+        print(f"⚠️ Retrieval error: {repr(e)}")
         return []
 
 # ==========================================================