Fix the token count at the iteration node

langgenius · Dec 1, 2024 · 30b87b7 · 30b87b7
1 parent ef20481
commit 30b87b7
Showing 1 changed file with 5 additions and 1 deletion.
diff --git 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: