Merge bf1aa35974 into bf542233a9

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

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(

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"

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

api/core/workflow/nodes/answer/answer_stream_generate_router.py

@@ -151,6 +151,18 @@ class AnswerStreamGeneratorRouter:
                 continue
             source_node_type = node_id_config_mapping[source_node_id].get("data", {}).get("type")
             source_node_data = node_id_config_mapping[source_node_id].get("data", {})
+
+            # Check if this answer node uses sys.total_tokens and add LLM dependencies
+            if source_node_type == NodeType.LLM.value:
+                answer_node_config = node_id_config_mapping.get(answer_node_id)
+                if answer_node_config:
+                    answer_data = answer_node_config.get("data", {})
+                    answer_text = answer_data.get("answer", "")
+                    # Check if the answer template contains sys.total_tokens
+                    if "{{#sys.total_tokens#}}" in answer_text:
+                        answer_dependencies[answer_node_id].append(source_node_id)
+                        continue
+
             if (
                 source_node_type
                 in {

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():

api/services/workflow_service.py

@@ -695,6 +695,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

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
     }

web/app/components/workflow/nodes/start/panel.tsx

@@ -155,8 +155,18 @@ const Panel: FC<NodePanelProps<StartNodeType>> = ({
                   </div>
                 }
               />
+              <VarItem
+                readonly
+                payload={{
+                  variable: 'sys.total_tokens',
+                } as any}
+                rightContent={
+                  <div className='text-xs font-normal text-text-tertiary'>
+                    Number
+                  </div>
+                }
+              />
             </div>
-
           </>
         </Field>
       </div>