fix: workflow log run time error (#7130)

langgenius · Aug 9, 2024 · 7201b56 · 7201b56
1 parent 34cab0e
commit 7201b56
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 9 deletions.
diff --git a/api/core/app/task_pipeline/workflow_cycle_manage.py b/api/core/app/task_pipeline/workflow_cycle_manage.py
@@ -40,6 +40,7 @@
     WorkflowRunStatus,
     WorkflowRunTriggeredFrom,
 )
+from services.workflow_service import WorkflowService
 
 
 class WorkflowCycleManage(WorkflowIterationCycleManage):
@@ -97,7 +98,6 @@ def _init_workflow_run(self, workflow: Workflow,
 
     def _workflow_run_success(
         self, workflow_run: WorkflowRun,
-        start_at: float,
         total_tokens: int,
         total_steps: int,
         outputs: Optional[str] = None,
@@ -107,7 +107,6 @@ def _workflow_run_success(
         """
         Workflow run success
         :param workflow_run: workflow run
-        :param start_at: start time
         :param total_tokens: total tokens
         :param total_steps: total steps
         :param outputs: outputs
@@ -116,7 +115,7 @@ def _workflow_run_success(
         """
         workflow_run.status = WorkflowRunStatus.SUCCEEDED.value
         workflow_run.outputs = outputs
-        workflow_run.elapsed_time = time.perf_counter() - start_at
+        workflow_run.elapsed_time = WorkflowService.get_elapsed_time(workflow_run_id=workflow_run.id)
         workflow_run.total_tokens = total_tokens
         workflow_run.total_steps = total_steps
         workflow_run.finished_at = datetime.now(timezone.utc).replace(tzinfo=None)
@@ -139,7 +138,6 @@ def _workflow_run_success(
 
     def _workflow_run_failed(
         self, workflow_run: WorkflowRun,
-        start_at: float,
         total_tokens: int,
         total_steps: int,
         status: WorkflowRunStatus,
@@ -150,7 +148,6 @@ def _workflow_run_failed(
         """
         Workflow run failed
         :param workflow_run: workflow run
-        :param start_at: start time
         :param total_tokens: total tokens
         :param total_steps: total steps
         :param status: status
@@ -159,7 +156,7 @@ def _workflow_run_failed(
         """
         workflow_run.status = status.value
         workflow_run.error = error
-        workflow_run.elapsed_time = time.perf_counter() - start_at
+        workflow_run.elapsed_time = WorkflowService.get_elapsed_time(workflow_run_id=workflow_run.id)
         workflow_run.total_tokens = total_tokens
         workflow_run.total_steps = total_steps
         workflow_run.finished_at = datetime.now(timezone.utc).replace(tzinfo=None)
@@ -542,7 +539,6 @@ def _handle_workflow_finished(
         if isinstance(event, QueueStopEvent):
             workflow_run = self._workflow_run_failed(
                 workflow_run=workflow_run,
-                start_at=self._task_state.start_at,
                 total_tokens=self._task_state.total_tokens,
                 total_steps=self._task_state.total_steps,
                 status=WorkflowRunStatus.STOPPED,
@@ -565,7 +561,6 @@ def _handle_workflow_finished(
         elif isinstance(event, QueueWorkflowFailedEvent):
             workflow_run = self._workflow_run_failed(
                 workflow_run=workflow_run,
-                start_at=self._task_state.start_at,
                 total_tokens=self._task_state.total_tokens,
                 total_steps=self._task_state.total_steps,
                 status=WorkflowRunStatus.FAILED,
@@ -583,7 +578,6 @@ def _handle_workflow_finished(
 
             workflow_run = self._workflow_run_success(
                 workflow_run=workflow_run,
-                start_at=self._task_state.start_at,
                 total_tokens=self._task_state.total_tokens,
                 total_steps=self._task_state.total_steps,
                 outputs=outputs,

diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py
@@ -319,3 +319,25 @@ def validate_features_structure(self, app_model: App, features: dict) -> dict:
             )
         else:
             raise ValueError(f"Invalid app mode: {app_model.mode}")
+
+    @classmethod
+    def get_elapsed_time(cls, workflow_run_id: str) -> float:
+        """
+        Get elapsed time
+        """
+        elapsed_time = 0.0
+
+        # fetch workflow node execution by workflow_run_id
+        workflow_nodes = (
+            db.session.query(WorkflowNodeExecution)
+            .filter(WorkflowNodeExecution.workflow_run_id == workflow_run_id)
+            .order_by(WorkflowNodeExecution.created_at.asc())
+            .all()
+        )
+        if not workflow_nodes:
+            return elapsed_time
+
+        for node in workflow_nodes:
+            elapsed_time += node.elapsed_time
+
+        return elapsed_time