Revert "fix: total tokens is wrong which is zero in inter way, close #…

…11221 (#11224)" This reverts commit fd4be36.
langgenius · Dec 10, 2024 · 4858cc1 · 4858cc1
1 parent 28231d3
commit 4858cc1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 8 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
@@ -127,7 +127,6 @@ def __init__(
 
         self._conversation_name_generate_thread = None
         self._recorded_files: list[Mapping[str, Any]] = []
-        self.total_tokens: int = 0
 
     def process(self):
         """
@@ -361,8 +360,6 @@ def _process_stream_response(
                 if not workflow_run:
                     raise Exception("Workflow run not initialized.")
 
-                # FIXME for issue #11221 quick fix maybe have a better solution
-                self.total_tokens += event.metadata.get("total_tokens", 0) if event.metadata else 0
                 yield self._workflow_iteration_completed_to_stream_response(
                     task_id=self._application_generate_entity.task_id, workflow_run=workflow_run, event=event
                 )
@@ -376,7 +373,7 @@ def _process_stream_response(
                 workflow_run = self._handle_workflow_run_success(
                     workflow_run=workflow_run,
                     start_at=graph_runtime_state.start_at,
-                    total_tokens=graph_runtime_state.total_tokens or self.total_tokens,
+                    total_tokens=graph_runtime_state.total_tokens,
                     total_steps=graph_runtime_state.node_run_steps,
                     outputs=event.outputs,
                     conversation_id=self._conversation.id,

diff --git a/api/core/app/apps/workflow/generate_task_pipeline.py b/api/core/app/apps/workflow/generate_task_pipeline.py
@@ -106,7 +106,6 @@ def __init__(
 
         self._task_state = WorkflowTaskState()
         self._wip_workflow_node_executions = {}
-        self.total_tokens: int = 0
 
     def process(self) -> Union[WorkflowAppBlockingResponse, Generator[WorkflowAppStreamResponse, None, None]]:
         """
@@ -320,8 +319,6 @@ def _process_stream_response(
                 if not workflow_run:
                     raise Exception("Workflow run not initialized.")
 
-                # FIXME for issue #11221 quick fix maybe have a better solution
-                self.total_tokens += event.metadata.get("total_tokens", 0) if event.metadata else 0
                 yield self._workflow_iteration_completed_to_stream_response(
                     task_id=self._application_generate_entity.task_id, workflow_run=workflow_run, event=event
                 )
@@ -335,7 +332,7 @@ def _process_stream_response(
                 workflow_run = self._handle_workflow_run_success(
                     workflow_run=workflow_run,
                     start_at=graph_runtime_state.start_at,
-                    total_tokens=graph_runtime_state.total_tokens or self.total_tokens,
+                    total_tokens=graph_runtime_state.total_tokens,
                     total_steps=graph_runtime_state.node_run_steps,
                     outputs=event.outputs,
                     conversation_id=None,