feat: add Parse_Node

ScrapeGraphAI · Jun 12, 2024 · e6c7940 · e6c7940
1 parent 79b8326
commit e6c7940
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 8 deletions.
diff --git a/scrapegraphai/graphs/pdf_scraper_graph.py b/scrapegraphai/graphs/pdf_scraper_graph.py
@@ -11,6 +11,7 @@
 
 from ..nodes import (
     FetchNode,
+    ParseNode,
     RAGNode,
     GenerateAnswerPDFNode
 )
@@ -66,6 +67,15 @@ def _create_graph(self) -> BaseGraph:
             output=["doc"],
         )
 
+        parse_node = ParseNode(
+            input="doc",
+            output=["parsed_doc"],
+            node_config={
+                "parse_html": False,
+                "chunk_size": self.model_token
+            }
+        )
+
         rag_node = RAGNode(
             input="user_prompt & (parsed_doc | doc)",
             output=["relevant_chunks"],
@@ -86,11 +96,13 @@ def _create_graph(self) -> BaseGraph:
         return BaseGraph(
             nodes=[
                 fetch_node,
+                parse_node,
                 rag_node,
                 generate_answer_node_pdf,
             ],
             edges=[
-                (fetch_node, rag_node),
+                (fetch_node, parse_node),
+                (parse_node, rag_node),
                 (rag_node, generate_answer_node_pdf)
             ],
             entry_point=fetch_node

diff --git a/scrapegraphai/graphs/smart_scraper_graph.py b/scrapegraphai/graphs/smart_scraper_graph.py
@@ -3,8 +3,8 @@
 """
 
 from typing import Optional
+import logging
 from pydantic import BaseModel
-
 from .base_graph import BaseGraph
 from .abstract_graph import AbstractGraph
 
@@ -70,6 +70,7 @@ def _create_graph(self) -> BaseGraph:
             }
         )
         logging.info("FetchNode configured with headless: %s", self.config.get("headless", True))
+
         parse_node = ParseNode(
             input="doc",
             output=["parsed_doc"],

diff --git a/scrapegraphai/nodes/parse_node.py b/scrapegraphai/nodes/parse_node.py
@@ -70,12 +70,20 @@ def execute(self, state: dict) -> dict:
         docs_transformed = input_data[0]
         if self.parse_html:
             docs_transformed = Html2TextTransformer().transform_documents(input_data[0])
-        docs_transformed = docs_transformed[0]
-
-        chunks = chunk(text=docs_transformed.page_content,
-                        chunk_size= self.node_config.get("chunk_size", 4096),
-                        token_counter=lambda x: len(x.split()),
-                        memoize=False)
+            docs_transformed = docs_transformed[0]
+
+            chunks = chunk(text=docs_transformed.page_content,
+                            chunk_size= self.node_config.get("chunk_size", 4096),
+                            token_counter=lambda x: len(x.split()),
+                            memoize=False)
+        else:
+            docs_transformed = docs_transformed[0]
+
+            chunks = chunk(text=docs_transformed,
+                            chunk_size= self.node_config.get("chunk_size", 4096),
+                            token_counter=lambda x: len(x.split()),
+                            memoize=False)
+
         state.update({self.output[0]: chunks})
 
         return state