From 7a13a6819ff35a6f6197ee837d0eb8ea65e31776 Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Tue, 4 Jun 2024 12:01:21 +0200
Subject: [PATCH] feat: refactoring of rag node

---
 .gitignore                      |  4 ++++
 scrapegraphai/nodes/rag_node.py | 16 ++++++++++++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index c1750078..aa84820c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -23,6 +23,7 @@ docs/source/_static/
 venv/
 .venv/
 .vscode/
+.conda/
 
 # exclude pdf, mp3
 *.pdf
@@ -38,3 +39,6 @@ lib/
 *.html
 .idea
 
+# extras
+cache/
+run_smart_scraper.py
diff --git a/scrapegraphai/nodes/rag_node.py b/scrapegraphai/nodes/rag_node.py
index 6d26bd1c..e9834693 100644
--- a/scrapegraphai/nodes/rag_node.py
+++ b/scrapegraphai/nodes/rag_node.py
@@ -3,6 +3,7 @@ RAGNode Module
 """
 
 from typing import List, Optional
+import os
 
 from langchain.docstore.document import Document
 from langchain.retrievers import ContextualCompressionRetriever
@@ -98,7 +99,18 @@ class RAGNode(BaseNode):
         )
         embeddings = self.embedder_model
 
-        retriever = FAISS.from_documents(chunked_docs, embeddings).as_retriever()
+        #------
+        index = FAISS.from_documents(chunked_docs, embeddings)
+        # Define the folder name
+        folder_name = "cache"
+        # Check if the folder exists, if not, create it
+        if not os.path.exists(folder_name):
+            os.makedirs(folder_name)
+        # Save the index to the folder
+        index.save_local(folder_name)
+
+        retriever = index.as_retriever()
+        #------
 
         redundant_filter = EmbeddingsRedundantFilter(embeddings=embeddings)
         # similarity_threshold could be set, now k=20
@@ -121,4 +133,4 @@ class RAGNode(BaseNode):
         self.logger.info("--- (tokens compressed and vector stored) ---")
 
         state.update({self.output[0]: compressed_docs})
-        return state
+        return state
\ No newline at end of file