Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

chat stream message

This commit is contained in:
Nate Kelley 2025-03-06 09:35:15 -07:00
parent 998b927ba7
commit 0728efe71a
No known key found for this signature in database
GPG Key ID: FD90372AB8D98B4F
7 changed files with 94 additions and 136 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
web/src/api/buster_rest/chats/queryRequests.ts
View File

@ -40,6 +40,7 @@ export const prefetchGetListChats = async (
export const useGetChat = (params: Parameters<typeof getChat>[0]) => { export const useGetChat = (params: Parameters<typeof getChat>[0]) => {
const queryFn = useMemoizedFn(async () => { const queryFn = useMemoizedFn(async () => {
return await getChat(params).then((chat) => { return await getChat(params).then((chat) => {
console.log('TODO move this to put message in a better spot');
return updateChatToIChat(chat, true).iChat; return updateChatToIChat(chat, true).iChat;
}); });
}); });
@ -67,6 +68,7 @@ export const prefetchGetChat = async (
...queryKeys.chatsGetChat(params.id), ...queryKeys.chatsGetChat(params.id),
queryFn: async () => { queryFn: async () => {
return await getChat_server(params).then((chat) => { return await getChat_server(params).then((chat) => {
console.log('TODO move this to put message in a better spot');
return updateChatToIChat(chat, true).iChat; return updateChatToIChat(chat, true).iChat;
}); });
} }

2
web/src/components/ui/streaming/StreamingMessageCode/StreamingMessageCode.tsx
View File

@ -51,7 +51,7 @@ export const StreamingMessageCode: React.FC<
version_id, version_id,
buttons buttons
}) => { }) => {
const showLoader = status === 'loading'; const showLoader = status === 'loading' && !isCompletedStream;
const { text = '', modified } = file; const { text = '', modified } = file;
const [lineSegments, setLineSegments] = useState<LineSegment[]>([]); const [lineSegments, setLineSegments] = useState<LineSegment[]>([]);

59
web/src/context/Chats/NewChatProvider/useAutoAppendThought.ts → web/src/context/Chats/NewChatProvider/useBlackBoxMessage.ts
View File

@ -5,39 +5,39 @@ import type {
import { useMemoizedFn } from 'ahooks'; import { useMemoizedFn } from 'ahooks';
import sample from 'lodash/sample'; import sample from 'lodash/sample';
import { useBusterChatContextSelector } from '../ChatProvider'; import { useBusterChatContextSelector } from '../ChatProvider';
import random from 'lodash/random';
import last from 'lodash/last';
import { timeout } from '@/lib/timeout';
import { useState } from 'react';
export const useBlackBoxMessage = () => {
const [boxBoxMessages, setBoxBoxMessages] = useState<Record<string, string>>({});
export const useAutoAppendThought = () => {
const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage); const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage);
const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized); const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized);
const removeAutoThoughts = useMemoizedFn( const removeAutoThoughts = useMemoizedFn(() => {
(reasoningMessages: BusterChatMessageReasoning[]): BusterChatMessageReasoning[] => { // return reasoningMessages.filter((rm) => rm.id !== AUTO_THOUGHT_ID);
return reasoningMessages.filter((rm) => rm.id !== AUTO_THOUGHT_ID); });
}
);
const autoAppendThought = useMemoizedFn( const autoAppendThought = useMemoizedFn(({ messageId }: { messageId: string }) => {
( //
reasoningMessages: BusterChatMessageReasoning[], // const lastReasoningMessage = reasoningMessages[reasoningMessages.length - 1];
chatId: string // const lastMessageIsCompleted =
): BusterChatMessageReasoning[] => { // !lastReasoningMessage || lastReasoningMessage?.status === 'completed';
const lastReasoningMessage = reasoningMessages[reasoningMessages.length - 1];
const lastMessageIsCompleted =
!lastReasoningMessage || lastReasoningMessage?.status === 'completed';
if (lastMessageIsCompleted) { // if (lastMessageIsCompleted) {
_loopAutoThought(chatId); // _loopAutoThought(chatId);
return [...reasoningMessages, createAutoThought()]; // return [...reasoningMessages, createAutoThought()];
} // }
return removeAutoThoughts(reasoningMessages); return removeAutoThoughts();
} });
);
const _loopAutoThought = useMemoizedFn(async (chatId: string) => { const _loopAutoThought = useMemoizedFn(async (chatId: string) => {
// const randomDelay = random(3000, 5000); const randomDelay = random(3000, 5000);
// await timeout(randomDelay); await timeout(randomDelay);
// const chatMessages = getChatMessagesMemoized(chatId); // const chatMessages = getChatMessagesMemoized(chatId);
// const lastMessage = last(chatMessages); // const lastMessage = last(chatMessages);
// const isCompletedStream = !!lastMessage?.isCompletedStream; // const isCompletedStream = !!lastMessage?.isCompletedStream;
@ -56,7 +56,6 @@ export const useAutoAppendThought = () => {
// isCompletedStream: false // isCompletedStream: false
// }); // });
// _loopAutoThought(chatId); // _loopAutoThought(chatId);
// }
}); });
return { autoAppendThought, removeAutoThoughts }; return { autoAppendThought, removeAutoThoughts };
@ -69,18 +68,6 @@ const getRandomThought = (currentThought?: string): string => {
return sample(thoughts) ?? DEFAULT_THOUGHTS[0]; return sample(thoughts) ?? DEFAULT_THOUGHTS[0];
}; };
const AUTO_THOUGHT_ID = 'stub-thought-id';
const createAutoThought = (currentThought?: string): BusterChatMessageReasoning_pills => {
return {
id: AUTO_THOUGHT_ID,
type: 'pills',
title: getRandomThought(currentThought),
secondary_title: '',
pill_containers: [],
status: 'loading'
};
};
const DEFAULT_THOUGHTS = [ const DEFAULT_THOUGHTS = [
'Thinking through next steps...', 'Thinking through next steps...',
'Looking through context...', 'Looking through context...',

156
web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
View File

@ -4,9 +4,7 @@ import type {
BusterChat, BusterChat,
BusterChatMessageReasoning_files, BusterChatMessageReasoning_files,
BusterChatMessageReasoning_text, BusterChatMessageReasoning_text,
BusterChatMessageReasoning_pills,
BusterChatResponseMessage_text, BusterChatResponseMessage_text,
BusterChatMessageResponse,
BusterChatMessageReasoning_file BusterChatMessageReasoning_file
} from '@/api/asset_interfaces'; } from '@/api/asset_interfaces';
import type { import type {
@ -15,7 +13,7 @@ import type {
ChatEvent_GeneratingTitle ChatEvent_GeneratingTitle
} from '@/api/buster_socket/chats'; } from '@/api/buster_socket/chats';
import { updateChatToIChat } from '@/lib/chat'; import { updateChatToIChat } from '@/lib/chat';
import { useAutoAppendThought } from './useAutoAppendThought'; import { useBlackBoxMessage } from './useBlackBoxMessage';
import { useAppLayoutContextSelector } from '@/context/BusterAppLayout'; import { useAppLayoutContextSelector } from '@/context/BusterAppLayout';
import { BusterRoutes } from '@/routes'; import { BusterRoutes } from '@/routes';
import { useSocketQueryOn } from '@/api/buster_socket_query'; import { useSocketQueryOn } from '@/api/buster_socket_query';
@ -31,13 +29,16 @@ export const useChatStreamMessage = () => {
const onChangePage = useAppLayoutContextSelector((x) => x.onChangePage); const onChangePage = useAppLayoutContextSelector((x) => x.onChangePage);
const onUpdateChat = useBusterChatContextSelector((x) => x.onUpdateChat); const onUpdateChat = useBusterChatContextSelector((x) => x.onUpdateChat);
const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage); const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage);
const chatRef = useRef<Record<string, Partial<IBusterChat>>>({}); const chatRef = useRef<Record<string, IBusterChat>>({});
const chatRefMessages = useRef<Record<string, IBusterChatMessage>>({});
const [isPending, startTransition] = useTransition(); const [isPending, startTransition] = useTransition();
const { autoAppendThought } = useBlackBoxMessage();
const onUpdateChatMessageTransition = useMemoizedFn( const onUpdateChatMessageTransition = useMemoizedFn(
(chatMessage: Parameters<typeof onUpdateChatMessage>[0], chatId: string) => { (chatMessage: Parameters<typeof onUpdateChatMessage>[0]) => {
const currentChatMessage = chatRef.current[chatId]?.messages?.[chatMessage.id]; const currentChatMessage = chatRefMessages.current[chatMessage.id];
const iChatMessage = create(currentChatMessage, (draft) => { const iChatMessage: IBusterChatMessage = create(currentChatMessage, (draft) => {
Object.assign(draft || {}, chatMessage); Object.assign(draft || {}, chatMessage);
})!; })!;
@ -49,32 +50,14 @@ export const useChatStreamMessage = () => {
} }
); );
const { autoAppendThought } = useAutoAppendThought();
const initializeOrUpdateMessage = useMemoizedFn( const initializeOrUpdateMessage = useMemoizedFn(
( (messageId: string, updateFn: (draft: IBusterChatMessage) => void) => {
chatId: string, const currentMessage = chatRefMessages.current[messageId];
messageId: string, const updatedMessage = create(currentMessage || {}, (draft) => {
updateFn: (draft: Record<string, Partial<IBusterChat>>) => void
) => {
chatRef.current = create(chatRef.current, (draft) => {
if (!draft[chatId]) draft[chatId] = {};
if (!draft[chatId].messages) draft[chatId].messages = {};
if (!draft[chatId].messages[messageId]) {
draft[chatId].messages[messageId] = {
id: messageId,
request_message: null,
response_message_ids: [],
response_messages: {},
reasoning_message_ids: [],
reasoning_messages: {},
created_at: new Date().toISOString(),
final_reasoning_message: null
};
}
updateFn(draft); updateFn(draft);
}); });
chatRefMessages.current[messageId] = updatedMessage;
onUpdateChatMessage(updatedMessage);
} }
); );
@ -84,6 +67,7 @@ export const useChatStreamMessage = () => {
const options = queryKeys.chatsMessages(message.id); const options = queryKeys.chatsMessages(message.id);
const queryKey = options.queryKey; const queryKey = options.queryKey;
queryClient.setQueryData(queryKey, message); queryClient.setQueryData(queryKey, message);
chatRefMessages.current[message.id] = message;
} }
} }
); );
@ -92,7 +76,6 @@ export const useChatStreamMessage = () => {
const { iChat, iChatMessages } = updateChatToIChat(d, false); const { iChat, iChatMessages } = updateChatToIChat(d, false);
chatRef.current = create(chatRef.current, (draft) => { chatRef.current = create(chatRef.current, (draft) => {
draft[iChat.id] = iChat; draft[iChat.id] = iChat;
draft[iChat.id].messages = iChatMessages;
}); });
normalizeChatMessage(iChatMessages); normalizeChatMessage(iChatMessages);
onUpdateChat(iChat); onUpdateChat(iChat);
@ -109,8 +92,8 @@ export const useChatStreamMessage = () => {
const { iChat, iChatMessages } = updateChatToIChat(d, true); const { iChat, iChatMessages } = updateChatToIChat(d, true);
chatRef.current = create(chatRef.current, (draft) => { chatRef.current = create(chatRef.current, (draft) => {
draft[iChat.id] = iChat; draft[iChat.id] = iChat;
draft[iChat.id].messages = iChatMessages;
}); });
normalizeChatMessage(iChatMessages); normalizeChatMessage(iChatMessages);
onUpdateChat(iChat); onUpdateChat(iChat);
onChangePage({ onChangePage({
@ -141,8 +124,7 @@ export const useChatStreamMessage = () => {
const currentTitle = chatRef.current[chat_id]?.title || ''; const currentTitle = chatRef.current[chat_id]?.title || '';
const newTitle = isCompleted ? title : currentTitle + title_chunk; const newTitle = isCompleted ? title : currentTitle + title_chunk;
chatRef.current = create(chatRef.current, (draft) => { chatRef.current = create(chatRef.current, (draft) => {
if (!draft[chat_id]) draft[chat_id] = {}; if (newTitle) draft[chat_id].title = newTitle;
draft[chat_id].title = newTitle;
}); });
onUpdateChat({ onUpdateChat({
id: chat_id, id: chat_id,
@ -158,28 +140,29 @@ export const useChatStreamMessage = () => {
const responseMessageId = response_message.id; const responseMessageId = response_message.id;
const existingMessage = const existingMessage =
chatRef.current[chat_id]?.messages?.[message_id]?.response_messages?.[responseMessageId]; chatRefMessages.current[message_id]?.response_messages?.[responseMessageId];
const isNewMessage = !existingMessage; const isNewMessage = !existingMessage;
if (isNewMessage) { if (isNewMessage) {
initializeOrUpdateMessage(chat_id, message_id, (draft) => { initializeOrUpdateMessage(message_id, (draft) => {
const chat = draft[chat_id]; if (!draft.response_messages) {
if (!chat?.messages?.[message_id]) return; draft.response_messages = {};
if (!chat.messages[message_id].response_messages) }
chat.messages[message_id].response_messages = {}; draft.response_messages[responseMessageId] = response_message;
chat.messages[message_id].response_messages[responseMessageId] = response_message; if (!draft.response_message_ids) {
chat.messages[message_id].response_message_ids.push(responseMessageId); draft.response_message_ids = [];
}
draft.response_message_ids.push(responseMessageId);
}); });
} }
if (response_message.type === 'text') { if (response_message.type === 'text') {
const existingResponseMessageText = existingMessage as BusterChatResponseMessage_text; const existingResponseMessageText = existingMessage as BusterChatResponseMessage_text;
const isStreaming = const isStreaming =
response_message.message_chunk !== undefined || response_message.message_chunk !== null; response_message.message_chunk !== undefined && response_message.message_chunk !== null;
initializeOrUpdateMessage(chat_id, message_id, (draft) => { initializeOrUpdateMessage(message_id, (draft) => {
const responseMessage = const responseMessage = draft.response_messages?.[responseMessageId];
draft[chat_id]?.messages?.[message_id]?.response_messages?.[responseMessageId];
if (!responseMessage) return; if (!responseMessage) return;
const messageText = responseMessage as BusterChatMessageReasoning_text; const messageText = responseMessage as BusterChatMessageReasoning_text;
Object.assign(messageText, { Object.assign(messageText, {
@ -193,18 +176,12 @@ export const useChatStreamMessage = () => {
}); });
} }
const response_messages = chatRef.current[chat_id]?.messages?.[message_id]?.response_messages; const currentMessage = chatRefMessages.current[message_id];
const response_message_ids = onUpdateChatMessageTransition({
chatRef.current[chat_id]?.messages?.[message_id]?.response_message_ids; id: message_id,
response_messages: currentMessage?.response_messages,
onUpdateChatMessageTransition( response_message_ids: currentMessage?.response_message_ids
{ });
id: message_id,
response_messages,
response_message_ids
},
chat_id
);
} }
); );
@ -214,17 +191,19 @@ export const useChatStreamMessage = () => {
const reasoningMessageId = reasoning.id; const reasoningMessageId = reasoning.id;
const existingMessage = const existingMessage =
chatRef.current[chat_id]?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId]; chatRefMessages.current[message_id]?.reasoning_messages?.[reasoningMessageId];
const isNewMessage = !existingMessage; const isNewMessage = !existingMessage;
if (isNewMessage) { if (isNewMessage) {
initializeOrUpdateMessage(chat_id, message_id, (draft) => { initializeOrUpdateMessage(message_id, (draft) => {
const chat = draft[chat_id]; if (!draft.reasoning_messages) {
if (!chat?.messages?.[message_id]) return; draft.reasoning_messages = {};
if (!chat.messages[message_id].reasoning_messages) }
chat.messages[message_id].reasoning_messages = {}; draft.reasoning_messages[reasoningMessageId] = reasoning;
chat.messages[message_id].reasoning_messages[reasoningMessageId] = reasoning; if (!draft.reasoning_message_ids) {
chat.messages[message_id].reasoning_message_ids.push(reasoningMessageId); draft.reasoning_message_ids = [];
}
draft.reasoning_message_ids.push(reasoningMessageId);
}); });
} }
@ -234,9 +213,8 @@ export const useChatStreamMessage = () => {
const isStreaming = const isStreaming =
reasoning.message_chunk !== null || reasoning.message_chunk !== undefined; reasoning.message_chunk !== null || reasoning.message_chunk !== undefined;
initializeOrUpdateMessage(chat_id, message_id, (draft) => { initializeOrUpdateMessage(message_id, (draft) => {
const reasoningMessage = const reasoningMessage = draft.reasoning_messages?.[reasoningMessageId];
draft[chat_id]?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId];
if (!reasoningMessage) return; if (!reasoningMessage) return;
const messageText = reasoningMessage as BusterChatMessageReasoning_text; const messageText = reasoningMessage as BusterChatMessageReasoning_text;
@ -254,14 +232,12 @@ export const useChatStreamMessage = () => {
case 'files': { case 'files': {
const existingReasoningMessageFiles = existingMessage as BusterChatMessageReasoning_files; const existingReasoningMessageFiles = existingMessage as BusterChatMessageReasoning_files;
initializeOrUpdateMessage(chat_id, message_id, (draft) => { initializeOrUpdateMessage(message_id, (draft) => {
const chat = draft[chat_id]; const reasoningMessage = draft.reasoning_messages?.[reasoningMessageId];
if (!chat?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId]) return; if (!reasoningMessage) return;
const messageFiles = create( const messageFiles = create(
chat.messages[message_id].reasoning_messages[ reasoningMessage as BusterChatMessageReasoning_files,
reasoningMessageId
] as BusterChatMessageReasoning_files,
(draft) => { (draft) => {
draft.file_ids = existingReasoningMessageFiles?.file_ids || []; draft.file_ids = existingReasoningMessageFiles?.file_ids || [];
@ -305,15 +281,14 @@ export const useChatStreamMessage = () => {
} }
); );
chat.messages[message_id].reasoning_messages[reasoningMessageId] = messageFiles; draft.reasoning_messages[reasoningMessageId] = messageFiles;
}); });
break; break;
} }
case 'pills': { case 'pills': {
initializeOrUpdateMessage(chat_id, message_id, (draft) => { initializeOrUpdateMessage(message_id, (draft) => {
if (!draft[chat_id]?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId]) if (!draft.reasoning_messages?.[reasoningMessageId]) return;
return; draft.reasoning_messages[reasoningMessageId] = reasoning;
draft[chat_id].messages[message_id].reasoning_messages[reasoningMessageId]! = reasoning;
}); });
break; break;
@ -324,20 +299,13 @@ export const useChatStreamMessage = () => {
} }
} }
const reasoning_messages = const currentMessage = chatRefMessages.current[message_id];
chatRef.current[chat_id]?.messages?.[message_id]?.reasoning_messages; onUpdateChatMessageTransition({
const reasoning_message_ids = id: message_id,
chatRef.current[chat_id]?.messages?.[message_id]?.reasoning_message_ids; reasoning_messages: currentMessage?.reasoning_messages,
reasoning_message_ids: currentMessage?.reasoning_message_ids,
onUpdateChatMessageTransition( isCompletedStream: false
{ });
id: message_id,
reasoning_messages,
reasoning_message_ids,
isCompletedStream: false
},
chat_id
);
} }
); );

2
web/src/context/Chats/interfaces.ts
View File

@ -1,6 +1,6 @@
import type { BusterChat, BusterChatMessage } from '@/api/asset_interfaces'; import type { BusterChat, BusterChatMessage } from '@/api/asset_interfaces';
export interface IBusterChat extends BusterChat { export interface IBusterChat extends Omit<BusterChat, 'messages'> {
isNewChat: boolean; isNewChat: boolean;
} }

6
web/src/controllers/ReasoningController/ReasoningMessages/ReasoningMessage_Files/ReasoningMessageFile.tsx
View File

@ -52,19 +52,19 @@ const StreamingMessageStatus = React.memo(
const content = useMemo(() => { const content = useMemo(() => {
if (status === 'loading') if (status === 'loading')
return ( return (
<Text variant={'secondary'} className="flex gap-1.5"> <Text variant={'secondary'} size={'sm'} className="flex gap-1.5">
Running SQL... <CircleSpinnerLoader size={9} fill={'var(--color-text-secondary)'} /> Running SQL... <CircleSpinnerLoader size={9} fill={'var(--color-text-secondary)'} />
</Text> </Text>
); );
if (status === 'completed') if (status === 'completed')
return ( return (
<Text variant={'secondary'} className="flex gap-1.5"> <Text variant={'secondary'} size={'sm'} className="flex gap-1.5">
Completed <CheckDouble /> Completed <CheckDouble />
</Text> </Text>
); );
if (status === 'failed') if (status === 'failed')
return ( return (
<Text variant={'danger'} className="flex gap-1.5"> <Text variant={'danger'} size={'sm'} className="flex gap-1.5">
Failed <AlertWarning /> Failed <AlertWarning />
</Text> </Text>
); );

3
web/src/lib/chat.ts
View File

@ -1,10 +1,11 @@
import type { BusterChat, BusterChatMessage } from '@/api/asset_interfaces'; import type { BusterChat, BusterChatMessage } from '@/api/asset_interfaces';
import type { IBusterChat, IBusterChatMessage } from '@/context/Chats/interfaces'; import type { IBusterChat, IBusterChatMessage } from '@/context/Chats/interfaces';
import { create } from 'mutative'; import { create } from 'mutative';
import omit from 'lodash/omit';
const chatUpgrader = (chat: BusterChat, { isNewChat }: { isNewChat: boolean }): IBusterChat => { const chatUpgrader = (chat: BusterChat, { isNewChat }: { isNewChat: boolean }): IBusterChat => {
return { return {
...chat, ...omit(chat, 'messages'),
isNewChat isNewChat
}; };
}; };