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]) => {
const queryFn = useMemoizedFn(async () => {
return await getChat(params).then((chat) => {
console.log('TODO move this to put message in a better spot');
return updateChatToIChat(chat, true).iChat;
});
});
@ -67,6 +68,7 @@ export const prefetchGetChat = async (
...queryKeys.chatsGetChat(params.id),
queryFn: async () => {
return await getChat_server(params).then((chat) => {
console.log('TODO move this to put message in a better spot');
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,
buttons
}) => {
const showLoader = status === 'loading';
const showLoader = status === 'loading' && !isCompletedStream;
const { text = '', modified } = file;
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 sample from 'lodash/sample';
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 getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized);
const removeAutoThoughts = useMemoizedFn(
(reasoningMessages: BusterChatMessageReasoning[]): BusterChatMessageReasoning[] => {
return reasoningMessages.filter((rm) => rm.id !== AUTO_THOUGHT_ID);
}
);
const removeAutoThoughts = useMemoizedFn(() => {
// return reasoningMessages.filter((rm) => rm.id !== AUTO_THOUGHT_ID);
});
const autoAppendThought = useMemoizedFn(
(
reasoningMessages: BusterChatMessageReasoning[],
chatId: string
): BusterChatMessageReasoning[] => {
const lastReasoningMessage = reasoningMessages[reasoningMessages.length - 1];
const lastMessageIsCompleted =
!lastReasoningMessage || lastReasoningMessage?.status === 'completed';
const autoAppendThought = useMemoizedFn(({ messageId }: { messageId: string }) => {
//
// const lastReasoningMessage = reasoningMessages[reasoningMessages.length - 1];
// const lastMessageIsCompleted =
// !lastReasoningMessage || lastReasoningMessage?.status === 'completed';
if (lastMessageIsCompleted) {
_loopAutoThought(chatId);
// if (lastMessageIsCompleted) {
// _loopAutoThought(chatId);
return [...reasoningMessages, createAutoThought()];
}
// return [...reasoningMessages, createAutoThought()];
// }
return removeAutoThoughts(reasoningMessages);
}
);
return removeAutoThoughts();
});
const _loopAutoThought = useMemoizedFn(async (chatId: string) => {
// const randomDelay = random(3000, 5000);
// await timeout(randomDelay);
const randomDelay = random(3000, 5000);
await timeout(randomDelay);
// const chatMessages = getChatMessagesMemoized(chatId);
// const lastMessage = last(chatMessages);
// const isCompletedStream = !!lastMessage?.isCompletedStream;
@ -56,7 +56,6 @@ export const useAutoAppendThought = () => {
// isCompletedStream: false
// });
// _loopAutoThought(chatId);
// }
});
return { autoAppendThought, removeAutoThoughts };
@ -69,18 +68,6 @@ const getRandomThought = (currentThought?: string): string => {
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 = [
'Thinking through next steps...',
'Looking through context...',

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

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

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

@ -1,6 +1,6 @@
import type { BusterChat, BusterChatMessage } from '@/api/asset_interfaces';
export interface IBusterChat extends BusterChat {
export interface IBusterChat extends Omit<BusterChat, 'messages'> {
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(() => {
if (status === 'loading')
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)'} />
</Text>
);
if (status === 'completed')
return (
<Text variant={'secondary'} className="flex gap-1.5">
<Text variant={'secondary'} size={'sm'} className="flex gap-1.5">
Completed <CheckDouble />
</Text>
);
if (status === 'failed')
return (
<Text variant={'danger'} className="flex gap-1.5">
<Text variant={'danger'} size={'sm'} className="flex gap-1.5">
Failed <AlertWarning />
</Text>
);

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

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