Update rag.py

rag.py

@@ -2,27 +2,44 @@ 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.index:
+        if not self.has_data():
             return []
+            
         query_embed = embedder.encode([query])
-        D, I = self.index.search(np.array(query_embed), k=top_k)
-        return [self.texts[i] for i in I[0]]
+        _, 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