Update rag.py

rag.py

@@ -1,45 +1,29 @@
+# rag.py
 from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 
-# Load model only once
+# load model only once
 embedder = SentenceTransformer('all-MiniLM-L6-v2')
-DIMENSION = 384  # Fixed dimension for all-MiniLM-L6-v2
 
 class VectorStore:
     def __init__(self):
         self.texts = []
+        self.embeddings = []
         self.index = None
-        self.embeddings = None
 
     def add_texts(self, texts):
         """Add list of texts to the store."""
-        if not texts:
-            return
-            
         new_embeds = embedder.encode(texts)
-        
-        # Initialize index if needed
-        if self.index is None:
-            self.index = faiss.IndexFlatL2(DIMENSION)
-            self.embeddings = new_embeds
-        else:
-            self.embeddings = np.vstack([self.embeddings, new_embeds])
-        
-        # Rebuild index with all embeddings
-        self.index.reset()
-        self.index.add(self.embeddings.astype('float32'))
         self.texts.extend(texts)
+        self.embeddings.extend(new_embeds)
+        self.index = faiss.IndexFlatL2(new_embeds.shape[1])
+        self.index.add(np.array(self.embeddings))
 
     def retrieve(self, query, top_k=3):
         """Return top-k relevant texts for the query."""
-        if not self.has_data():
+        if not self.index:
             return []
-            
         query_embed = embedder.encode([query])
-        _, I = self.index.search(query_embed.astype('float32'), top_k)
-        return [self.texts[i] for i in I[0] if i < len(self.texts)]
-
-    def has_data(self):
-        """Check if we have any data stored"""
-        return self.index is not None and self.index.ntotal > 0
+        D, I = self.index.search(np.array(query_embed), k=top_k)
+        return [self.texts[i] for i in I[0]]