refactoring of the test engine

examples/openai/custom_graph_openai.py

@@ -62,19 +62,19 @@
 # ************************************************
 
 graph = BaseGraph(
-    nodes={
+    nodes=[
         robot_node,
         fetch_node,
         parse_node,
         rag_node,
         generate_answer_node,
-    },
-    edges={
+    ],
+    edges=[
         (robot_node, fetch_node),
         (fetch_node, parse_node),
         (parse_node, rag_node),
         (rag_node, generate_answer_node)
-    },
+    ],
     entry_point=robot_node
 )
 

manual deployment/commit_and_push_with_tests.sh

@@ -13,6 +13,8 @@ pylint pylint scrapegraphai/**/*.py scrapegraphai/*.py tests/**/*.py
 
 cd tests
 
+poetry install
+
 # Run pytest
 if ! pytest; then
     echo "Pytest failed. Aborting commit and push."

scrapegraphai/graphs/abstract_graph.py

@@ -41,7 +41,7 @@ def _create_llm(self, llm_config: dict):
             try:
                 self.model_token = models_tokens["openai"][llm_params["model"]]
             except KeyError:
-                raise ValueError("Model not supported")
+                raise KeyError("Model not supported")
             return OpenAI(llm_params)
 
         elif "azure" in llm_params["model"]:
@@ -50,14 +50,14 @@ def _create_llm(self, llm_config: dict):
             try:
                 self.model_token = models_tokens["azure"][llm_params["model"]]
             except KeyError:
-                raise ValueError("Model not supported")
+                raise KeyError("Model not supported")
             return AzureOpenAI(llm_params)
 
         elif "gemini" in llm_params["model"]:
             try:
                 self.model_token = models_tokens["gemini"][llm_params["model"]]
             except KeyError:
-                raise ValueError("Model not supported")
+                raise KeyError("Model not supported")
             return Gemini(llm_params)
 
         elif "ollama" in llm_params["model"]:
@@ -70,19 +70,27 @@ def _create_llm(self, llm_config: dict):
                 try:
                     self.model_token = models_tokens["ollama"][llm_params["model"]]
                 except KeyError:
-                    raise ValueError("Model not supported")
+                    raise KeyError("Model not supported")
 
             return Ollama(llm_params)
         elif "hugging_face" in llm_params["model"]:
             try:
                 self.model_token = models_tokens["hugging_face"][llm_params["model"]]
             except KeyError:
-                raise ValueError("Model not supported")
+                raise KeyError("Model not supported")
             return HuggingFace(llm_params)
         else:
             raise ValueError(
                 "Model provided by the configuration not supported")
 
+    def get_state(self, key=None) -> dict:
+        """""
+        Obtain the current state
+        """
+        if key is not None:
+            return self.final_state[key]
+        return self.final_state
+
     def get_execution_info(self):
         """
         Returns the execution information of the graph.

scrapegraphai/graphs/base_graph.py

@@ -2,6 +2,7 @@
 Module for creating the base graphs
  """
 import time
+import warnings
 from langchain_community.callbacks import get_openai_callback
 
 
@@ -10,31 +11,37 @@ class BaseGraph:
     BaseGraph manages the execution flow of a graph composed of interconnected nodes.
 
     Attributes:
-        nodes (dict): A dictionary mapping each node's name to its corresponding node instance.
-        edges (dict): A dictionary representing the directed edges of the graph where each 
+        nodes (list): A dictionary mapping each node's name to its corresponding node instance.
+        edges (list): A dictionary representing the directed edges of the graph where each
                       key-value pair corresponds to the from-node and to-node relationship.
         entry_point (str): The name of the entry point node from which the graph execution begins.
 
     Methods:
-        execute(initial_state): Executes the graph's nodes starting from the entry point and 
+        execute(initial_state): Executes the graph's nodes starting from the entry point and
                                 traverses the graph based on the provided initial state.
 
     Args:
         nodes (iterable): An iterable of node instances that will be part of the graph.
-        edges (iterable): An iterable of tuples where each tuple represents a directed edge 
+        edges (iterable): An iterable of tuples where each tuple represents a directed edge
                           in the graph, defined by a pair of nodes (from_node, to_node).
         entry_point (BaseNode): The node instance that represents the entry point of the graph.
     """
 
-    def __init__(self, nodes: dict, edges: dict, entry_point: str):
+    def __init__(self, nodes: list, edges: list, entry_point: str):
         """
         Initializes the graph with nodes, edges, and the entry point.
         """
-        self.nodes = {node.node_name: node for node in nodes}
-        self.edges = self._create_edges(edges)
+
+        self.nodes = nodes
+        self.edges = self._create_edges({e for e in edges})
         self.entry_point = entry_point.node_name
 
-    def _create_edges(self, edges: dict) -> dict:
+        if nodes[0].node_name != entry_point.node_name:
+            # raise a warning if the entry point is not the first node in the list
+            warnings.warn(
+                "Careful! The entry point node is different from the first node if the graph.")
+
+    def _create_edges(self, edges: list) -> dict:
         """
         Helper method to create a dictionary of edges from the given iterable of tuples.
 
@@ -51,8 +58,8 @@ def _create_edges(self, edges: dict) -> dict:
 
     def execute(self, initial_state: dict) -> dict:
         """
-        Executes the graph by traversing nodes starting from the entry point. The execution 
-        follows the edges based on the result of each node's execution and continues until 
+        Executes the graph by traversing nodes starting from the entry point. The execution
+        follows the edges based on the result of each node's execution and continues until
         it reaches a node with no outgoing edges.
 
         Args:
@@ -61,7 +68,8 @@ def execute(self, initial_state: dict) -> dict:
         Returns:
             dict: The state after execution has completed, which may have been altered by the nodes.
         """
-        current_node_name = self.entry_point
+        print(self.nodes)
+        current_node_name = self.nodes[0]
         state = initial_state
 
         # variables for tracking execution info
@@ -75,10 +83,10 @@ def execute(self, initial_state: dict) -> dict:
             "total_cost_USD": 0.0,
         }
 
-        while current_node_name is not None:
+        for index in self.nodes:
 
             curr_time = time.time()
-            current_node = self.nodes[current_node_name]
+            current_node = index
 
             with get_openai_callback() as cb:
                 result = current_node.execute(state)

scrapegraphai/graphs/script_creator_graph.py

@@ -1,4 +1,4 @@
-""" 
+"""
 Module for creating the smart scraper
 """
 from .base_graph import BaseGraph
@@ -57,17 +57,17 @@ def _create_graph(self):
         )
 
         return BaseGraph(
-            nodes={
+            nodes=[
                 fetch_node,
                 parse_node,
                 rag_node,
                 generate_scraper_node,
-            },
-            edges={
+            ],
+            edges=[
                 (fetch_node, parse_node),
                 (parse_node, rag_node),
                 (rag_node, generate_scraper_node)
-            },
+            ],
             entry_point=fetch_node
         )
 

scrapegraphai/graphs/search_graph.py

@@ -11,6 +11,7 @@
 )
 from .abstract_graph import AbstractGraph
 
+
 class SearchGraph(AbstractGraph):
     """ 
     Module for searching info on the internet
@@ -49,19 +50,19 @@ def _create_graph(self):
         )
 
         return BaseGraph(
-            nodes={
+            nodes=[
                 search_internet_node,
                 fetch_node,
                 parse_node,
                 rag_node,
                 generate_answer_node,
-            },
-            edges={
+            ],
+            edges=[
                 (search_internet_node, fetch_node),
                 (fetch_node, parse_node),
                 (parse_node, rag_node),
                 (rag_node, generate_answer_node)
-            },
+            ],
             entry_point=search_internet_node
         )
 

scrapegraphai/graphs/smart_scraper_graph.py

@@ -1,4 +1,4 @@
-""" 
+"""
 Module for creating the smart scraper
 """
 from .base_graph import BaseGraph
@@ -10,6 +10,7 @@
 )
 from .abstract_graph import AbstractGraph
 
+
 class SmartScraperGraph(AbstractGraph):
     """
     SmartScraper is a comprehensive web scraping tool that automates the process of extracting
@@ -52,25 +53,25 @@ def _create_graph(self):
         )
 
         return BaseGraph(
-            nodes={
+            nodes=[
                 fetch_node,
                 parse_node,
                 rag_node,
                 generate_answer_node,
-            },
-            edges={
+            ],
+            edges=[
                 (fetch_node, parse_node),
                 (parse_node, rag_node),
                 (rag_node, generate_answer_node)
-            },
+            ],
             entry_point=fetch_node
         )
 
     def run(self) -> str:
         """
         Executes the web scraping process and returns the answer to the prompt.
         """
-        inputs = {"user_prompt": self.prompt, self.input_key: self.source}  
+        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.")

scrapegraphai/graphs/speech_graph.py

@@ -62,19 +62,19 @@ def _create_graph(self):
         )
 
         return BaseGraph(
-            nodes={
+            nodes=[
                 fetch_node,
                 parse_node,
                 rag_node,
                 generate_answer_node,
                 text_to_speech_node
-            },
-            edges={
+            ],
+            edges=[
                 (fetch_node, parse_node),
                 (parse_node, rag_node),
                 (rag_node, generate_answer_node),
                 (generate_answer_node, text_to_speech_node)
-            },
+            ],
             entry_point=fetch_node
         )