Skip to content

Commit c787090

Browse files
committed
fix: total tokens and docs
1 parent ef8be1c commit c787090

File tree

3 files changed

+14
-2
lines changed

3 files changed

+14
-2
lines changed

docs/source/scrapers/telemetry.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ Additionally, the following properties are collected:
2929
"source_type": source_type,
3030
"execution_time": execution_time,
3131
"error_node": error_node_name,
32+
"total_tokens": total_tokens,
3233
}
3334
3435
For more details, refer to the `telemetry.py <https://github.com/VinciGit00/Scrapegraph-ai/blob/main/scrapegraphai/telemetry/telemetry.py>`_ module.

scrapegraphai/graphs/base_graph.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,16 @@ def _execute_standard(self, initial_state: dict) -> Tuple[dict, list]:
140140
result = current_node.execute(state)
141141
except Exception as e:
142142
error_node = current_node.node_name
143+
144+
graph_execution_time = time.time() - start_time
145+
log_graph_execution(
146+
graph_name=self.graph_name,
147+
llm_model=llm_model,
148+
embedder_model=embedder_model,
149+
source_type=source_type,
150+
execution_time=graph_execution_time,
151+
error_node=error_node
152+
)
143153
raise e
144154
node_exec_time = time.time() - curr_time
145155
total_exec_time += node_exec_time
@@ -187,7 +197,7 @@ def _execute_standard(self, initial_state: dict) -> Tuple[dict, list]:
187197
embedder_model=embedder_model,
188198
source_type=source_type,
189199
execution_time=graph_execution_time,
190-
error_node=error_node
200+
total_tokens=cb_total["total_tokens"] if cb_total["total_tokens"] > 0 else None,
191201
)
192202

193203
return state, exec_info

scrapegraphai/telemetry/telemetry.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,14 +156,15 @@ def log_event(event: str, properties: Dict[str, any]):
156156
send_event_json(event_json)
157157

158158

159-
def log_graph_execution(graph_name: str, llm_model: str, embedder_model: str, source_type: str, execution_time: float, error_node: str = None):
159+
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):
160160
properties = {
161161
"graph_name": graph_name,
162162
"llm_model": llm_model,
163163
"embedder_model": embedder_model,
164164
"source_type": source_type,
165165
"execution_time": execution_time,
166166
"error_node": error_node,
167+
"total_tokens": total_tokens,
167168
}
168169
log_event("graph_execution", properties)
169170

0 commit comments

Comments
 (0)