Create rag.py

rag.py

@@ -0,0 +1,37 @@
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
+
+class RAG:
+    def __init__(self):
+        self.model = SentenceTransformer("all-MiniLM-L6-v2")
+        self.index = None
+        self.text_chunks = []
+
+    def chunk_text(self, text, chunk_size=500, overlap=50):
+        chunks = []
+        start = 0
+        while start < len(text):
+            end = start + chunk_size
+            chunks.append(text[start:end])
+            start += chunk_size - overlap
+        return chunks
+
+    def add_document(self, text):
+        chunks = self.chunk_text(text)
+        embeddings = self.model.encode(chunks)
+
+        if self.index is None:
+            self.index = faiss.IndexFlatL2(len(embeddings[0]))
+
+        self.index.add(np.array(embeddings))
+        self.text_chunks.extend(chunks)
+
+    def retrieve(self, query, k=3):
+        if self.index is None:
+            return []
+
+        query_embedding = self.model.encode([query])
+        distances, indices = self.index.search(np.array(query_embedding), k)
+
+        return [self.text_chunks[i] for i in indices[0]]