fix: wrong token usage in iteration node for streaming result (langge…

…nius#5336)
HuberyHuV1 · Jun 18, 2024 · 7051ad6 · 7051ad6
1 parent 3005dd8
commit 7051ad6
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/api/core/app/entities/task_entities.py b/api/core/app/entities/task_entities.py
@@ -367,7 +367,7 @@ class Data(BaseModel):
 
 class IterationNodeCompletedStreamResponse(StreamResponse):
  """
- NodeStartStreamResponse entity
+ NodeCompletedStreamResponse entity
  """
  class Data(BaseModel):
  """
@@ -385,6 +385,7 @@ class Data(BaseModel):
  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 Data(BaseModel):
  total_tokens: int = 0
  node_data: BaseNodeData
 
- current_iterations: dict[str, Data] = None
+ 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
@@ -95,6 +95,9 @@ def _handle_iteration_to_stream_response(self, task_id: str, event: QueueIterati
  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 @@ def _handle_iteration_exception(self, task_id: str, error: str) -> Generator[Ite
  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
  )
- )
+ )