Update rag.py

rag.py

@@ -1,35 +1,28 @@
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
+
+# load model only once
+embedder = SentenceTransformer('all-MiniLM-L6-v2')
 
 class VectorStore:
     def __init__(self):
-        self.vectorstore = None
-        self.embedder = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-        self.text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=500,
-            chunk_overlap=50
-        )
+        self.texts = []
+        self.embeddings = []
+        self.index = None
 
     def add_texts(self, texts):
-        if not texts:
-            return
-            
-        # Split and add texts
-        if self.vectorstore is None:
-            self.vectorstore = FAISS.from_texts(
-                self.text_splitter.split_text("\n\n".join(texts)),
-                self.embedder
-            )
-        else:
-            self.vectorstore.add_texts(
-                self.text_splitter.split_text("\n\n".join(texts))
-            )
+        """Add list of texts to the store."""
+        new_embeds = embedder.encode(texts)
+        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):
-        if self.vectorstore is None:
+        """Return top-k relevant texts for the query."""
+        if not self.index:
             return []
-        return [
-            doc.page_content 
-            for doc in self.vectorstore.similarity_search(query, k=top_k)
-        ]
+        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]]