From 30b87b72b00aa62a502944e49ecc8e304551b696 Mon Sep 17 00:00:00 2001 From: s2terminal Date: Sun, 1 Dec 2024 13:44:45 +0900 Subject: [PATCH 1/2] Fix the token count at the iteration node --- api/core/workflow/nodes/iteration/iteration_node.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/api/core/workflow/nodes/iteration/iteration_node.py b/api/core/workflow/nodes/iteration/iteration_node.py index 22f242a42fb56f..cba4f8ac025f62 100644 --- a/api/core/workflow/nodes/iteration/iteration_node.py +++ b/api/core/workflow/nodes/iteration/iteration_node.py @@ -235,7 +235,11 @@ def _run(self) -> Generator[NodeEvent | InNodeEvent, None, None]: run_result=NodeRunResult( status=WorkflowNodeExecutionStatus.SUCCEEDED, outputs={"output": jsonable_encoder(outputs)}, - metadata={NodeRunMetadataKey.ITERATION_DURATION_MAP: iter_run_map}, + metadata={ + NodeRunMetadataKey.ITERATION_DURATION_MAP: iter_run_map, + NodeRunMetadataKey.TOTAL_TOKENS: graph_engine.graph_runtime_state.total_tokens, + }, + ) ) except IterationNodeError as e: From 1a9339739fb2aa2288b6ce460911ca354f3a2389 Mon Sep 17 00:00:00 2001 From: s2terminal Date: Tue, 3 Dec 2024 18:11:33 +0900 Subject: [PATCH 2/2] fix lint --- api/core/workflow/nodes/iteration/iteration_node.py | 1 - 1 file changed, 1 deletion(-) diff --git a/api/core/workflow/nodes/iteration/iteration_node.py b/api/core/workflow/nodes/iteration/iteration_node.py index cba4f8ac025f62..a4f8b6f4ae77b5 100644 --- a/api/core/workflow/nodes/iteration/iteration_node.py +++ b/api/core/workflow/nodes/iteration/iteration_node.py @@ -239,7 +239,6 @@ def _run(self) -> Generator[NodeEvent | InNodeEvent, None, None]: NodeRunMetadataKey.ITERATION_DURATION_MAP: iter_run_map, NodeRunMetadataKey.TOTAL_TOKENS: graph_engine.graph_runtime_state.total_tokens, }, - ) ) except IterationNodeError as e: