feat(append_node): append node to existing graph

ScrapeGraphAI · Jun 4, 2024 · f8b08e0 · f8b08e0
1 parent 376f758
commit f8b08e0
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 1 deletion.
diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py
@@ -378,6 +378,16 @@ def get_state(self, key=None) -> dict:
             return self.final_state[key]
         return self.final_state
 
+    def append_node(self, node):
+        """
+        Add a node to the graph.
+
+        Args:
+            node (BaseNode): The node to add to the graph.
+        """
+
+        self.graph.append_node(node)
+
     def get_execution_info(self):
         """
         Returns the execution information of the graph.

diff --git a/scrapegraphai/graphs/base_graph.py b/scrapegraphai/graphs/base_graph.py
@@ -49,6 +49,7 @@ class BaseGraph:
     def __init__(self, nodes: list, edges: list, entry_point: str, use_burr: bool = False, burr_config: dict = None):
 
         self.nodes = nodes
+        self.raw_edges = edges
         self.edges = self._create_edges({e for e in edges})
         self.entry_point = entry_point.node_name
         self.initial_state = {}
@@ -168,4 +169,25 @@ def execute(self, initial_state: dict) -> Tuple[dict, list]:
             result = bridge.execute(initial_state)
             return (result["_state"], [])
         else:
-            return self._execute_standard(initial_state)
+            return self._execute_standard(initial_state)
+
+    def append_node(self, node):
+        """
+        Adds a node to the graph.
+
+        Args:
+            node (BaseNode): The node instance to add to the graph.
+        """
+
+        # if node name already exists in the graph, raise an exception
+        if node.node_name in {n.node_name for n in self.nodes}:
+            raise ValueError(f"Node with name '{node.node_name}' already exists in the graph. You can change it by setting the 'node_name' attribute.")
+
+        # get the last node in the list
+        last_node = self.nodes[-1]
+        # add the edge connecting the last node to the new node
+        self.raw_edges.append((last_node, node))
+        # add the node to the list of nodes
+        self.nodes.append(node)
+        # update the edges connecting the last node to the new node
+        self.edges = self._create_edges({e for e in self.raw_edges})