diff --git a/api/core/workflow/nodes/question_classifier/question_classifier_node.py b/api/core/workflow/nodes/question_classifier/question_classifier_node.py index 882a115acc..5ad2e1d593 100644 --- a/api/core/workflow/nodes/question_classifier/question_classifier_node.py +++ b/api/core/workflow/nodes/question_classifier/question_classifier_node.py @@ -13,7 +13,7 @@ from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, Comp from core.prompt.simple_prompt_transform import ModelMode from core.prompt.utils.prompt_message_util import PromptMessageUtil from core.workflow.entities.base_node_data_entities import BaseNodeData -from core.workflow.entities.node_entities import NodeRunResult, NodeType +from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult, NodeType from core.workflow.entities.variable_pool import VariablePool from core.workflow.nodes.llm.llm_node import LLMNode from core.workflow.nodes.question_classifier.entities import QuestionClassifierNodeData @@ -91,14 +91,24 @@ class QuestionClassifierNode(LLMNode): inputs=variables, process_data=process_data, outputs=outputs, - edge_source_handle=classes_map.get(categories[0], None) + edge_source_handle=classes_map.get(categories[0], None), + metadata={ + NodeRunMetadataKey.TOTAL_TOKENS: usage.total_tokens, + NodeRunMetadataKey.TOTAL_PRICE: usage.total_price, + NodeRunMetadataKey.CURRENCY: usage.currency + } ) except ValueError as e: return NodeRunResult( status=WorkflowNodeExecutionStatus.FAILED, inputs=variables, - error=str(e) + error=str(e), + metadata={ + NodeRunMetadataKey.TOTAL_TOKENS: usage.total_tokens, + NodeRunMetadataKey.TOTAL_PRICE: usage.total_price, + NodeRunMetadataKey.CURRENCY: usage.currency + } ) @classmethod