CodeCow/Scrapegraph-ai
1
0
Fork 0
mirror of https://github.com/VinciGit00/Scrapegraph-ai.git synced 2026-06-23 21:00:30 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: bug for calculate costs

Browse Source
This commit is contained in:
VinciGit00 2024-04-28 14:36:49 +02:00
parent 9cd516507c
commit a9b11e433a
2 changed files with 21 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
scrapegraphai/graphs/base_graph.py
View File

@ -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

35
scrapegraphai/utils/prettify_exec_info.py
View File

@ -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