Fix the token count at the iteration node (#11235)

Co-authored-by: -LAN- <laipz8200@outlook.com>
2024-12-09 16:02:04 +09:00
parent 5c166b3f40
commit c3c6a48059
2 changed files with 5 additions and 3 deletions
--- a/api/core/workflow/nodes/iteration/iteration_node.py
+++ b/api/core/workflow/nodes/iteration/iteration_node.py
@@ -242,7 +242,10 @@ class IterationNode(BaseNode[IterationNodeData]):
                run_result=NodeRunResult(
                    status=WorkflowNodeExecutionStatus.SUCCEEDED,
                    outputs={"output": 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: