From 3fedad086640355e21c987728f2672c82823a295 Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Thu, 6 Mar 2025 11:17:53 -0700 Subject: [PATCH] better handling for black box messages --- .../chatStreamMessageHelper.test.ts | 36 +++++++++++++++---- .../NewChatProvider/useBlackBoxMessage.ts | 28 +++++++++------ .../NewChatProvider/useChatStreamMessage.ts | 16 +++++---- .../ReasoningBlackBoxMessage.tsx | 2 +- 4 files changed, 58 insertions(+), 24 deletions(-) diff --git a/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts index 07043015d..6fa63b75f 100644 --- a/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts +++ b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts @@ -522,7 +522,8 @@ describe('updateReasoningMessage', () => { }); it('should handle multiple updates and append text to existing reasoning message', () => { - const reasoning: BusterChatMessageReasoning_text = { + const messageId = 'test-message-id'; + let reasoning: BusterChatMessageReasoning_text = { id: 'reasoning-1', type: 'text', message: '', @@ -533,23 +534,46 @@ describe('updateReasoningMessage', () => { }; // First update with "Hello" - let result = updateReasoningMessage('test-message-id', undefined, reasoning); + let result = updateReasoningMessage(messageId, undefined, reasoning); expect( (result.reasoning_messages['reasoning-1'] as BusterChatMessageReasoning_text).message ).toBe('Hello'); // Second update with ", how" - reasoning.message_chunk = ', how'; - result = updateReasoningMessage('test-message-id', result, reasoning); + reasoning = { + ...reasoning, + message_chunk: ', how' + }; + result = updateReasoningMessage(messageId, result, reasoning); expect( (result.reasoning_messages['reasoning-1'] as BusterChatMessageReasoning_text).message ).toBe('Hello, how'); // Third update with " are you doing today?" - reasoning.message_chunk = ' are you doing today?'; - result = updateReasoningMessage('test-message-id', result, reasoning); + reasoning = { + ...reasoning, + message_chunk: ' are you doing today?' + }; + result = updateReasoningMessage(messageId, result, reasoning); expect( (result.reasoning_messages['reasoning-1'] as BusterChatMessageReasoning_text).message ).toBe('Hello, how are you doing today?'); + + let reasoning2: BusterChatMessageReasoning_text = { + ...reasoning, + message_chunk: 'new reasoning message baby', + id: 'reasoning-2' + }; + + result = updateReasoningMessage(messageId, result, reasoning2); + expect( + (result.reasoning_messages['reasoning-1'] as BusterChatMessageReasoning_text).message + ).toBe('Hello, how are you doing today?'); + + expect(result.reasoning_message_ids).toContain('reasoning-2'); + expect(result.reasoning_message_ids.length).toBe(2); + expect( + (result.reasoning_messages['reasoning-2'] as BusterChatMessageReasoning_text).message + ).toBe('new reasoning message baby'); }); }); diff --git a/web/src/context/Chats/NewChatProvider/useBlackBoxMessage.ts b/web/src/context/Chats/NewChatProvider/useBlackBoxMessage.ts index 4998401a6..0a0cb5e29 100644 --- a/web/src/context/Chats/NewChatProvider/useBlackBoxMessage.ts +++ b/web/src/context/Chats/NewChatProvider/useBlackBoxMessage.ts @@ -16,55 +16,63 @@ export const useBlackBoxMessage = () => { const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized); const queryClient = useQueryClient(); - const removeAutoThought = useMemoizedFn(({ messageId }: { messageId: string }) => { - console.log('removeAutoThought', messageId); + const clearTimeoutRef = useMemoizedFn((messageId: string) => { if (timeoutRef.current[messageId]) { clearTimeout(timeoutRef.current[messageId]); delete timeoutRef.current[messageId]; } + }); + + const removeBlackBoxMessage = useMemoizedFn(({ messageId }: { messageId: string }) => { + console.log('removeBlackBoxMessage', messageId); + clearTimeoutRef(messageId); const options = queryKeys.chatsBlackBoxMessages(messageId); queryClient.setQueryData(options.queryKey, null); }); - const addAutoThought = useMemoizedFn(({ messageId }: { messageId: string }) => { + const addBlackBoxMessage = useMemoizedFn(({ messageId }: { messageId: string }) => { const randomThought = getRandomThought(); console.log(messageId, randomThought); const options = queryKeys.chatsBlackBoxMessages(messageId); queryClient.setQueryData(options.queryKey, randomThought); }); - const checkAutoThought = useMemoizedFn( + const checkBlackBoxMessage = useMemoizedFn( (message: IBusterChatMessage, event: ChatEvent_GeneratingReasoningMessage) => { const isFinishedReasoningMessage = event.reasoning.status !== 'loading'; if (isFinishedReasoningMessage) { - addAutoThought({ messageId: message.id }); + clearTimeoutRef(message.id); + addBlackBoxMessage({ messageId: message.id }); _loopAutoThought({ messageId: message.id }); } else { - removeAutoThought({ messageId: message.id }); + removeBlackBoxMessage({ messageId: message.id }); } } ); const _loopAutoThought = useMemoizedFn(async ({ messageId }: { messageId: string }) => { - const randomDelay = random(3000, 5000); + const randomDelay = random(5000, 5000); timeoutRef.current[messageId] = setTimeout(() => { const message = getChatMessageMemoized(messageId); - console.log('loopAutoThought', messageId, !!message); if (!message) return; + if (!timeoutRef.current[messageId]) return; + + console.log('loopAutoThought', messageId, !!message); + const isMessageCompletedStream = !!message?.isCompletedStream; const lastReasoningMessageId = last(message?.reasoning_message_ids) || ''; const lastReasoningMessage = message?.reasoning_messages[lastReasoningMessageId]; const isLastReasoningMessageCompleted = lastReasoningMessage?.status === 'completed'; if (!isMessageCompletedStream && isLastReasoningMessageCompleted) { - addAutoThought({ messageId }); + addBlackBoxMessage({ messageId }); _loopAutoThought({ messageId }); } }, randomDelay); }); - return { checkAutoThought, removeAutoThought }; + return { checkBlackBoxMessage, removeBlackBoxMessage }; }; const getRandomThought = (currentThought?: string): string => { diff --git a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts index 17cf3495e..a8f20c6bd 100644 --- a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts +++ b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts @@ -32,7 +32,7 @@ export const useChatStreamMessage = () => { const chatRefMessages = useRef>({}); const [isPending, startTransition] = useTransition(); - const { checkAutoThought, removeAutoThought } = useBlackBoxMessage(); + const { checkBlackBoxMessage, removeBlackBoxMessage } = useBlackBoxMessage(); const onUpdateChatMessageTransition = useMemoizedFn( (chatMessage: Parameters[0]) => { @@ -40,6 +40,7 @@ export const useChatStreamMessage = () => { const iChatMessage: IBusterChatMessage = create(currentChatMessage, (draft) => { Object.assign(draft || {}, chatMessage); })!; + chatRefMessages.current[chatMessage.id] = iChatMessage; onUpdateChatMessage(iChatMessage!); @@ -65,7 +66,7 @@ export const useChatStreamMessage = () => { chatRef.current[iChat.id] = iChat; normalizeChatMessage(iChatMessages); onUpdateChat(iChat); - removeAutoThought({ messageId: iChat.message_ids[iChat.message_ids.length - 1] }); + removeBlackBoxMessage({ messageId: iChat.message_ids[iChat.message_ids.length - 1] }); }); const stopChatCallback = useMemoizedFn((chatId: string) => { @@ -127,12 +128,13 @@ export const useChatStreamMessage = () => { const _generatingReasoningMessageCallback = useMemoizedFn( (_: null, d: ChatEvent_GeneratingReasoningMessage) => { const { message_id, reasoning } = d; - const currentMessage = chatRefMessages.current[message_id]; - const updatedMessage = updateReasoningMessage(message_id, currentMessage, reasoning); + const updatedMessage = updateReasoningMessage( + message_id, + chatRefMessages.current[message_id], + reasoning + ); - checkAutoThought(updatedMessage, d); - - //TRIGGER + checkBlackBoxMessage(updatedMessage, d); //we only trigger black box message if the reasoning message is completed onUpdateChatMessageTransition({ id: message_id, diff --git a/web/src/controllers/ReasoningController/ReasoningMessages/ReasoningBlackBoxMessage.tsx b/web/src/controllers/ReasoningController/ReasoningMessages/ReasoningBlackBoxMessage.tsx index a92ae5d03..3962d226d 100644 --- a/web/src/controllers/ReasoningController/ReasoningMessages/ReasoningBlackBoxMessage.tsx +++ b/web/src/controllers/ReasoningController/ReasoningMessages/ReasoningBlackBoxMessage.tsx @@ -23,7 +23,7 @@ export const BlackBoxMessage: React.FC<{ messageId: string }> = React.memo(({ me ; } - return null; + return no black box message?; }); BlackBoxMessage.displayName = 'BlackBoxMessage';