From 2a6b7d57cbf791e79b2f7164f262b0c99bf43cc7 Mon Sep 17 00:00:00 2001 From: takatost Date: Tue, 9 Apr 2024 19:25:08 +0800 Subject: [PATCH] fix: token is not logging of question classifier node (#3249) --- .../question_classifier_node.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) 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