From 1329f690722f7921efc26b40e0c1bb5091841946 Mon Sep 17 00:00:00 2001 From: baonudesifeizhai Date: Fri, 11 Jul 2025 02:06:37 -0400 Subject: [PATCH] feat: add sys.total_tokens system variable for workflow token usage tracking - Add TOTAL_TOKENS to SystemVariableKey enum - Initialize total_tokens to 0 in all workflow runners - Update total_tokens dynamically during workflow execution - Make total_tokens editable system variable - Add frontend support for sys.total_tokens variable Closes #21848 --- api/core/app/apps/advanced_chat/app_runner.py | 1 + api/core/app/apps/workflow/app_runner.py | 1 + api/core/workflow/enums.py | 1 + api/core/workflow/graph_engine/graph_engine.py | 7 +++++++ api/models/workflow.py | 4 ++-- api/services/workflow_service.py | 2 ++ .../workflow/nodes/_base/components/variable/utils.ts | 4 ++++ 7 files changed, 18 insertions(+), 2 deletions(-) diff --git a/api/core/app/apps/advanced_chat/app_runner.py b/api/core/app/apps/advanced_chat/app_runner.py index 840a3c9d3b..6a2b96dc73 100644 --- a/api/core/app/apps/advanced_chat/app_runner.py +++ b/api/core/app/apps/advanced_chat/app_runner.py @@ -145,6 +145,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner): SystemVariableKey.APP_ID: app_config.app_id, SystemVariableKey.WORKFLOW_ID: app_config.workflow_id, SystemVariableKey.WORKFLOW_EXECUTION_ID: self.application_generate_entity.workflow_run_id, + SystemVariableKey.TOTAL_TOKENS: 0, } # init variable pool diff --git a/api/core/app/apps/workflow/app_runner.py b/api/core/app/apps/workflow/app_runner.py index 07aeb57fa3..a41861dda7 100644 --- a/api/core/app/apps/workflow/app_runner.py +++ b/api/core/app/apps/workflow/app_runner.py @@ -101,6 +101,7 @@ class WorkflowAppRunner(WorkflowBasedAppRunner): SystemVariableKey.APP_ID: app_config.app_id, SystemVariableKey.WORKFLOW_ID: app_config.workflow_id, SystemVariableKey.WORKFLOW_EXECUTION_ID: self.application_generate_entity.workflow_execution_id, + SystemVariableKey.TOTAL_TOKENS: 0, } variable_pool = VariablePool( diff --git a/api/core/workflow/enums.py b/api/core/workflow/enums.py index b52a2b0e6e..755a6712e9 100644 --- a/api/core/workflow/enums.py +++ b/api/core/workflow/enums.py @@ -14,3 +14,4 @@ class SystemVariableKey(StrEnum): APP_ID = "app_id" WORKFLOW_ID = "workflow_id" WORKFLOW_EXECUTION_ID = "workflow_run_id" + TOTAL_TOKENS = "total_tokens" diff --git a/api/core/workflow/graph_engine/graph_engine.py b/api/core/workflow/graph_engine/graph_engine.py index 5a2915e2d3..aacd22cb57 100644 --- a/api/core/workflow/graph_engine/graph_engine.py +++ b/api/core/workflow/graph_engine/graph_engine.py @@ -14,9 +14,11 @@ from flask import Flask, current_app from configs import dify_config from core.app.apps.base_app_queue_manager import GenerateTaskStoppedError from core.app.entities.app_invoke_entities import InvokeFrom +from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID from core.workflow.entities.node_entities import AgentNodeStrategyInit, NodeRunResult from core.workflow.entities.variable_pool import VariablePool, VariableValue from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus +from core.workflow.enums import SystemVariableKey from core.workflow.graph_engine.condition_handlers.condition_manager import ConditionManager from core.workflow.graph_engine.entities.event import ( BaseAgentEvent, @@ -748,6 +750,11 @@ class GraphEngine: self.graph_runtime_state.total_tokens += int( run_result.metadata.get(WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS) # type: ignore[arg-type] ) + # Update system variable total_tokens + self.graph_runtime_state.variable_pool.add( + (SYSTEM_VARIABLE_NODE_ID, SystemVariableKey.TOTAL_TOKENS.value), + self.graph_runtime_state.total_tokens, + ) if run_result.llm_usage: # use the latest usage diff --git a/api/models/workflow.py b/api/models/workflow.py index 77d48bec4f..736e93a1ac 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -897,8 +897,8 @@ class ConversationVariable(Base): return variable_factory.build_conversation_variable_from_mapping(mapping) -# Only `sys.query` and `sys.files` could be modified. -_EDITABLE_SYSTEM_VARIABLE = frozenset(["query", "files"]) +# Only `sys.query`, `sys.files`, and `sys.total_tokens` could be modified. +_EDITABLE_SYSTEM_VARIABLE = frozenset(["query", "files", "total_tokens"]) def _naive_utc_datetime(): diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index 2be57fd51c..f95edc56a8 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -681,6 +681,8 @@ def _setup_variable_pool( SystemVariableKey.WORKFLOW_ID: workflow.id, # Randomly generated. SystemVariableKey.WORKFLOW_EXECUTION_ID: str(uuid.uuid4()), + # Initialize total tokens + SystemVariableKey.TOTAL_TOKENS: 0, } # Only add chatflow-specific variables for non-workflow types diff --git a/web/app/components/workflow/nodes/_base/components/variable/utils.ts b/web/app/components/workflow/nodes/_base/components/variable/utils.ts index ac95f54757..d14094c977 100644 --- a/web/app/components/workflow/nodes/_base/components/variable/utils.ts +++ b/web/app/components/workflow/nodes/_base/components/variable/utils.ts @@ -225,6 +225,10 @@ const formatItem = ( variable: 'sys.workflow_run_id', type: VarType.string, }) + res.vars.push({ + variable: 'sys.total_tokens', + type: VarType.number, + }) break }