From d7f6ec1fba447f2764d11fffa5dc6a4f271aa990 Mon Sep 17 00:00:00 2001
From: shaohuzhang1 <shaohu.zhang@fit2cloud.com>
Date: Fri, 15 Dec 2023 18:25:08 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E5=AF=B9=E8=AF=9D=E6=B5=81=E5=A4=84?=
 =?UTF-8?q?=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../serializers/chat_message_serializers.py   |  4 +++-
 ui/src/api/type/application.ts                |  2 --
 ui/src/components/ai-chat/index.vue           | 22 ++++++++++++++-----
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/apps/application/serializers/chat_message_serializers.py b/apps/application/serializers/chat_message_serializers.py
index 7d8c62927..981a4e380 100644
--- a/apps/application/serializers/chat_message_serializers.py
+++ b/apps/application/serializers/chat_message_serializers.py
@@ -165,8 +165,10 @@ class ChatMessageSerializer(serializers.Serializer):
                 for chunk in response:
                     all_text += chunk.content
                     yield 'data: ' + json.dumps({'chat_id': chat_id, 'id': _id, 'operate': paragraph is not None,
-                                                 'content': chunk.content}) + "\n\n"
+                                                 'content': chunk.content, 'is_end': False}) + "\n\n"
 
+                yield 'data: ' + json.dumps({'chat_id': chat_id, 'id': _id, 'operate': paragraph is not None,
+                                             'content': '', 'is_end': True}) + "\n\n"
                 chat_info.append_chat_message(
                     ChatMessage(_id, message, title, content, embedding_id, dataset_id, document_id,
                                 paragraph_id,
diff --git a/ui/src/api/type/application.ts b/ui/src/api/type/application.ts
index 883f7ab71..a9a0a3de4 100644
--- a/ui/src/api/type/application.ts
+++ b/ui/src/api/type/application.ts
@@ -57,7 +57,6 @@ export class ChatRecordManage {
           this.chat.write_ed = true
           this.write_ed = true
           if (this.loading) {
-            console.log('停止')
             this.loading.value = false
           }
           if (this.id) {
@@ -76,7 +75,6 @@ export class ChatRecordManage {
     }
   }
   close() {
-    console.log('close')
     this.is_close = true
   }
   append(answer_text_block: string) {
diff --git a/ui/src/components/ai-chat/index.vue b/ui/src/components/ai-chat/index.vue
index d22ad0444..51b1cfccc 100644
--- a/ui/src/components/ai-chat/index.vue
+++ b/ui/src/components/ai-chat/index.vue
@@ -282,13 +282,17 @@ function chatMessage() {
             if (str && str.startsWith('data:')) {
               const split = str.match(/data:.*}\n\n/g)
               if (split) {
-                split.forEach((item_str) => {
-                  row.record_id = JSON?.parse(item_str.replace('data:', '')).id
-                  const content = JSON?.parse(item_str.replace('data:', ''))?.content
+                for (const index in split) {
+                  const chunk = JSON?.parse(split[index].replace('data:', ''))
+                  row.record_id = chunk.id
+                  const content = chunk?.content
                   if (content) {
                     ChatManagement.append(id, content)
                   }
-                })
+                  if (chunk.is_end) {
+                    ChatManagement.close(id)
+                  }
+                }
               }
             }
           } catch (e) {
@@ -297,7 +301,15 @@ function chatMessage() {
           }
           return reader.read().then(write)
         }
-        reader.read().then(write)
+        reader
+          .read()
+          .then(write)
+          .finally((ok: any) => {
+            ChatManagement.close(id)
+          })
+          .catch((e: any) => {
+            ChatManagement.close(id)
+          })
       }
     })
   }