From a7f3f422d295c5efd2142b93d2c8b9ef08789c28 Mon Sep 17 00:00:00 2001
From: orenzhang <orenzhang@ovinc.cn>
Date: Fri, 7 Jun 2024 16:18:39 +0800
Subject: [PATCH 1/3] feat: support for websocket #21

---
 src/api/chat.js                   |  15 -----
 src/components/ChatInput.vue      | 104 ++++++++++++++++++------------
 src/components/MessageContent.vue |   2 +
 src/context/index.js              |   1 +
 src/locale/en-us.js               |   3 +
 src/locale/zh-cn.js               |   3 +
 vite.config.js                    |   1 +
 7 files changed, 74 insertions(+), 55 deletions(-)

diff --git a/src/api/chat.js b/src/api/chat.js
index b6a3cfe..71fffc1 100644
--- a/src/api/chat.js
+++ b/src/api/chat.js
@@ -1,20 +1,5 @@
-import globalContext from '../context';
 import http from './index';
 
-export const createChatAPI = (data) => new Promise((resolve, reject) => {
-  fetch(`${globalContext.backendUrl}/chat/`, {
-    method: 'POST',
-    headers: {
-      'Content-Type': 'application/json',
-    },
-    body: JSON.stringify(data),
-    credentials: 'include',
-  }).then(
-      (res) => resolve(res),
-      (err) => reject(err),
-  );
-});
-
 export const preCheckAPI = (data) => new Promise((resolve, reject) => {
   http.post('/chat/pre_check/', data).then(
       (res) => resolve(res),
diff --git a/src/components/ChatInput.vue b/src/components/ChatInput.vue
index e0d8b54..c62fdd7 100644
--- a/src/components/ChatInput.vue
+++ b/src/components/ChatInput.vue
@@ -1,10 +1,11 @@
 <script setup>
-import {computed, ref, watch} from 'vue';
-import {createChatAPI, preCheckAPI} from '../api/chat';
+import {computed, onBeforeUnmount, ref, watch} from 'vue';
+import {preCheckAPI} from '../api/chat';
 import {Message} from '@arco-design/web-vue';
 import {Role} from '../constants';
 import {useStore} from 'vuex';
 import {useI18n} from 'vue-i18n';
+import globalContext from '../context';
 
 // props
 const props = defineProps({
@@ -78,6 +79,7 @@ watch(() => model.value, () => {
 });
 
 // chat
+const lastResponseContent = ref(null);
 const doChat = async () => {
   // set loading
   emits('setChatLoading', true);
@@ -98,46 +100,68 @@ const doChat = async () => {
   if (!key) {
     return;
   }
-  createChatAPI({key})
-      .then((res) => {
-        // check success
-        if (!res.ok) {
-          emits('setChatLoading', false);
-          Message.error(res.statusText);
-          return;
-        }
-        // push message
-        emits('addMessage', {role: 'user', content: promptForm.value.content});
-        // auto scroll
-        emits('toggleUserBehavior', false);
-        // init response content
-        const lastResponseContent = ref({role: Role.Assistant, content: ''});
-        emits('addMessage', lastResponseContent.value);
-        // clear input
-        promptForm.value.content = '';
-        // decode chunk
-        const decoder = new TextDecoder('utf-8');
-        const reader = res.body.getReader();
-        // eslint-disable-next-line require-jsdoc
-        function read() {
-          reader.read().then((chunk) => {
-            if (!chunk.done) {
-              const value = decoder.decode(chunk.value);
-              lastResponseContent.value.content += value;
-              read();
-            }
-            if (chunk.done) {
-              emits('setChatLoading', false);
-            }
-            emits('saveMessage');
-          });
+  // add message to display
+  emits('addMessage', {role: 'user', content: promptForm.value.content});
+  // auto scroll
+  emits('toggleUserBehavior', false);
+  // init response content
+  lastResponseContent.value = {role: Role.Assistant, content: ''};
+  emits('addMessage', lastResponseContent.value);
+  // clear input
+  promptForm.value.content = '';
+  // send message
+  sendMessage(JSON.stringify({key}), true);
+};
+const onMessage = (event) => {
+  const data = JSON.parse(event.data);
+  if (data.is_finished) {
+    emits('setChatLoading', false);
+  }
+  if (data.data) {
+    lastResponseContent.value.content += data.data;
+    emits('saveMessage');
+  }
+};
+const sendMessage = (message) => {
+  initWebSocket().then(
+      () => {
+        if (webSocket.value && webSocket.value.readyState === WebSocket.OPEN) {
+          webSocket.value.send(message);
+        } else {
+          Message.error(i18n.t('ConnectionClosedPleaseRetry'));
         }
-        read();
-      }, (err) =>{
-        emits('setChatLoading', false);
-        Message.error(err.message);
-      });
+      },
+  );
+};
+
+// webSocket
+const webSocket = ref(null);
+const retryTimes = ref(0);
+const maxRetryTimes = ref(1000);
+const initWebSocket = () => {
+  if (webSocket.value && webSocket.value.readyState === WebSocket.OPEN) {
+    return new Promise();
+  }
+  webSocket.value = new WebSocket(`${globalContext.webSocketUrl}/chat/`);
+  webSocket.value.onmessage = (e) => {
+    onMessage(e);
+  };
+  webSocket.value.onclose = () => {
+    emits('setChatLoading', false);
+  };
+  webSocket.value.onerror = (e) => {
+    emits('setChatLoading', false);
+  };
+  return new Promise((resolve) => {
+    webSocket.value.onopen = (e) => resolve(e);
+  });
+};
+const closeWebSocket = () => {
+  if (webSocket.value && webSocket.value.readyState === WebSocket.OPEN) {
+    webSocket.value.close();
+  }
 };
+onBeforeUnmount(() => closeWebSocket());
 
 const reGenerate = () => {
   if (!props.localMessages.length) {
diff --git a/src/components/MessageContent.vue b/src/components/MessageContent.vue
index eb0d1ec..b957675 100644
--- a/src/components/MessageContent.vue
+++ b/src/components/MessageContent.vue
@@ -49,9 +49,11 @@ const emits = defineEmits(['reGenerate']);
       }"
     >
       <v-md-preview
+        v-show="message.content"
         :text="message.content"
         class="v-md-preview"
       />
+      <icon-loading v-show="!message.content" />
     </div>
     <a-avatar
       v-show="false"
diff --git a/src/context/index.js b/src/context/index.js
index 8f82a4c..95e288b 100644
--- a/src/context/index.js
+++ b/src/context/index.js
@@ -1,6 +1,7 @@
 const globalContext = {
   siteUrl: process.env.SITE_URL,
   backendUrl: process.env.BACKEND_URL,
+  webSocketUrl: process.env.WEB_SOCKET_URL,
   ovincUrl: process.env.OVINC_URL,
   ovincWebUrl: process.env.OVINC_WEB_URL,
 };
diff --git a/src/locale/en-us.js b/src/locale/en-us.js
index aca9d72..3e1b8dc 100644
--- a/src/locale/en-us.js
+++ b/src/locale/en-us.js
@@ -13,6 +13,9 @@ const mEnUS = {
   NoModelChoosed: 'No Model Choosed',
   VerifyingUserInfo: 'Verifying...',
   Logout: 'Logout',
+  ConnectionClosedPleaseRetry: 'Connection is closed, please try again',
+  ConnectionClosed: 'Connection is closed',
+  ConnectionError: 'Connection Error',
 };
 
 export default mEnUS;
diff --git a/src/locale/zh-cn.js b/src/locale/zh-cn.js
index 3376d35..4626023 100644
--- a/src/locale/zh-cn.js
+++ b/src/locale/zh-cn.js
@@ -13,6 +13,9 @@ const mZhCN = {
   NoModelChoosed: '未选择模型',
   VerifyingUserInfo: '登录中',
   Logout: '登出',
+  ConnectionClosedPleaseRetry: '连接异常，请重试',
+  ConnectionClosed: '连接已关闭',
+  ConnectionError: '连接异常',
 };
 
 export default mZhCN;
diff --git a/vite.config.js b/vite.config.js
index 66e5406..d928f7e 100644
--- a/vite.config.js
+++ b/vite.config.js
@@ -7,6 +7,7 @@ export default defineConfig({
     'process.env': {
       BACKEND_URL: process.env.BACKEND_URL,
       SITE_URL: process.env.SITE_URL,
+      WEB_SOCKET_URL: process.env.WEB_SOCKET_URL,
       OVINC_URL: process.env.OVINC_URL,
       OVINC_WEB_URL: process.env.OVINC_WEB_URL,
     },

From 9911ba55df3c101c9d46a60cc388260c7ee84948 Mon Sep 17 00:00:00 2001
From: orenzhang <orenzhang@ovinc.cn>
Date: Fri, 7 Jun 2024 17:18:18 +0800
Subject: [PATCH 2/3] fix: promise has not resolve #21

---
 src/components/ChatInput.vue | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/components/ChatInput.vue b/src/components/ChatInput.vue
index c62fdd7..9091381 100644
--- a/src/components/ChatInput.vue
+++ b/src/components/ChatInput.vue
@@ -125,11 +125,10 @@ const onMessage = (event) => {
 const sendMessage = (message) => {
   initWebSocket().then(
       () => {
-        if (webSocket.value && webSocket.value.readyState === WebSocket.OPEN) {
-          webSocket.value.send(message);
-        } else {
-          Message.error(i18n.t('ConnectionClosedPleaseRetry'));
-        }
+        webSocket.value.send(message);
+      },
+      () => {
+        Message.error(i18n.t('ConnectionClosedPleaseRetry'));
       },
   );
 };
@@ -138,9 +137,20 @@ const sendMessage = (message) => {
 const webSocket = ref(null);
 const retryTimes = ref(0);
 const maxRetryTimes = ref(1000);
+const buildWebSocketPromise = () => {
+  return new Promise((resolve, reject) => {
+    webSocket.value.onopen = (e) => {
+      resolve(e);
+    };
+    webSocket.value.onerror = (e) => {
+      emits('setChatLoading', false);
+      reject(e);
+    };
+  });
+};
 const initWebSocket = () => {
   if (webSocket.value && webSocket.value.readyState === WebSocket.OPEN) {
-    return new Promise();
+    return buildWebSocketPromise();
   }
   webSocket.value = new WebSocket(`${globalContext.webSocketUrl}/chat/`);
   webSocket.value.onmessage = (e) => {
@@ -149,12 +159,7 @@ const initWebSocket = () => {
   webSocket.value.onclose = () => {
     emits('setChatLoading', false);
   };
-  webSocket.value.onerror = (e) => {
-    emits('setChatLoading', false);
-  };
-  return new Promise((resolve) => {
-    webSocket.value.onopen = (e) => resolve(e);
-  });
+  return buildWebSocketPromise();
 };
 const closeWebSocket = () => {
   if (webSocket.value && webSocket.value.readyState === WebSocket.OPEN) {

From c3ad6ecaf74a06ed066f8cdd99a507dbc44264b2 Mon Sep 17 00:00:00 2001
From: orenzhang <orenzhang@ovinc.cn>
Date: Fri, 7 Jun 2024 17:27:08 +0800
Subject: [PATCH 3/3] fix: initialized ws cannot reuse #21

---
 src/components/ChatInput.vue | 23 ++++++++++-------------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/src/components/ChatInput.vue b/src/components/ChatInput.vue
index 9091381..7e099a9 100644
--- a/src/components/ChatInput.vue
+++ b/src/components/ChatInput.vue
@@ -137,20 +137,9 @@ const sendMessage = (message) => {
 const webSocket = ref(null);
 const retryTimes = ref(0);
 const maxRetryTimes = ref(1000);
-const buildWebSocketPromise = () => {
-  return new Promise((resolve, reject) => {
-    webSocket.value.onopen = (e) => {
-      resolve(e);
-    };
-    webSocket.value.onerror = (e) => {
-      emits('setChatLoading', false);
-      reject(e);
-    };
-  });
-};
 const initWebSocket = () => {
   if (webSocket.value && webSocket.value.readyState === WebSocket.OPEN) {
-    return buildWebSocketPromise();
+    return new Promise((resolve) => resolve());
   }
   webSocket.value = new WebSocket(`${globalContext.webSocketUrl}/chat/`);
   webSocket.value.onmessage = (e) => {
@@ -159,7 +148,15 @@ const initWebSocket = () => {
   webSocket.value.onclose = () => {
     emits('setChatLoading', false);
   };
-  return buildWebSocketPromise();
+  return new Promise((resolve, reject) => {
+    webSocket.value.onopen = (e) => {
+      resolve(e);
+    };
+    webSocket.value.onerror = (e) => {
+      emits('setChatLoading', false);
+      reject(e);
+    };
+  });
 };
 const closeWebSocket = () => {
   if (webSocket.value && webSocket.value.readyState === WebSocket.OPEN) {