From c7870905e10da85b81761ab2c3f71220bafe9f22 Mon Sep 17 00:00:00 2001 From: Marco Perini Date: Tue, 18 Jun 2024 23:22:23 +0200 Subject: [PATCH] fix: total tokens and docs --- docs/source/scrapers/telemetry.rst | 1 + scrapegraphai/graphs/base_graph.py | 12 +++++++++++- scrapegraphai/telemetry/telemetry.py | 3 ++- 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/docs/source/scrapers/telemetry.rst b/docs/source/scrapers/telemetry.rst index a6598092..f5e9f27c 100644 --- a/docs/source/scrapers/telemetry.rst +++ b/docs/source/scrapers/telemetry.rst @@ -29,6 +29,7 @@ Additionally, the following properties are collected: "source_type": source_type, "execution_time": execution_time, "error_node": error_node_name, + "total_tokens": total_tokens, } For more details, refer to the `telemetry.py `_ module. diff --git a/scrapegraphai/graphs/base_graph.py b/scrapegraphai/graphs/base_graph.py index 90585e6a..33c259bb 100644 --- a/scrapegraphai/graphs/base_graph.py +++ b/scrapegraphai/graphs/base_graph.py @@ -140,6 +140,16 @@ class BaseGraph: result = current_node.execute(state) except Exception as e: 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 node_exec_time = time.time() - curr_time total_exec_time += node_exec_time @@ -187,7 +197,7 @@ class BaseGraph: embedder_model=embedder_model, source_type=source_type, 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 diff --git a/scrapegraphai/telemetry/telemetry.py b/scrapegraphai/telemetry/telemetry.py index 73e7c9cb..20d38186 100644 --- a/scrapegraphai/telemetry/telemetry.py +++ b/scrapegraphai/telemetry/telemetry.py @@ -156,7 +156,7 @@ def log_event(event: str, properties: Dict[str, any]): 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 = { "graph_name": graph_name, "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, "execution_time": execution_time, "error_node": error_node, + "total_tokens": total_tokens, } log_event("graph_execution", properties)