From ffb7dd63dbd1a051d5d93862ca921cd9a044efda Mon Sep 17 00:00:00 2001 From: shaohuzhang1 Date: Wed, 26 Jun 2024 11:54:17 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E5=9B=9E=E7=AD=94=E9=94=99=E8=AF=AF=E6=97=B6,=E6=8E=A5?= =?UTF-8?q?=E5=8F=A3=E6=8A=A5=E9=94=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ai_chat_step_node/impl/base_chat_node.py | 11 +++++++--- apps/application/flow/tools.py | 21 ++++++++++++------- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py index b33736aa7..6265704ea 100644 --- a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py +++ b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py @@ -65,10 +65,15 @@ def write_context(node_variable: Dict, workflow_variable: Dict, node: INode, wor def get_to_response_write_context(node_variable: Dict, node: INode): - def _write_context(answer): + def _write_context(answer, status=200): chat_model = node_variable.get('chat_model') message_tokens = chat_model.get_num_tokens_from_messages(node_variable.get('message_list')) - answer_tokens = chat_model.get_num_tokens(answer) + if status == 200: + answer_tokens = chat_model.get_num_tokens(answer) + else: + answer_tokens = 0 + node.err_message = answer + node.status = status node.context['message_tokens'] = message_tokens node.context['answer_tokens'] = answer_tokens node.context['answer'] = answer @@ -170,7 +175,7 @@ class BaseChatNode(IChatNode): def get_details(self, index: int, **kwargs): return { - 'name':self.node.properties.get('stepName'), + 'name': self.node.properties.get('stepName'), "index": index, 'run_time': self.context.get('run_time'), 'system': self.node_params.get('system'), diff --git a/apps/application/flow/tools.py b/apps/application/flow/tools.py index e791f6c11..839aae8da 100644 --- a/apps/application/flow/tools.py +++ b/apps/application/flow/tools.py @@ -29,14 +29,21 @@ def event_content(chat_id, chat_record_id, response, workflow, @param post_handler: 后置处理器 """ answer = '' - for chunk in response: - answer += chunk.content + try: + for chunk in response: + answer += chunk.content + yield 'data: ' + json.dumps({'chat_id': str(chat_id), 'id': str(chat_record_id), 'operate': True, + 'content': chunk.content, 'is_end': False}, ensure_ascii=False) + "\n\n" + write_context(answer, 200) + post_handler.handler(chat_id, chat_record_id, answer, workflow) yield 'data: ' + json.dumps({'chat_id': str(chat_id), 'id': str(chat_record_id), 'operate': True, - 'content': chunk.content, 'is_end': False}, ensure_ascii=False) + "\n\n" - write_context(answer) - post_handler.handler(chat_id, chat_record_id, answer, workflow) - yield 'data: ' + json.dumps({'chat_id': str(chat_id), 'id': str(chat_record_id), 'operate': True, - 'content': '', 'is_end': True}, ensure_ascii=False) + "\n\n" + 'content': '', 'is_end': True}, ensure_ascii=False) + "\n\n" + except Exception as e: + answer = str(e) + write_context(answer, 500) + post_handler.handler(chat_id, chat_record_id, answer, workflow) + yield 'data: ' + json.dumps({'chat_id': str(chat_id), 'id': str(chat_record_id), 'operate': True, + 'content': answer, 'is_end': True}, ensure_ascii=False) + "\n\n" def to_stream_response(chat_id, chat_record_id, response: Iterator[BaseMessageChunk], workflow, write_context,