DB-GPT/web/hooks/use-chat.ts
Dreammy23 471689ba20
feat(web): Unified frontend code style (#1923)
Co-authored-by: Fangyin Cheng <staneyffer@gmail.com>
Co-authored-by: 谨欣 <echo.cmy@antgroup.com>
Co-authored-by: 严志勇 <yanzhiyong@tiansuixiansheng.com>
Co-authored-by: yanzhiyong <932374019@qq.com>
2024-08-30 14:03:06 +08:00

113 lines
3.4 KiB
TypeScript

import { ChatContext } from '@/app/chat-context';
import i18n from '@/app/i18n';
import { getUserId } from '@/utils';
import { HEADER_USER_ID_KEY } from '@/utils/constants/index';
import { EventStreamContentType, fetchEventSource } from '@microsoft/fetch-event-source';
import { message } from 'antd';
import { useCallback, useContext, useState } from 'react';
type Props = {
queryAgentURL?: string;
app_code?: string;
};
type ChatParams = {
chatId: string;
ctrl?: AbortController;
data?: any;
query?: Record<string, string>;
onMessage: (message: string) => void;
onClose?: () => void;
onDone?: () => void;
onError?: (content: string, error?: Error) => void;
};
const useChat = ({ queryAgentURL = '/api/v1/chat/completions', app_code }: Props) => {
const [ctrl, setCtrl] = useState<AbortController>({} as AbortController);
const { scene } = useContext(ChatContext);
const chat = useCallback(
async ({ data, chatId, onMessage, onClose, onDone, onError, ctrl }: ChatParams) => {
ctrl && setCtrl(ctrl);
if (!data?.user_input && !data?.doc_id) {
message.warning(i18n.t('no_context_tip'));
return;
}
const params = {
...data,
conv_uid: chatId,
app_code,
};
// if (!params.conv_uid) {
// message.error('conv_uid 不存在,请刷新后重试');
// return;
// }
try {
await fetchEventSource(`${process.env.API_BASE_URL ?? ''}${queryAgentURL}`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
[HEADER_USER_ID_KEY]: getUserId() ?? '',
},
body: JSON.stringify(params),
signal: ctrl ? ctrl.signal : null,
openWhenHidden: true,
async onopen(response) {
if (response.ok && response.headers.get('content-type') === EventStreamContentType) {
return;
}
if (response.headers.get('content-type') === 'application/json') {
response.json().then(data => {
onMessage?.(data);
onDone?.();
ctrl && ctrl.abort();
});
}
},
onclose() {
ctrl && ctrl.abort();
onClose?.();
},
onerror(err) {
throw new Error(err);
},
onmessage: event => {
let message = event.data;
try {
if (scene === 'chat_agent') {
message = JSON.parse(message).vis;
} else {
message = JSON.parse(message);
}
} catch {
message.replaceAll('\\n', '\n');
}
if (typeof message === 'string') {
if (message === '[DONE]') {
onDone?.();
} else if (message?.startsWith('[ERROR]')) {
onError?.(message?.replace('[ERROR]', ''));
} else {
onMessage?.(message);
}
} else {
onMessage?.(message);
onDone?.();
}
},
});
} catch (err) {
ctrl && ctrl.abort();
onError?.('Sorry, We meet some error, please try agin later.', err as Error);
}
},
[queryAgentURL, app_code, scene],
);
return { chat, ctrl };
};
export default useChat;