adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(workflow): optimize token consumption tracking in question classifier and parameter extractor nodes

Browse Source
This commit is contained in:
Eternity
2026-02-09 20:19:15 +08:00
parent aa42cd0aec
commit d57e3b3f64
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
api/app/core/workflow/nodes/parameter_extractor/node.py
View File

@@ -32,8 +32,8 @@ class ParameterExtractorNode(BaseNode):
usage = self.response_metadata.get('token_usage') usage = self.response_metadata.get('token_usage')
if usage: if usage:
return { return {
"prompt_tokens": usage.get('prompt_tokens', 0), "prompt_tokens": usage.get('input_tokens', 0),
"completion_tokens": usage.get('completion_tokens', 0), "completion_tokens": usage.get('output_tokens', 0),
"total_tokens": usage.get('total_tokens', 0) "total_tokens": usage.get('total_tokens', 0)
} }
return None return None

4
api/app/core/workflow/nodes/question_classifier/node.py
View File

@@ -32,8 +32,8 @@ class QuestionClassifierNode(BaseNode):
usage = self.response_metadata.get('token_usage') usage = self.response_metadata.get('token_usage')
if usage: if usage:
return { return {
"prompt_tokens": usage.get('prompt_tokens', 0), "prompt_tokens": usage.get('input_tokens', 0),
"completion_tokens": usage.get('completion_tokens', 0), "completion_tokens": usage.get('output_tokens', 0),
"total_tokens": usage.get('total_tokens', 0) "total_tokens": usage.get('total_tokens', 0)
} }
return None return None