diff --git a/scrapegraphai/graphs/base_graph.py b/scrapegraphai/graphs/base_graph.py index c2ebfb0b..47e6a158 100644 --- a/scrapegraphai/graphs/base_graph.py +++ b/scrapegraphai/graphs/base_graph.py @@ -56,7 +56,7 @@ class BaseGraph: 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 @@ class BaseGraph: 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,17 +93,18 @@ class BaseGraph: 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_time": node_exec_time, } - 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"] @@ -119,10 +119,14 @@ class BaseGraph: else: current_node_name = None - execution_info = { - "total_exec_time": total_exec_time, - "total_model_info": cb_total, - "nodes_info": exec_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, execution_info + return state, exec_info diff --git a/scrapegraphai/utils/prettify_exec_info.py b/scrapegraphai/utils/prettify_exec_info.py index 7023d6df..21004b71 100644 --- a/scrapegraphai/utils/prettify_exec_info.py +++ b/scrapegraphai/utils/prettify_exec_info.py @@ -5,44 +5,17 @@ Prettify the execution information of the graph. 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