feat: Support for conversational streaming (infiniflow#809)

### What problem does this PR solve? feat: Support for conversational streaming infiniflow#709 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
aigc-hf · May 16, 2024 · c6c9dbd · c6c9dbd
1 parent 95f8091
commit c6c9dbd
Show file tree

Hide file tree

Showing 21 changed files with 426 additions and 257 deletions.
diff --git a/web/.env b/web/.env
@@ -0,0 +1 @@
+PORT=9222
diff --git a/web/package-lock.json b/web/package-lock.json
diff --git a/web/package.json b/web/package.json
@@ -3,7 +3,7 @@
   "author": "zhaofengchao <[email protected]>",
   "scripts": {
     "build": "umi build",
-    "dev": "cross-env PORT=9200 UMI_DEV_SERVER_COMPRESS=none umi dev",
+    "dev": "cross-env UMI_DEV_SERVER_COMPRESS=none umi dev",
     "postinstall": "umi setup",
     "lint": "umi lint --eslint-only",
     "setup": "umi setup",
@@ -19,6 +19,7 @@
     "axios": "^1.6.3",
     "classnames": "^2.5.1",
     "dayjs": "^1.11.10",
+    "eventsource-parser": "^1.1.2",
     "i18next": "^23.7.16",
     "js-base64": "^3.7.5",
     "jsencrypt": "^3.3.2",

diff --git a/web/src/components/new-document-link.tsx b/web/src/components/new-document-link.tsx
@@ -18,7 +18,7 @@ const NewDocumentLink = ({
       onClick={!preventDefault ? undefined : (e) => e.preventDefault()}
       href={link}
       rel="noreferrer"
-      style={{ color }}
+      style={{ color, wordBreak: 'break-all' }}
     >
       {children}
     </a>

diff --git a/web/src/hooks/chatHooks.ts b/web/src/hooks/chatHooks.ts
@@ -154,6 +154,9 @@ export const useRemoveConversation = () => {
   return removeConversation;
 };
 
+/*
+@deprecated
+ */
 export const useCompleteConversation = () => {
   const dispatch = useDispatch();
 
@@ -283,20 +286,4 @@ export const useFetchSharedConversation = () => {
   return fetchSharedConversation;
 };
 
-export const useCompleteSharedConversation = () => {
-  const dispatch = useDispatch();
-
-  const completeSharedConversation = useCallback(
-    (payload: any) => {
-      return dispatch<any>({
-        type: 'chatModel/completeExternalConversation',
-        payload: payload,
-      });
-    },
-    [dispatch],
-  );
-
-  return completeSharedConversation;
-};
-
 //#endregion
diff --git a/web/src/hooks/logicHooks.ts b/web/src/hooks/logicHooks.ts
@@ -1,13 +1,14 @@
 import { Authorization } from '@/constants/authorization';
 import { LanguageTranslationMap } from '@/constants/common';
 import { Pagination } from '@/interfaces/common';
+import { IAnswer } from '@/interfaces/database/chat';
 import { IKnowledgeFile } from '@/interfaces/database/knowledge';
 import { IChangeParserConfigRequestBody } from '@/interfaces/request/document';
 import api from '@/utils/api';
-import authorizationUtil from '@/utils/authorizationUtil';
-import { getSearchValue } from '@/utils/commonUtil';
+import { getAuthorization } from '@/utils/authorizationUtil';
 import { PaginationProps } from 'antd';
 import axios from 'axios';
+import { EventSourceParserStream } from 'eventsource-parser/stream';
 import { useCallback, useEffect, useMemo, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useDispatch } from 'umi';
@@ -138,62 +139,60 @@ export const useFetchAppConf = () => {
   return appConf;
 };
 
-export const useConnectWithSse = (url: string) => {
-  const [content, setContent] = useState<string>('');
-
-  const connect = useCallback(() => {
-    const source = new EventSource(
-      url || '/sse/createSseEmitter?clientId=123456',
-    );
-
-    source.onopen = function () {
-      console.log('Connection to the server was opened.');
-    };
-
-    source.onmessage = function (event: any) {
-      setContent(event.data);
-    };
-
-    source.onerror = function (error) {
-      console.error('Error occurred:', error);
-    };
-  }, [url]);
-
-  return { connect, content };
-};
+export const useSendMessageWithSse = (
+  url: string = api.completeConversation,
+) => {
+  const [answer, setAnswer] = useState<IAnswer>({} as IAnswer);
+  const [done, setDone] = useState(true);
 
-export const useConnectWithSseNext = () => {
-  const [content, setContent] = useState<string>('');
-  const sharedId = getSearchValue('shared_id');
-  const authorization = sharedId
-    ? 'Bearer ' + sharedId
-    : authorizationUtil.getAuthorization();
   const send = useCallback(
     async (body: any) => {
-      const response = await fetch(api.completeConversation, {
-        method: 'POST',
-        headers: {
-          [Authorization]: authorization,
-          'Content-Type': 'application/json',
-        },
-        body: JSON.stringify(body),
-      });
-      const reader = response?.body
-        ?.pipeThrough(new TextDecoderStream())
-        .getReader();
-
-      // const reader = response.body.getReader();
-
-      while (true) {
-        const { value, done } = await reader?.read();
-        console.log('Received', value);
-        setContent(value);
-        if (done) break;
+      try {
+        setDone(false);
+        const response = await fetch(url, {
+          method: 'POST',
+          headers: {
+            [Authorization]: getAuthorization(),
+            'Content-Type': 'application/json',
+          },
+          body: JSON.stringify(body),
+        });
+
+        const reader = response?.body
+          ?.pipeThrough(new TextDecoderStream())
+          .pipeThrough(new EventSourceParserStream())
+          .getReader();
+
+        while (true) {
+          const x = await reader?.read();
+          if (x) {
+            const { done, value } = x;
+            try {
+              const val = JSON.parse(value?.data || '');
+              const d = val?.data;
+              if (typeof d !== 'boolean') {
+                console.info('data:', d);
+                setAnswer(d);
+              }
+            } catch (e) {
+              console.warn(e);
+            }
+            if (done) {
+              console.info('done');
+              break;
+            }
+          }
+        }
+        console.info('done?');
+        setDone(true);
+        return response;
+      } catch (e) {
+        setDone(true);
+        console.warn(e);
       }
-      return response;
     },
-    [authorization],
+    [url],
   );
 
-  return { send, content };
+  return { send, answer, done };
 };
diff --git a/web/src/interfaces/database/chat.ts b/web/src/interfaces/database/chat.ts
@@ -72,6 +72,11 @@ export interface IReference {
   total: number;
 }
 
+export interface IAnswer {
+  answer: string;
+  reference: IReference;
+}
+
 export interface Docagg {
   count: number;
   doc_id: string;

diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts
@@ -25,6 +25,7 @@ export default {
       comingSoon: 'Coming Soon',
       download: 'Download',
       close: 'Close',
+      preview: 'Preview',
     },
     login: {
       login: 'Sign in',
@@ -381,6 +382,7 @@ export default {
       partialTitle: 'Partial Embed',
       extensionTitle: 'Chrome Extension',
       tokenError: 'Please create API Token first!',
+      searching: 'searching...',
     },
     setting: {
       profile: 'Profile',

diff --git a/web/src/locales/zh-traditional.ts b/web/src/locales/zh-traditional.ts
@@ -25,6 +25,7 @@ export default {
       comingSoon: '即將推出',
       download: '下載',
       close: '关闭',
+      preview: '預覽',
     },
     login: {
       login: '登入',
@@ -352,6 +353,7 @@ export default {
       partialTitle: '部分嵌入',
       extensionTitle: 'Chrome 插件',
       tokenError: '請先創建 Api Token!',
+      searching: '搜索中',
     },
     setting: {
       profile: '概述',

diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts
@@ -25,6 +25,7 @@ export default {
       comingSoon: '即将推出',
       download: '下载',
       close: '关闭',
+      preview: '预览',
     },
     login: {
       login: '登录',
@@ -369,6 +370,7 @@ export default {
       partialTitle: '部分嵌入',
       extensionTitle: 'Chrome 插件',
       tokenError: '请先创建 Api Token!',
+      searching: '搜索中',
     },
     setting: {
       profile: '概要',

diff --git a/web/src/pages/chat/chat-container/index.tsx b/web/src/pages/chat/chat-container/index.tsx
@@ -6,16 +6,7 @@ import { useSelectFileThumbnails } from '@/hooks/knowledgeHook';
 import { useSelectUserInfo } from '@/hooks/userSettingHook';
 import { IReference, Message } from '@/interfaces/database/chat';
 import { IChunk } from '@/interfaces/database/knowledge';
-import {
-  Avatar,
-  Button,
-  Drawer,
-  Flex,
-  Input,
-  List,
-  Skeleton,
-  Spin,
-} from 'antd';
+import { Avatar, Button, Drawer, Flex, Input, List, Spin } from 'antd';
 import classNames from 'classnames';
 import { useMemo } from 'react';
 import {
@@ -32,27 +23,39 @@ import SvgIcon from '@/components/svg-icon';
 import { useTranslate } from '@/hooks/commonHooks';
 import { useGetDocumentUrl } from '@/hooks/documentHooks';
 import { getExtension, isPdf } from '@/utils/documentUtils';
+import { buildMessageItemReference } from '../utils';
 import styles from './index.less';
 
 const MessageItem = ({
   item,
   reference,
+  loading = false,
   clickDocumentButton,
 }: {
   item: Message;
   reference: IReference;
+  loading?: boolean;
   clickDocumentButton: (documentId: string, chunk: IChunk) => void;
 }) => {
   const userInfo = useSelectUserInfo();
   const fileThumbnails = useSelectFileThumbnails();
   const getDocumentUrl = useGetDocumentUrl();
+  const { t } = useTranslate('chat');
 
   const isAssistant = item.role === MessageType.Assistant;
 
   const referenceDocumentList = useMemo(() => {
     return reference?.doc_aggs ?? [];
   }, [reference?.doc_aggs]);
 
+  const content = useMemo(() => {
+    let text = item.content;
+    if (text === '') {
+      text = t('searching');
+    }
+    return loading ? text?.concat('~~2$$') : text;
+  }, [item.content, loading, t]);
+
   return (
     <div
       className={classNames(styles.messageItem, {
@@ -85,15 +88,11 @@ const MessageItem = ({
           <Flex vertical gap={8} flex={1}>
             <b>{isAssistant ? '' : userInfo.nickname}</b>
             <div className={styles.messageText}>
-              {item.content !== '' ? (
-                <MarkdownContent
-                  content={item.content}
-                  reference={reference}
-                  clickDocumentButton={clickDocumentButton}
-                ></MarkdownContent>
-              ) : (
-                <Skeleton active className={styles.messageEmpty} />
-              )}
+              <MarkdownContent
+                content={content}
+                reference={reference}
+                clickDocumentButton={clickDocumentButton}
+              ></MarkdownContent>
             </div>
             {isAssistant && referenceDocumentList.length > 0 && (
               <List
@@ -139,13 +138,19 @@ const ChatContainer = () => {
     currentConversation: conversation,
     addNewestConversation,
     removeLatestMessage,
+    addNewestAnswer,
   } = useFetchConversationOnMount();
   const {
     handleInputChange,
     handlePressEnter,
     value,
     loading: sendLoading,
-  } = useSendMessage(conversation, addNewestConversation, removeLatestMessage);
+  } = useSendMessage(
+    conversation,
+    addNewestConversation,
+    removeLatestMessage,
+    addNewestAnswer,
+  );
   const { visible, hideModal, documentId, selectedChunk, clickDocumentButton } =
     useClickDrawer();
   const disabled = useGetSendButtonDisabled();
@@ -159,19 +164,17 @@ const ChatContainer = () => {
         <Flex flex={1} vertical className={styles.messageContainer}>
           <div>
             <Spin spinning={loading}>
-              {conversation?.message?.map((message) => {
-                const assistantMessages = conversation?.message
-                  ?.filter((x) => x.role === MessageType.Assistant)
-                  .slice(1);
-                const referenceIndex = assistantMessages.findIndex(
-                  (x) => x.id === message.id,
-                );
-                const reference = conversation.reference[referenceIndex];
+              {conversation?.message?.map((message, i) => {
                 return (
                   <MessageItem
+                    loading={
+                      message.role === MessageType.Assistant &&
+                      sendLoading &&
+                      conversation?.message.length - 1 === i
+                    }
                     key={message.id}
                     item={message}
-                    reference={reference}
+                    reference={buildMessageItemReference(conversation, message)}
                     clickDocumentButton={clickDocumentButton}
                   ></MessageItem>
                 );