Update src/search.py

src/search.py

@@ -0,0 +1,85 @@
+import faiss
+import numpy as np
+import pickle
+from sentence_transformers import SentenceTransformer, CrossEncoder
+
+INDEX_FILE = "navy_index.faiss"
+META_FILE = "navy_metadata.pkl" # We still use this for fast mapping, or we could query SQL.
+
+class SearchEngine:
+    def __init__(self):
+        # Force CPU
+        self.bi_encoder = SentenceTransformer('all-MiniLM-L6-v2', device="cpu")
+        self.cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device="cpu")
+        self.index = None
+        self.metadata = [] # List of dicts: {'doc_id':..., 'text':...}
+        self.load_index()
+
+    def load_index(self):
+        if os.path.exists(INDEX_FILE) and os.path.exists(META_FILE):
+            try:
+                self.index = faiss.read_index(INDEX_FILE)
+                with open(META_FILE, "rb") as f: self.metadata = pickle.load(f)
+            except:
+                self.reset_index()
+        else:
+            self.reset_index()
+
+    def reset_index(self):
+        self.index = faiss.IndexIDMap(faiss.IndexFlatIP(384))
+        self.metadata = []
+
+    def add_features(self, chunks):
+        """
+        Embeds chunks and adds to FAISS.
+        chunks = [{'text':..., 'doc_id':...}]
+        """
+        texts = [c["text"] for c in chunks]
+        embeddings = self.bi_encoder.encode(texts)
+        faiss.normalize_L2(embeddings)
+        
+        start_id = len(self.metadata)
+        ids = np.arange(start_id, start_id + len(chunks)).astype('int64')
+        
+        self.index.add_with_ids(embeddings, ids)
+        self.metadata.extend(chunks)
+        self.save()
+
+    def save(self):
+        faiss.write_index(self.index, INDEX_FILE)
+        with open(META_FILE, "wb") as f: pickle.dump(self.metadata, f)
+
+    def search(self, query, top_k=5):
+        if not self.index or self.index.ntotal == 0: return []
+        
+        q_vec = self.bi_encoder.encode([query])
+        faiss.normalize_L2(q_vec)
+        
+        # 1. Retrieve Candidate Vectors
+        scores, indices = self.index.search(q_vec, min(self.index.ntotal, top_k * 10))
+        
+        candidates = []
+        for i, idx in enumerate(indices[0]):
+            if idx != -1:
+                item = self.metadata[idx]
+                candidates.append([query, item['text']])
+        
+        if not candidates: return []
+
+        # 2. Re-Rank with Cross Encoder
+        cross_scores = self.cross_encoder.predict(candidates)
+        
+        # 3. Format Results
+        results = []
+        for i, idx in enumerate(indices[0]):
+            if idx != -1:
+                meta = self.metadata[idx]
+                results.append({
+                    "score": cross_scores[i],
+                    "doc_id": meta["doc_id"],
+                    "source": meta["source"],
+                    "snippet": meta["text"]
+                })
+        
+        # Sort by Cross-Encoder Score
+        return sorted(results, key=lambda x: x['score'], reverse=True)[:top_k]