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

pull/5349/head
rerorero 2 years ago committed by GitHub
parent 132f5fb3de
commit 9b7fdadce4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -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

@ -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
)
)
)

Loading…
Cancel
Save