fix: bug for calculate costs

Author: VinciGit00

scrapegraphai/graphs/base_graph.py

@@ -56,7 +56,7 @@ def _create_edges(self, edges: list) -> dict:
             edge_dict[from_node.node_name] = to_node.node_name
         return edge_dict
 
-    def execute(self, initial_state: dict) -> dict:
+    def execute(self, initial_state: dict) -> (dict, list):
         """
         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
@@ -68,13 +68,12 @@ def execute(self, initial_state: dict) -> dict:
         Returns:
             dict: The state after execution has completed, which may have been altered by the nodes.
         """
-        print(self.nodes)
         current_node_name = self.nodes[0]
         state = initial_state
 
         # variables for tracking execution info
         total_exec_time = 0.0
-        exec_info = {}
+        exec_info = []
         cb_total = {
             "total_tokens": 0,
             "prompt_tokens": 0,
@@ -94,18 +93,19 @@ def execute(self, initial_state: dict) -> dict:
                 total_exec_time += node_exec_time
 
                 cb = {
+                    "node_name": index.node_name,
                     "total_tokens": cb.total_tokens,
                     "prompt_tokens": cb.prompt_tokens,
                     "completion_tokens": cb.completion_tokens,
                     "successful_requests": cb.successful_requests,
                     "total_cost_USD": cb.total_cost,
-                }
-
-                exec_info[current_node_name] = {
                     "exec_time": node_exec_time,
-                    "model_info": cb
                 }
 
+                exec_info.append(
+                    cb
+                )
+
                 cb_total["total_tokens"] += cb["total_tokens"]
                 cb_total["prompt_tokens"] += cb["prompt_tokens"]
                 cb_total["completion_tokens"] += cb["completion_tokens"]
@@ -119,10 +119,14 @@ def execute(self, initial_state: dict) -> dict:
             else:
                 current_node_name = None
 
-        execution_info = {
-            "total_exec_time": total_exec_time,
-            "total_model_info": cb_total,
-            "nodes_info": exec_info
-        }
-
-        return state, execution_info
+        exec_info.append({
+            "node_name": "TOTAL RESULT",
+            "total_tokens":  cb_total["total_tokens"],
+            "prompt_tokens":  cb_total["prompt_tokens"],
+            "completion_tokens": cb_total["completion_tokens"],
+            "successful_requests": cb_total["successful_requests"],
+            "total_cost_USD":   cb_total["total_cost_USD"],
+            "exec_time": total_exec_time,
+        })
+
+        return state, exec_info

scrapegraphai/utils/prettify_exec_info.py

@@ -5,44 +5,17 @@
 import pandas as pd
 
 
-def prettify_exec_info(complete_result: dict) -> pd.DataFrame:
+def prettify_exec_info(complete_result: list[dict]) -> pd.DataFrame:
     """
     Transform the execution information of the graph into a DataFrame for better visualization.
 
     Args:
-    - complete_result (dict): The complete execution information of the graph.
+    - complete_result (list[dict]): The complete execution information of the graph.
 
     Returns:
     - pd.DataFrame: The execution information of the graph in a DataFrame.
     """
 
-    nodes_info = complete_result['nodes_info']
-    total_info = {
-        'total_exec_time': complete_result['total_exec_time'],
-        'total_model_info': complete_result['total_model_info']
-    }
+    df_nodes = pd.DataFrame(complete_result)
 
-    # Convert node-specific information to DataFrame
-    flat_data = []
-    for node_name, node_info in nodes_info.items():
-        flat_data.append({
-            'Node': node_name,
-            'Execution Time': node_info['exec_time'],
-            # Unpack the model_info dict into the row
-            **node_info['model_info']
-        })
-
-    df_nodes = pd.DataFrame(flat_data)
-
-    # Add a row for the total execution time and total model info
-    total_row = {
-        'Node': 'Total',
-        'Execution Time': total_info['total_exec_time'],
-        # Unpack the total_model_info dict into the row
-        **total_info['total_model_info']
-    }
-    df_total = pd.DataFrame([total_row])
-
-    # Combine the nodes DataFrame with the total info DataFrame
-    df_combined_with_total = pd.concat([df_nodes, df_total], ignore_index=True)
-    return df_combined_with_total
+    return df_nodes