fix(fetch_node): bug in handling local files

PeriniM · PeriniM · commit a6e1813ddd36 · 2024-05-14T16:51:10.000+02:00
diff --git a/examples/openai/custom_graph_openai copy.py b/examples/openai/custom_graph_openai copy.py
diff --git a/examples/openai/omni_scraper_openai.py b/examples/openai/omni_scraper_openai.py
@@ -19,7 +19,7 @@
 graph_config = {
     "llm": {
         "api_key": openai_key,
-        "model": "gpt-4o",
+        "model": "gpt-4-turbo",
     },
     "verbose": True,
     "headless": True,
diff --git a/scrapegraphai/graphs/csv_scraper_graph.py b/scrapegraphai/graphs/csv_scraper_graph.py
@@ -30,8 +30,8 @@ def _create_graph(self):
         Creates the graph of nodes representing the workflow for web scraping.
         """
         fetch_node = FetchNode(
-            input="csv",
-            output=["doc"],
+            input="csv | csv_dir",
+            output=["doc", "link_urls", "img_urls"],
         )
         parse_node = ParseNode(
             input="doc",
diff --git a/scrapegraphai/graphs/deep_scraper_graph.py b/scrapegraphai/graphs/deep_scraper_graph.py
@@ -61,7 +61,7 @@ def _create_graph(self) -> BaseGraph:
         """
         fetch_node = FetchNode(
             input="url | local_dir",
-            output=["doc"]
+            output=["doc", "link_urls", "img_urls"]
         )
         parse_node = ParseNode(
             input="doc",
diff --git a/scrapegraphai/graphs/json_scraper_graph.py b/scrapegraphai/graphs/json_scraper_graph.py
@@ -54,8 +54,8 @@ def _create_graph(self) -> BaseGraph:
         """
 
         fetch_node = FetchNode(
-            input="json",
-            output=["doc"],
+            input="json | json_dir",
+            output=["doc", "link_urls", "img_urls"],
         )
         parse_node = ParseNode(
             input="doc",
diff --git a/scrapegraphai/graphs/pdf_scraper_graph.py b/scrapegraphai/graphs/pdf_scraper_graph.py
@@ -56,8 +56,8 @@ def _create_graph(self) -> BaseGraph:
         """
 
         fetch_node = FetchNode(
-            input='pdf',
-            output=["doc"],
+            input='pdf | pdf_dir',
+            output=["doc", "link_urls", "img_urls"],
         )
         parse_node = ParseNode(
             input="doc",
diff --git a/scrapegraphai/graphs/script_creator_graph.py b/scrapegraphai/graphs/script_creator_graph.py
@@ -59,7 +59,7 @@ def _create_graph(self) -> BaseGraph:
 
         fetch_node = FetchNode(
             input="url | local_dir",
-            output=["doc"],
+            output=["doc", "link_urls", "img_urls"],
         )
         parse_node = ParseNode(
             input="doc",
diff --git a/scrapegraphai/graphs/smart_scraper_graph.py b/scrapegraphai/graphs/smart_scraper_graph.py
@@ -57,7 +57,7 @@ def _create_graph(self) -> BaseGraph:
         """
         fetch_node = FetchNode(
             input="url | local_dir",
-            output=["doc"],
+            output=["doc", "link_urls", "img_urls"],
             node_config={
                 "loader_kwargs": self.config.get("loader_kwargs", {}),
             }
diff --git a/scrapegraphai/graphs/speech_graph.py b/scrapegraphai/graphs/speech_graph.py
@@ -56,7 +56,7 @@ def _create_graph(self) -> BaseGraph:
 
         fetch_node = FetchNode(
             input="url | local_dir",
-            output=["doc"]
+            output=["doc", "link_urls", "img_urls"]
         )
         parse_node = ParseNode(
             input="doc",
diff --git a/scrapegraphai/graphs/xml_scraper_graph.py b/scrapegraphai/graphs/xml_scraper_graph.py
@@ -56,8 +56,8 @@ def _create_graph(self) -> BaseGraph:
         """
 
         fetch_node = FetchNode(
-            input="xml",
-            output=["doc"]
+            input="xml | xml_dir",
+            output=["doc", "link_urls", "img_urls"]
         )
         parse_node = ParseNode(
             input="doc",
diff --git a/scrapegraphai/nodes/fetch_node.py b/scrapegraphai/nodes/fetch_node.py
@@ -83,37 +83,49 @@ def execute(self, state):
 
         source = input_data[0]
         if (
-            self.input == "json_dir"
-            or self.input == "xml_dir"
-            or self.input == "csv_dir"
+            input_keys[0] == "json_dir"
+            or input_keys[0] == "xml_dir"
+            or input_keys[0] == "csv_dir"
         ):
             compressed_document = [
                 Document(page_content=source, metadata={"source": "local_dir"})
             ]
-        # if it is a local directory
-
+            state.update({self.output[0]: compressed_document})
+            return state
+        
         # handling for pdf
-        elif self.input == "pdf":
+        elif input_keys[0] == "pdf":
             loader = PyPDFLoader(source)
             compressed_document = loader.load()
+            state.update({self.output[0]: compressed_document})
+            return state
 
-        elif self.input == "csv":
+        elif input_keys[0] == "csv":
             compressed_document = [
                 Document(
                     page_content=str(pd.read_csv(source)), metadata={"source": "csv"}
                 )
             ]
-        elif self.input == "json":
+            state.update({self.output[0]: compressed_document})
+            return state
+        
+        elif input_keys[0] == "json":
             f = open(source)
             compressed_document = [
                 Document(page_content=str(json.load(f)), metadata={"source": "json"})
             ]
-        elif self.input == "xml":
+            state.update({self.output[0]: compressed_document})
+            return state
+        
+        elif input_keys[0] == "xml":
             with open(source, "r", encoding="utf-8") as f:
                 data = f.read()
             compressed_document = [
                 Document(page_content=data, metadata={"source": "xml"})
             ]
+            state.update({self.output[0]: compressed_document})
+            return state
+        
         elif self.input == "pdf_dir":
             pass
 

Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ def _create_graph(self) -> BaseGraph:`
`61`	`61`	`"""`
`62`	`62`	`fetch_node = FetchNode(`
`63`	`63`	`input="url \| local_dir",`
`64`		`- output=["doc"]`
	`64`	`+ output=["doc", "link_urls", "img_urls"]`
`65`	`65`	`)`
`66`	`66`	`parse_node = ParseNode(`
`67`	`67`	`input="doc",`
Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,7 @@ def _create_graph(self) -> BaseGraph:`
`59`	`59`
`60`	`60`	`fetch_node = FetchNode(`
`61`	`61`	`input="url \| local_dir",`
`62`		`- output=["doc"],`
	`62`	`+ output=["doc", "link_urls", "img_urls"],`
`63`	`63`	`)`
`64`	`64`	`parse_node = ParseNode(`
`65`	`65`	`input="doc",`
Original file line number	Diff line number	Diff line change
`@@ -57,7 +57,7 @@ def _create_graph(self) -> BaseGraph:`
`57`	`57`	`"""`
`58`	`58`	`fetch_node = FetchNode(`
`59`	`59`	`input="url \| local_dir",`
`60`		`- output=["doc"],`
	`60`	`+ output=["doc", "link_urls", "img_urls"],`
`61`	`61`	`node_config={`
`62`	`62`	`"loader_kwargs": self.config.get("loader_kwargs", {}),`
`63`	`63`	`}`
Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@ def _create_graph(self) -> BaseGraph:`
`56`	`56`
`57`	`57`	`fetch_node = FetchNode(`
`58`	`58`	`input="url \| local_dir",`
`59`		`- output=["doc"]`
	`59`	`+ output=["doc", "link_urls", "img_urls"]`
`60`	`60`	`)`
`61`	`61`	`parse_node = ParseNode(`
`62`	`62`	`input="doc",`