Skip to content

Commit

Permalink
feat: Support for conversational streaming (infiniflow#809)
Browse files Browse the repository at this point in the history
### 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)
  • Loading branch information
cike8899 authored May 16, 2024
1 parent 95f8091 commit c6c9dbd
Show file tree
Hide file tree
Showing 21 changed files with 426 additions and 257 deletions.
1 change: 1 addition & 0 deletions web/.env
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
PORT=9222
9 changes: 9 additions & 0 deletions web/package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion web/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand Down
2 changes: 1 addition & 1 deletion web/src/components/new-document-link.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
19 changes: 3 additions & 16 deletions web/src/hooks/chatHooks.ts
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,9 @@ export const useRemoveConversation = () => {
return removeConversation;
};

/*
@deprecated
*/
export const useCompleteConversation = () => {
const dispatch = useDispatch();

Expand Down Expand Up @@ -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
105 changes: 52 additions & 53 deletions web/src/hooks/logicHooks.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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 };
};
5 changes: 5 additions & 0 deletions web/src/interfaces/database/chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,11 @@ export interface IReference {
total: number;
}

export interface IAnswer {
answer: string;
reference: IReference;
}

export interface Docagg {
count: number;
doc_id: string;
Expand Down
2 changes: 2 additions & 0 deletions web/src/locales/en.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ export default {
comingSoon: 'Coming Soon',
download: 'Download',
close: 'Close',
preview: 'Preview',
},
login: {
login: 'Sign in',
Expand Down Expand Up @@ -381,6 +382,7 @@ export default {
partialTitle: 'Partial Embed',
extensionTitle: 'Chrome Extension',
tokenError: 'Please create API Token first!',
searching: 'searching...',
},
setting: {
profile: 'Profile',
Expand Down
2 changes: 2 additions & 0 deletions web/src/locales/zh-traditional.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ export default {
comingSoon: '即將推出',
download: '下載',
close: '关闭',
preview: '預覽',
},
login: {
login: '登入',
Expand Down Expand Up @@ -352,6 +353,7 @@ export default {
partialTitle: '部分嵌入',
extensionTitle: 'Chrome 插件',
tokenError: '請先創建 Api Token!',
searching: '搜索中',
},
setting: {
profile: '概述',
Expand Down
2 changes: 2 additions & 0 deletions web/src/locales/zh.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ export default {
comingSoon: '即将推出',
download: '下载',
close: '关闭',
preview: '预览',
},
login: {
login: '登录',
Expand Down Expand Up @@ -369,6 +370,7 @@ export default {
partialTitle: '部分嵌入',
extensionTitle: 'Chrome 插件',
tokenError: '请先创建 Api Token!',
searching: '搜索中',
},
setting: {
profile: '概要',
Expand Down
61 changes: 32 additions & 29 deletions web/src/pages/chat/chat-container/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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, {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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();
Expand All @@ -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>
);
Expand Down
Loading

0 comments on commit c6c9dbd

Please sign in to comment.