From 1c8a83140bfda937985bf1597f6a8b692c0c088e Mon Sep 17 00:00:00 2001
From: Eternity <1533512157@qq.com>
Date: Wed, 4 Feb 2026 18:08:02 +0800
Subject: [PATCH] feat(workflow): add token usage statistics for question
 classifier and parameter extraction

---
 .../core/workflow/nodes/parameter_extractor/node.py | 13 +++++++++++++
 .../core/workflow/nodes/question_classifier/node.py | 13 +++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/api/app/core/workflow/nodes/parameter_extractor/node.py b/api/app/core/workflow/nodes/parameter_extractor/node.py
index ec58d96c..079cd4cc 100644
--- a/api/app/core/workflow/nodes/parameter_extractor/node.py
+++ b/api/app/core/workflow/nodes/parameter_extractor/node.py
@@ -23,6 +23,18 @@ class ParameterExtractorNode(BaseNode):
     def __init__(self, node_config: dict[str, Any], workflow_config: dict[str, Any]):
         super().__init__(node_config, workflow_config)
         self.typed_config: ParameterExtractorNodeConfig | None = None
+        self.response_metadata = {}
+
+    def _extract_token_usage(self, business_result: Any) -> dict[str, int] | None:
+        if self.response_metadata:
+            usage = self.response_metadata.get('token_usage')
+            if usage:
+                return {
+                    "prompt_tokens": usage.get('prompt_tokens', 0),
+                    "completion_tokens": usage.get('completion_tokens', 0),
+                    "total_tokens": usage.get('total_tokens', 0)
+                }
+        return None
 
     @staticmethod
     def _get_prompt():
@@ -171,6 +183,7 @@ class ParameterExtractorNode(BaseNode):
             ])
 
         model_resp = await llm.ainvoke(messages)
+        self.response_metadata = model_resp.response_metadata
         result = json_repair.repair_json(model_resp.content, return_objects=True)
         logger.info(f"node: {self.node_id} get params:{result}")
 
diff --git a/api/app/core/workflow/nodes/question_classifier/node.py b/api/app/core/workflow/nodes/question_classifier/node.py
index 6df410cb..8076dc9d 100644
--- a/api/app/core/workflow/nodes/question_classifier/node.py
+++ b/api/app/core/workflow/nodes/question_classifier/node.py
@@ -23,6 +23,18 @@ class QuestionClassifierNode(BaseNode):
         super().__init__(node_config, workflow_config)
         self.typed_config: QuestionClassifierNodeConfig | None = None
         self.category_to_case_map = {}
+        self.response_metadata = {}
+
+    def _extract_token_usage(self, business_result: Any) -> dict[str, int] | None:
+        if self.response_metadata:
+            usage = self.response_metadata.get('token_usage')
+            if usage:
+                return {
+                    "prompt_tokens": usage.get('prompt_tokens', 0),
+                    "completion_tokens": usage.get('completion_tokens', 0),
+                    "total_tokens": usage.get('total_tokens', 0)
+                }
+        return None
 
     def _get_llm_instance(self) -> RedBearLLM:
         """获取LLM实例"""
@@ -112,6 +124,7 @@ class QuestionClassifierNode(BaseNode):
 
             response = await llm.ainvoke(messages)
             result = response.content.strip()
+            self.response_metadata = response.response_metadata
 
             if result in category_names:
                 category = result