Feat: Interrupt streaming #6515 (#6723)

### What problem does this PR solve?

Feat: Interrupt streaming #6515
### Type of change


- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
balibabu 2025-04-01 17:26:54 +08:00 committed by GitHub
parent ead5f7aba9
commit 132eae9d5b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 86 additions and 18 deletions

View File

@ -29,6 +29,7 @@ import {
UploadProps,
} from 'antd';
import get from 'lodash/get';
import { CircleStop } from 'lucide-react';
import {
ChangeEventHandler,
memo,
@ -72,6 +73,7 @@ interface IProps {
isShared?: boolean;
showUploadIcon?: boolean;
createConversationBeforeUploadDocument?(message: string): Promise<any>;
stopOutputMessage?(): void;
}
const getBase64 = (file: FileType): Promise<string> =>
@ -94,6 +96,7 @@ const MessageInput = ({
showUploadIcon = true,
createConversationBeforeUploadDocument,
uploadMethod = 'upload_and_parse',
stopOutputMessage,
}: IProps) => {
const { t } = useTranslate('chat');
const { removeDocument } = useRemoveNextDocument();
@ -160,7 +163,7 @@ const MessageInput = ({
event.preventDefault();
handlePressEnter();
},
[fileList, onPressEnter, isUploadingFile],
[sendDisabled, isUploadingFile, sendLoading, handlePressEnter],
);
const handlePressEnter = useCallback(async () => {
@ -199,6 +202,10 @@ const MessageInput = ({
[removeDocument, deleteDocument, isShared],
);
const handleStopOutputMessage = useCallback(() => {
stopOutputMessage?.();
}, [stopOutputMessage]);
const getDocumentInfoById = useCallback(
(id: string) => {
return documentInfos.find((x) => x.id === id);
@ -346,14 +353,20 @@ const MessageInput = ({
</Button>
</Upload>
)}
<Button
type="primary"
onClick={handlePressEnter}
loading={sendLoading}
disabled={sendDisabled || isUploadingFile || sendLoading}
>
<SendOutlined />
</Button>
{sendLoading ? (
<Button onClick={handleStopOutputMessage}>
<CircleStop />
</Button>
) : (
<Button
type="primary"
onClick={handlePressEnter}
loading={sendLoading}
disabled={sendDisabled || isUploadingFile || sendLoading}
>
<SendOutlined />
</Button>
)}
</Flex>
</Flex>
</Flex>

View File

@ -160,6 +160,11 @@ export const useSendMessageWithSse = (
const [answer, setAnswer] = useState<IAnswer>({} as IAnswer);
const [done, setDone] = useState(true);
const timer = useRef<any>();
const sseRef = useRef<AbortController>();
const initializeSseRef = useCallback(() => {
sseRef.current = new AbortController();
}, []);
const resetAnswer = useCallback(() => {
if (timer.current) {
@ -176,6 +181,7 @@ export const useSendMessageWithSse = (
body: any,
controller?: AbortController,
): Promise<{ response: Response; data: ResponseType } | undefined> => {
initializeSseRef();
try {
setDone(false);
const response = await fetch(url, {
@ -185,7 +191,7 @@ export const useSendMessageWithSse = (
'Content-Type': 'application/json',
},
body: JSON.stringify(body),
signal: controller?.signal,
signal: controller?.signal || sseRef.current?.signal,
});
const res = response.clone().json();
@ -230,10 +236,14 @@ export const useSendMessageWithSse = (
console.warn(e);
}
},
[url, resetAnswer],
[initializeSseRef, url, resetAnswer],
);
return { send, answer, done, setDone, resetAnswer };
const stopOutputMessage = useCallback(() => {
sseRef.current?.abort();
}, []);
return { send, answer, done, setDone, resetAnswer, stopOutputMessage };
};
export const useSpeechWithSse = (url: string = api.tts) => {

View File

@ -40,6 +40,7 @@ const ChatContainer = ({ controller }: IProps) => {
handlePressEnter,
regenerateMessage,
removeMessageById,
stopOutputMessage,
} = useSendNextMessage(controller);
const { visible, hideModal, documentId, selectedChunk, clickDocumentButton } =
@ -100,6 +101,7 @@ const ChatContainer = ({ controller }: IProps) => {
createConversationBeforeUploadDocument={
createConversationBeforeUploadDocument
}
stopOutputMessage={stopOutputMessage}
></MessageInput>
</Flex>
<PdfDrawer

View File

@ -375,6 +375,10 @@ export const useSendNextMessage = (controller: AbortController) => {
const { setConversationIsNew, getConversationIsNew } =
useSetChatRouteParams();
const stopOutputMessage = useCallback(() => {
controller.abort();
}, [controller]);
const sendMessage = useCallback(
async ({
message,
@ -490,6 +494,7 @@ export const useSendNextMessage = (controller: AbortController) => {
ref,
derivedMessages,
removeMessageById,
stopOutputMessage,
};
};

View File

@ -37,6 +37,7 @@ const ChatContainer = () => {
ref,
derivedMessages,
hasError,
stopOutputMessage,
} = useSendSharedMessage();
const sendDisabled = useSendButtonDisabled(value);
@ -105,6 +106,7 @@ const ChatContainer = () => {
sendLoading={sendLoading}
uploadMethod="external_upload_and_parse"
showUploadIcon={false}
stopOutputMessage={stopOutputMessage}
></MessageInput>
</Flex>
{visible && (

View File

@ -49,7 +49,7 @@ export const useSendSharedMessage = () => {
const { createSharedConversation: setConversation } =
useCreateNextSharedConversation();
const { handleInputChange, value, setValue } = useHandleMessageInputChange();
const { send, answer, done } = useSendMessageWithSse(
const { send, answer, done, stopOutputMessage } = useSendMessageWithSse(
`/api/v1/${from === SharedFrom.Agent ? 'agentbots' : 'chatbots'}/${conversationId}/completions`,
);
const {
@ -144,5 +144,6 @@ export const useSendSharedMessage = () => {
loading: false,
derivedMessages,
hasError,
stopOutputMessage,
};
};

View File

@ -24,6 +24,7 @@ const FlowChatBox = () => {
ref,
derivedMessages,
reference,
stopOutputMessage,
} = useSendNextMessage();
const { visible, hideModal, documentId, selectedChunk, clickDocumentButton } =
@ -75,6 +76,7 @@ const FlowChatBox = () => {
conversationId=""
onPressEnter={handlePressEnter}
onInputChange={handleInputChange}
stopOutputMessage={stopOutputMessage}
/>
</Flex>
<PdfDrawer

View File

@ -57,7 +57,9 @@ export const useSendNextMessage = () => {
const { handleInputChange, value, setValue } = useHandleMessageInputChange();
const { refetch } = useFetchFlow();
const { send, answer, done } = useSendMessageWithSse(api.runCanvas);
const { send, answer, done, stopOutputMessage } = useSendMessageWithSse(
api.runCanvas,
);
const sendMessage = useCallback(
async ({ message }: { message: Message; messages?: Message[] }) => {
@ -134,5 +136,6 @@ export const useSendNextMessage = () => {
derivedMessages,
ref,
removeMessageById,
stopOutputMessage,
};
};

View File

@ -17,7 +17,9 @@ import {
} from 'react';
export const useSendQuestion = (kbIds: string[]) => {
const { send, answer, done } = useSendMessageWithSse(api.ask);
const { send, answer, done, stopOutputMessage } = useSendMessageWithSse(
api.ask,
);
const { testChunk, loading } = useTestChunkRetrieval();
const [sendingLoading, setSendingLoading] = useState(false);
const [currentAnswer, setCurrentAnswer] = useState({} as IAnswer);
@ -116,6 +118,7 @@ export const useSendQuestion = (kbIds: string[]) => {
isFirstRender,
selectedDocumentIds,
isSearchStrEmpty: isEmpty(trim(searchStr)),
stopOutputMessage,
};
};

View File

@ -137,6 +137,12 @@
.input();
}
.searchInput {
:global(.ant-input-search-button) {
display: none;
}
}
.appIcon {
display: inline-block;
vertical-align: middle;

View File

@ -12,6 +12,7 @@ import {
import { useGetPaginationWithRouter } from '@/hooks/logic-hooks';
import { IReference } from '@/interfaces/database/chat';
import {
Button,
Card,
Divider,
Flex,
@ -28,9 +29,11 @@ import {
Tag,
Tooltip,
} from 'antd';
import classNames from 'classnames';
import DOMPurify from 'dompurify';
import { isEmpty } from 'lodash';
import { useMemo, useState } from 'react';
import { CircleStop, SendHorizontal } from 'lucide-react';
import { useCallback, useMemo, useState } from 'react';
import { useTranslation } from 'react-i18next';
import MarkdownContent from '../chat/markdown-content';
import { useSendQuestion, useShowMindMapDrawer } from './hooks';
@ -64,6 +67,7 @@ const SearchPage = () => {
isFirstRender,
selectedDocumentIds,
isSearchStrEmpty,
stopOutputMessage,
} = useSendQuestion(checkedWithoutEmbeddingIdList);
const { visible, hideModal, documentId, selectedChunk, clickDocumentButton } =
useClickDrawer();
@ -81,18 +85,35 @@ const SearchPage = () => {
handleTestChunk(selectedDocumentIds, pageNumber, pageSize);
};
const handleSearch = useCallback(() => {
sendQuestion(searchStr);
}, [searchStr, sendQuestion]);
const InputSearch = (
<Search
value={searchStr}
onChange={handleSearchStrChange}
placeholder={t('header.search')}
allowClear
enterButton
addonAfter={
sendingLoading ? (
<Button onClick={stopOutputMessage}>
<CircleStop />
</Button>
) : (
<Button onClick={handleSearch}>
<SendHorizontal className="size-5 text-blue-500" />
</Button>
)
}
onSearch={sendQuestion}
size="large"
loading={sendingLoading}
disabled={checkedWithoutEmbeddingIdList.length === 0}
className={isFirstRender ? styles.globalInput : styles.partialInput}
className={classNames(
styles.searchInput,
isFirstRender ? styles.globalInput : styles.partialInput,
)}
/>
);