From 4858cc19cdffdc80d6d991c5b125537d7006ad55 Mon Sep 17 00:00:00 2001 From: yihong0618 Date: Tue, 10 Dec 2024 22:14:35 +0800 Subject: [PATCH] Revert "fix: total tokens is wrong which is zero in inter way, close #11221 (#11224)" This reverts commit fd4be36991c2c1e604e73f5b8132123594e78ada. --- api/core/app/apps/advanced_chat/generate_task_pipeline.py | 5 +---- api/core/app/apps/workflow/generate_task_pipeline.py | 5 +---- 2 files changed, 2 insertions(+), 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 index cd12690e286103..0c8cd384d5a930 100644 --- 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 index 9966a1a9d13cf5..9e4921d6a22c5a 100644 --- 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,