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

fix: total tokens and docs

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/source/scrapers/telemetry.rst
View File

@ -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 <https://github.com/VinciGit00/Scrapegraph-ai/blob/main/scrapegraphai/telemetry/telemetry.py>`_ module.

12
scrapegraphai/graphs/base_graph.py
View File

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

3
scrapegraphai/telemetry/telemetry.py
View File

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