Procházet zdrojové kódy

fix: wrong token usage in iteration node for streaming result (#5336)

rerorero před 10 měsíci
rodič
revize
9b7fdadce4

+ 3 - 2
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
+    current_iterations: dict[str, Data] = None

+ 7 - 1
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
                 )
-            )
+            )