Fix(workflow): Prevent token overcount caused by loop/iteration (#28406)

This commit is contained in:
Jax
2025-11-25 09:56:59 +08:00
committed by GitHub
parent 6bd114285c
commit eed38c8b2a
2 changed files with 2 additions and 12 deletions

View File

@@ -237,8 +237,7 @@ class IterationNode(LLMUsageTrackingMixin, Node):
)
)
# Update the total tokens from this iteration
self.graph_runtime_state.total_tokens += graph_engine.graph_runtime_state.total_tokens
# Accumulate usage from this iteration
usage_accumulator[0] = self._merge_usage(
usage_accumulator[0], graph_engine.graph_runtime_state.llm_usage
)
@@ -265,7 +264,6 @@ class IterationNode(LLMUsageTrackingMixin, Node):
datetime,
list[GraphNodeEventBase],
object | None,
int,
dict[str, VariableUnion],
LLMUsage,
]
@@ -292,7 +290,6 @@ class IterationNode(LLMUsageTrackingMixin, Node):
iter_start_at,
events,
output_value,
tokens_used,
conversation_snapshot,
iteration_usage,
) = result
@@ -304,7 +301,6 @@ class IterationNode(LLMUsageTrackingMixin, Node):
yield from events
# Update tokens and timing
self.graph_runtime_state.total_tokens += tokens_used
iter_run_map[str(index)] = (datetime.now(UTC).replace(tzinfo=None) - iter_start_at).total_seconds()
usage_accumulator[0] = self._merge_usage(usage_accumulator[0], iteration_usage)
@@ -336,7 +332,7 @@ class IterationNode(LLMUsageTrackingMixin, Node):
item: object,
flask_app: Flask,
context_vars: contextvars.Context,
) -> tuple[datetime, list[GraphNodeEventBase], object | None, int, dict[str, VariableUnion], LLMUsage]:
) -> tuple[datetime, list[GraphNodeEventBase], object | None, dict[str, VariableUnion], LLMUsage]:
"""Execute a single iteration in parallel mode and return results."""
with preserve_flask_contexts(flask_app=flask_app, context_vars=context_vars):
iter_start_at = datetime.now(UTC).replace(tzinfo=None)
@@ -363,7 +359,6 @@ class IterationNode(LLMUsageTrackingMixin, Node):
iter_start_at,
events,
output_value,
graph_engine.graph_runtime_state.total_tokens,
conversation_snapshot,
graph_engine.graph_runtime_state.llm_usage,
)