Update rag.py

rag.py

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