diff --git a/api/core/app/entities/task_entities.py b/api/core/app/entities/task_entities.py index cb023ae9f..af93399a2 100644 --- a/api/core/app/entities/task_entities.py +++ b/api/core/app/entities/task_entities.py @@ -367,7 +367,7 @@ class IterationNodeNextStreamResponse(StreamResponse): class IterationNodeCompletedStreamResponse(StreamResponse): """ - NodeStartStreamResponse entity + NodeCompletedStreamResponse entity """ class Data(BaseModel): """ @@ -385,6 +385,7 @@ class IterationNodeCompletedStreamResponse(StreamResponse): error: Optional[str] = None elapsed_time: float total_tokens: int + execution_metadata: Optional[dict] = None finished_at: int steps: int @@ -545,4 +546,4 @@ class WorkflowIterationState(BaseModel): total_tokens: int = 0 node_data: BaseNodeData - current_iterations: dict[str, Data] = None \ No newline at end of file + current_iterations: dict[str, Data] = None diff --git a/api/core/app/task_pipeline/workflow_iteration_cycle_manage.py b/api/core/app/task_pipeline/workflow_iteration_cycle_manage.py index 19bb7788c..69af81d02 100644 --- a/api/core/app/task_pipeline/workflow_iteration_cycle_manage.py +++ b/api/core/app/task_pipeline/workflow_iteration_cycle_manage.py @@ -95,6 +95,9 @@ class WorkflowIterationCycleManage(WorkflowCycleStateManager): error=None, elapsed_time=time.perf_counter() - current_iteration.started_at, total_tokens=current_iteration.total_tokens, + execution_metadata={ + 'total_tokens': current_iteration.total_tokens, + }, finished_at=int(time.time()), steps=current_iteration.current_index ) @@ -276,7 +279,10 @@ class WorkflowIterationCycleManage(WorkflowCycleStateManager): error=error, elapsed_time=time.perf_counter() - current_iteration.started_at, total_tokens=current_iteration.total_tokens, + execution_metadata={ + 'total_tokens': current_iteration.total_tokens, + }, finished_at=int(time.time()), steps=current_iteration.current_index ) - ) \ No newline at end of file + )