fix: total tokens and docs

This commit is contained in:
Marco Perini 2024-06-18 23:22:23 +02:00
parent ef8be1c115
commit c7870905e1
3 changed files with 14 additions and 2 deletions

View File

@ -29,6 +29,7 @@ Additionally, the following properties are collected:
"source_type": source_type, "source_type": source_type,
"execution_time": execution_time, "execution_time": execution_time,
"error_node": error_node_name, "error_node": error_node_name,
"total_tokens": total_tokens,
} }
For more details, refer to the `telemetry.py <https://github.com/VinciGit00/Scrapegraph-ai/blob/main/scrapegraphai/telemetry/telemetry.py>`_ module. For more details, refer to the `telemetry.py <https://github.com/VinciGit00/Scrapegraph-ai/blob/main/scrapegraphai/telemetry/telemetry.py>`_ module.

View File

@ -140,6 +140,16 @@ class BaseGraph:
result = current_node.execute(state) result = current_node.execute(state)
except Exception as e: except Exception as e:
error_node = current_node.node_name error_node = current_node.node_name
graph_execution_time = time.time() - start_time
log_graph_execution(
graph_name=self.graph_name,
llm_model=llm_model,
embedder_model=embedder_model,
source_type=source_type,
execution_time=graph_execution_time,
error_node=error_node
)
raise e raise e
node_exec_time = time.time() - curr_time node_exec_time = time.time() - curr_time
total_exec_time += node_exec_time total_exec_time += node_exec_time
@ -187,7 +197,7 @@ class BaseGraph:
embedder_model=embedder_model, embedder_model=embedder_model,
source_type=source_type, source_type=source_type,
execution_time=graph_execution_time, execution_time=graph_execution_time,
error_node=error_node total_tokens=cb_total["total_tokens"] if cb_total["total_tokens"] > 0 else None,
) )
return state, exec_info return state, exec_info

View File

@ -156,7 +156,7 @@ def log_event(event: str, properties: Dict[str, any]):
send_event_json(event_json) send_event_json(event_json)
def log_graph_execution(graph_name: str, llm_model: str, embedder_model: str, source_type: str, execution_time: float, error_node: str = None): def log_graph_execution(graph_name: str, llm_model: str, embedder_model: str, source_type: str, execution_time: float, error_node: str = None, total_tokens: int = None):
properties = { properties = {
"graph_name": graph_name, "graph_name": graph_name,
"llm_model": llm_model, "llm_model": llm_model,
@ -164,6 +164,7 @@ def log_graph_execution(graph_name: str, llm_model: str, embedder_model: str, so
"source_type": source_type, "source_type": source_type,
"execution_time": execution_time, "execution_time": execution_time,
"error_node": error_node, "error_node": error_node,
"total_tokens": total_tokens,
} }
log_event("graph_execution", properties) log_event("graph_execution", properties)