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