@@ -359,9 +359,19 @@ def await_result(self, result):
359359 break
360360 yield chunk
361361 finally :
362+ details = self .get_runtime_details ()
363+ message_tokens = sum ([row .get ('message_tokens' ) for row in details .values () if
364+ 'message_tokens' in row and row .get ('message_tokens' ) is not None ])
365+ answer_tokens = sum ([row .get ('answer_tokens' ) for row in details .values () if
366+ 'answer_tokens' in row and row .get ('answer_tokens' ) is not None ])
362367 self .work_flow_post_handler .handler (self .params ['chat_id' ], self .params ['chat_record_id' ],
363368 self .answer ,
364369 self )
370+ yield self .base_to_response .to_stream_chunk_response (self .params ['chat_id' ],
371+ self .params ['chat_record_id' ],
372+ '' ,
373+ [],
374+ '' , True , message_tokens , answer_tokens , {})
365375
366376 def run_chain_async (self , current_node , node_result_future ):
367377 future = executor .submit (self .run_chain , current_node , node_result_future )
@@ -430,14 +440,14 @@ def append_node(self, current_node):
430440
431441 def hand_event_node_result (self , current_node , node_result_future ):
432442 node_chunk = NodeChunk ()
443+ real_node_id = current_node .runtime_node_id
444+ child_node = {}
433445 try :
434446 current_result = node_result_future .result ()
435447 result = current_result .write_context (current_node , self )
436448 if result is not None :
437449 if self .is_result (current_node , current_result ):
438450 self .node_chunk_manage .add_node_chunk (node_chunk )
439- child_node = {}
440- real_node_id = current_node .runtime_node_id
441451 for r in result :
442452 content = r
443453 child_node = {}
@@ -487,8 +497,12 @@ def hand_event_node_result(self, current_node, node_result_future):
487497 current_node .id ,
488498 current_node .up_node_id_list ,
489499 str (e ), False , 0 , 0 ,
490- {'node_is_end' : True , 'node_type' : current_node .type ,
491- 'view_type' : current_node .view_type })
500+ {'node_is_end' : True ,
501+ 'runtime_node_id' : current_node .runtime_node_id ,
502+ 'node_type' : current_node .type ,
503+ 'view_type' : current_node .view_type ,
504+ 'child_node' : {},
505+ 'real_node_id' : real_node_id })
492506 if not self .node_chunk_manage .contains (node_chunk ):
493507 self .node_chunk_manage .add_node_chunk (node_chunk )
494508 node_chunk .end (chunk )
0 commit comments