fix: missing usage of metadata in the chatflow app (langgenius#9167)

liuhaoran1212 · Oct 10, 2024 · 477beae · 477beae
1 parent 6689b59
commit 477beae
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/api/core/app/apps/advanced_chat/generate_task_pipeline.py b/api/core/app/apps/advanced_chat/generate_task_pipeline.py
@@ -45,6 +45,7 @@
 from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
 from core.app.task_pipeline.message_cycle_manage import MessageCycleManage
 from core.app.task_pipeline.workflow_cycle_manage import WorkflowCycleManage
+from core.model_runtime.entities.llm_entities import LLMUsage
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.workflow.enums import SystemVariableKey
@@ -509,6 +510,10 @@ def _save_message(self, graph_runtime_state: Optional[GraphRuntimeState] = None)
             self._message.total_price = usage.total_price
             self._message.currency = usage.currency
 
+            self._task_state.metadata["usage"] = jsonable_encoder(usage)
+        else:
+            self._task_state.metadata["usage"] = jsonable_encoder(LLMUsage.empty_usage())
+
         db.session.commit()
 
         message_was_created.send(