Scrapegraph-ai/scrapegraphai/utils/prettify_exec_info.py
2025-01-06 15:10:35 +01:00

52 lines
1.5 KiB
Python

"""
Prettify the execution information of the graph.
"""
from typing import Union
def prettify_exec_info(
complete_result: list[dict], as_string: bool = True
) -> Union[str, list[dict]]:
"""
Formats the execution information of a graph showing node statistics.
Args:
complete_result (list[dict]): The execution information containing node statistics.
as_string (bool, optional): If True, returns a formatted string table.
If False, returns the original list. Defaults to True.
Returns:
Union[str, list[dict]]: A formatted string table if as_string=True,
otherwise the original list of dictionaries.
"""
if not as_string:
return complete_result
if not complete_result:
return "Empty result"
# Format the table
lines = []
lines.append("Node Statistics:")
lines.append("-" * 100)
lines.append(
f"{'Node':<20} {'Tokens':<10} {'Prompt':<10} {'Compl.':<10} {'Requests':<10} {'Cost ($)':<10} {'Time (s)':<10}"
)
lines.append("-" * 100)
for item in complete_result:
node = item["node_name"]
tokens = item["total_tokens"]
prompt = item["prompt_tokens"]
completion = item["completion_tokens"]
requests = item["successful_requests"]
cost = f"{item['total_cost_USD']:.4f}"
time = f"{item['exec_time']:.2f}"
lines.append(
f"{node:<20} {tokens:<10} {prompt:<10} {completion:<10} {requests:<10} {cost:<10} {time:<10}"
)
return "\n".join(lines)