feat: support for websocket #21

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

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,70 @@ 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');
-          });
-        }
-        read();
-      }, (err) =>{
-        emits('setChatLoading', false);
-        Message.error(err.message);
-      });
+  // 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(
+      () => {
+        webSocket.value.send(message);
+      },
+      () => {
+        Message.error(i18n.t('ConnectionClosedPleaseRetry'));
+      },
+  );
+};
+
+// 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((resolve) => resolve());
+  }
+  webSocket.value = new WebSocket(`${globalContext.webSocketUrl}/chat/`);
+  webSocket.value.onmessage = (e) => {
+    onMessage(e);
+  };
+  webSocket.value.onclose = () => {
+    emits('setChatLoading', false);
+  };
+  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) {
+    webSocket.value.close();
+  }
 };
+onBeforeUnmount(() => closeWebSocket());
 
 const reGenerate = () => {
   if (!props.localMessages.length) {

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"

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,
 };

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;

src/locale/zh-cn.js

@@ -13,6 +13,9 @@ const mZhCN = {
   NoModelChoosed: '未选择模型',
   VerifyingUserInfo: '登录中',
   Logout: '登出',
+  ConnectionClosedPleaseRetry: '连接异常，请重试',
+  ConnectionClosed: '连接已关闭',
+  ConnectionError: '连接异常',
 };
 
 export default mZhCN;

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,
     },