feat: removed rag node

VinciGit00 · May 31, 2024 · 930f673 · 930f673
1 parent 6d1d91a
commit 930f673
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 18 deletions.
diff --git a/scrapegraphai/graphs/pdf_scraper_graph.py b/scrapegraphai/graphs/pdf_scraper_graph.py
@@ -1,3 +1,4 @@
+
 """
 PDFScraperGraph Module
 """
@@ -9,7 +10,6 @@
 
 from ..nodes import (
     FetchNode,
-    RAGNode,
     GenerateAnswerPDFNode
 )
 
@@ -63,14 +63,7 @@ def _create_graph(self) -> BaseGraph:
             input='pdf | pdf_dir',
             output=["doc"],
         )
-        rag_node = RAGNode(
-            input="user_prompt & doc",
-            output=["relevant_chunks"],
-            node_config={
-                "llm_model": self.llm_model,
-                "embedder_model": self.embedder_model
-            }
-        )
+
         generate_answer_node_pdf = GenerateAnswerPDFNode(
             input="user_prompt & (relevant_chunks | doc)",
             output=["answer"],
@@ -83,12 +76,10 @@ def _create_graph(self) -> BaseGraph:
         return BaseGraph(
             nodes=[
                 fetch_node,
-                rag_node,
                 generate_answer_node_pdf,
             ],
             edges=[
-                (fetch_node, rag_node),
-                (rag_node, generate_answer_node_pdf)
+                (fetch_node, generate_answer_node_pdf)
             ],
             entry_point=fetch_node
         )
@@ -104,4 +95,4 @@ def run(self) -> str:
         inputs = {"user_prompt": self.prompt, self.input_key: self.source}
         self.final_state, self.execution_info = self.graph.execute(inputs)
 
-        return self.final_state.get("answer", "No answer found.")
+        return self.final_state.get("answer", "No answer found.")
diff --git a/scrapegraphai/graphs/smart_scraper_graph.py b/scrapegraphai/graphs/smart_scraper_graph.py
@@ -117,4 +117,4 @@ def run(self) -> str:
         inputs = {"user_prompt": self.prompt, self.input_key: self.source}
         self.final_state, self.execution_info = self.graph.execute(inputs)
 
-        return self.final_state.get("answer", "No answer found.")
+        return self.final_state.get("answer", "No answer found.")
diff --git a/scrapegraphai/nodes/generate_answer_pdf_node.py b/scrapegraphai/nodes/generate_answer_pdf_node.py
@@ -95,9 +95,7 @@ def execute(self, state):
         output_parser = JsonOutputParser()
         format_instructions = output_parser.get_format_instructions()
 
-
         chains_dict = {}
-
         # Use tqdm to add progress bar
         for i, chunk in enumerate(
             tqdm(doc, desc="Processing chunks", disable=not self.verbose)
@@ -107,7 +105,7 @@ def execute(self, state):
                     template=template_no_chunks_pdf,
                     input_variables=["question"],
                     partial_variables={
-                        "context": chunk.page_content,
+                        "context":chunk,
                         "format_instructions": format_instructions,
                     },
                 )
@@ -116,7 +114,7 @@ def execute(self, state):
                     template=template_chunks_pdf,
                     input_variables=["question"],
                     partial_variables={
-                        "context": chunk.page_content,
+                        "context":chunk,
                         "chunk_id": i + 1,
                         "format_instructions": format_instructions,
                     },