thinking step

This commit is contained in:
Nate Kelley 2025-02-11 15:02:02 -07:00
parent bf24b249d4
commit 9f914a34f9
No known key found for this signature in database
GPG Key ID: FD90372AB8D98B4F
12 changed files with 128 additions and 96 deletions

View File

@ -22,6 +22,7 @@ export type BusterChatMessage_text = {
type: 'text';
message: string;
message_chunk?: string;
is_final_message?: boolean; //defaults to false
};
export type BusterChatMessage_fileMetadata = {

View File

@ -198,8 +198,6 @@ const NewChatInput: React.FC<{
const onSelectSearchAsset = useBusterNewChatContextSelector((x) => x.onSelectSearchAsset);
const [loadingNewChat, setLoadingNewChat] = useState(false);
console.log(selectedChatDataSource);
const onStartNewChatPreflight = useMemoizedFn(async () => {
setLoadingNewChat(true);
await onStartNewChat({ prompt, datasetId: selectedChatDataSource?.id });

View File

@ -1,28 +1,30 @@
import React, { useEffect, useState } from 'react';
import React, { useEffect, useRef, useTransition } from 'react';
import { AnimatePresence, motion } from 'framer-motion';
import { itemAnimationConfig } from './animationConfig';
import { useMemoizedFn } from 'ahooks';
interface StreamingMessage_TextProps {
isCompletedStream: boolean;
message: {
message_chunk?: string;
message?: string;
};
message: string;
}
export const StreamingMessage_Text: React.FC<StreamingMessage_TextProps> = React.memo(
({ message: messageProp, isCompletedStream }) => {
const { message_chunk, message } = messageProp;
({ message, isCompletedStream }) => {
const [isPending, startTransition] = useTransition();
const textChunksRef = useRef<string[]>([]);
const [textChunks, setTextChunks] = useState<string[]>([]);
const setTextChunks = useMemoizedFn((updater: (prevChunks: string[]) => string[]) => {
textChunksRef.current = updater(textChunksRef.current);
startTransition(() => {
//just used to trigger UI update
});
});
useEffect(() => {
if (message_chunk && !message) {
// Handle streaming chunks
setTextChunks((prevChunks) => [...prevChunks, message_chunk || '']);
} else if (message) {
if (message) {
// Handle complete message
const currentText = textChunks.join('');
const currentText = textChunksRef.current.join('');
if (message.startsWith(currentText)) {
const remainingText = message.slice(currentText.length);
if (remainingText) {
@ -30,14 +32,14 @@ export const StreamingMessage_Text: React.FC<StreamingMessage_TextProps> = React
}
} else {
// If there's a mismatch, just use the complete message
setTextChunks([message]);
setTextChunks(() => [message]);
}
}
}, [message_chunk, message]);
}, [message]);
return (
<div className={''}>
{textChunks.map((chunk, index) => (
{textChunksRef.current.map((chunk, index) => (
<AnimatePresence key={index} initial={!isCompletedStream}>
<motion.span {...itemAnimationConfig}>{chunk}</motion.span>
</AnimatePresence>

View File

@ -22,7 +22,7 @@ const ReasoningMessageRecord: Record<
text: (props) => (
<StreamingMessage_Text
{...props}
message={props.reasoningMessage as BusterChatMessageReasoning_text}
message={(props.reasoningMessage as BusterChatMessageReasoning_text).message}
/>
),
file: ReasoningMessage_File

View File

@ -16,7 +16,10 @@ const ChatResponseMessageRecord: Record<
React.FC<ChatResponseMessageProps>
> = {
text: (props) => (
<StreamingMessage_Text {...props} message={props.responseMessage as BusterChatMessage_text} />
<StreamingMessage_Text
{...props}
message={(props.responseMessage as BusterChatMessage_text).message}
/>
),
file: ChatResponseMessage_File
};
@ -27,37 +30,37 @@ export interface ChatResponseMessageSelectorProps {
isLastMessageItem: boolean;
}
export const ChatResponseMessageSelector: React.FC<ChatResponseMessageSelectorProps> = ({
responseMessage,
isCompletedStream,
isLastMessageItem
}) => {
const messageType = responseMessage.type;
const ChatResponseMessage = ChatResponseMessageRecord[messageType];
const { cx, styles } = useStyles();
export const ChatResponseMessageSelector: React.FC<ChatResponseMessageSelectorProps> = React.memo(
({ responseMessage, isCompletedStream, isLastMessageItem }) => {
const { cx, styles } = useStyles();
const messageType = responseMessage.type;
const ChatResponseMessage = ChatResponseMessageRecord[messageType];
const typeClassRecord: Record<BusterChatMessageResponse['type'], string> = useMemo(() => {
return {
text: cx(styles.textCard, 'text-card'),
file: cx(styles.fileCard, 'file-card')
};
}, []);
const typeClassRecord: Record<BusterChatMessageResponse['type'], string> = useMemo(() => {
return {
text: cx(styles.textCard, 'text-card'),
file: cx(styles.fileCard, 'file-card')
};
}, []);
const getContainerClass = useMemoizedFn((item: BusterChatMessageResponse) => {
return typeClassRecord[item.type];
});
const getContainerClass = useMemoizedFn((item: BusterChatMessageResponse) => {
return typeClassRecord[item.type];
});
return (
<div key={responseMessage.id} className={getContainerClass(responseMessage)}>
<ChatResponseMessage
responseMessage={responseMessage}
isCompletedStream={isCompletedStream}
isLastMessageItem={isLastMessageItem}
/>
<VerticalDivider />
</div>
);
};
return (
<div key={responseMessage.id} className={getContainerClass(responseMessage)}>
<ChatResponseMessage
responseMessage={responseMessage}
isCompletedStream={isCompletedStream}
isLastMessageItem={isLastMessageItem}
/>
<VerticalDivider />
</div>
);
}
);
ChatResponseMessageSelector.displayName = 'ChatResponseMessageSelector';
const VerticalDivider: React.FC<{ className?: string }> = React.memo(({ className }) => {
const { cx, styles } = useStyles();

View File

@ -7,6 +7,7 @@ import type {
import { MessageContainer } from '../MessageContainer';
import { ChatResponseMessageSelector } from './ChatResponseMessageSelector';
import { ChatResponseReasoning } from './ChatResponseReasoning';
import { ShimmerText } from '@/components/text';
interface ChatResponseMessagesProps {
responseMessages: BusterChatMessageResponse[];
@ -17,40 +18,41 @@ interface ChatResponseMessagesProps {
export const ChatResponseMessages: React.FC<ChatResponseMessagesProps> = React.memo(
({ responseMessages, reasoningMessages, isCompletedStream, messageId }) => {
const firstResponseMessage = responseMessages[0] as BusterChatMessage_text;
const restResponseMessages = useMemo(() => {
if (!firstResponseMessage) return [];
return responseMessages.slice(1);
}, [firstResponseMessage, responseMessages]);
const lastMessageIndex = responseMessages.length - 1;
const showDefaultMessage = responseMessages.length === 0;
const reasonginStepIndex = useMemo(() => {
const lastTextMessage = responseMessages.findLast(
(message) => message.type === 'text' && message.is_final_message !== false
) as BusterChatMessage_text;
if (!lastTextMessage) return -1;
if (lastTextMessage?.message_chunk) return -1;
return responseMessages.findIndex((message) => message.id === lastTextMessage.id);
}, [responseMessages]);
return (
<MessageContainer className="flex w-full flex-col overflow-hidden">
{firstResponseMessage && (
<ChatResponseMessageSelector
key={firstResponseMessage.id}
responseMessage={firstResponseMessage}
isCompletedStream={isCompletedStream}
isLastMessageItem={false}
/>
)}
{showDefaultMessage && <DefaultFirstMessage />}
{firstResponseMessage && (
<ChatResponseReasoning
reasoningMessages={reasoningMessages}
isCompletedStream={isCompletedStream}
messageId={messageId}
/>
)}
{responseMessages.map((responseMessage, index) => (
<React.Fragment key={responseMessage.id}>
<ChatResponseMessageSelector
responseMessage={responseMessage}
isCompletedStream={isCompletedStream}
isLastMessageItem={index === lastMessageIndex}
/>
{restResponseMessages.map((responseMessage, index) => (
<ChatResponseMessageSelector
key={responseMessage.id}
responseMessage={responseMessage}
isCompletedStream={isCompletedStream}
isLastMessageItem={index === lastMessageIndex}
/>
{index === reasonginStepIndex && (
<ChatResponseReasoning
reasoningMessages={reasoningMessages}
isCompletedStream={isCompletedStream}
messageId={messageId}
/>
)}
</React.Fragment>
))}
</MessageContainer>
);
@ -58,3 +60,11 @@ export const ChatResponseMessages: React.FC<ChatResponseMessagesProps> = React.m
);
ChatResponseMessages.displayName = 'ChatResponseMessages';
const DefaultFirstMessage: React.FC = () => {
return (
<div>
<ShimmerText text="Thinking..." />
</div>
);
};

View File

@ -11,7 +11,7 @@ export const MessageContainer: React.FC<{
}> = React.memo(({ children, senderName, senderId, senderAvatar, className = '' }) => {
const { cx } = useStyles();
return (
<div className={cx('flex w-full space-x-1 overflow-hidden')}>
<div className={cx('flex w-full space-x-2 overflow-hidden')}>
{senderName ? (
<BusterUserAvatar size={24} name={senderName} src={senderAvatar} useToolTip={true} />
) : (

View File

@ -17,25 +17,25 @@ export const useChatSelectors = ({
const getChatMessages = useCallback(
(chatId: string): IBusterChatMessage[] => {
const chatMessageIds = chatsRef.current[chatId].messages || [];
return chatMessageIds.map((messageId) => chatsMessagesRef.current[messageId]);
return getChatMessagesMemoized(chatId);
},
[chatsMessagesRef, isPending, chatsRef]
);
const getChatMessage = useCallback(
(messageId: string): IBusterChatMessage | undefined => {
return chatsMessagesRef.current[messageId];
return getChatMessageMemoized(messageId);
},
[chatsMessagesRef, isPending]
);
const getChatMessagesMemoized = useMemoizedFn((chatId: string) => {
return getChatMessages(chatId);
const chatMessageIds = chatsRef.current[chatId].messages || [];
return chatMessageIds.map((messageId) => chatsMessagesRef.current[messageId]);
});
const getChatMessageMemoized = useMemoizedFn((messageId: string) => {
return getChatMessage(messageId);
return chatsMessagesRef.current[messageId];
});
return {

View File

@ -17,7 +17,7 @@ export const useChatSubscriptions = ({
const busterSocket = useBusterWebSocket();
const _onGetChat = useMemoizedFn((chat: BusterChat): IBusterChat => {
const { iChat, iChatMessages } = updateChatToIChat(chat);
const { iChat, iChatMessages } = updateChatToIChat(chat, false);
chatsRef.current[chat.id] = iChat;
chatsMessagesRef.current = {

View File

@ -11,9 +11,7 @@ import random from 'lodash/random';
export const useAutoAppendThought = () => {
const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage);
const getChatMemoized = useBusterChatContextSelector((x) => x.getChatMemoized);
const getChatMessagesMemoized = useBusterChatContextSelector((x) => x.getChatMessagesMemoized);
const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized);
const removeAutoThoughts = useMemoizedFn(
(reasoningMessages: BusterChatMessageReasoning[]): BusterChatMessageReasoning[] => {
@ -26,8 +24,9 @@ export const useAutoAppendThought = () => {
reasoningMessages: BusterChatMessageReasoning[],
chatId: string
): BusterChatMessageReasoning[] => {
const lastReasoningMessage = reasoningMessages[reasoningMessages.length - 1];
const lastMessageIsCompleted =
reasoningMessages[reasoningMessages.length - 1].status === 'completed';
!lastReasoningMessage || lastReasoningMessage?.status === 'completed';
if (lastMessageIsCompleted) {
_loopAutoThought(chatId);
@ -40,13 +39,14 @@ export const useAutoAppendThought = () => {
);
const _loopAutoThought = useMemoizedFn(async (chatId: string) => {
const randomDelay = random(3500, 5500);
const randomDelay = random(3000, 5000);
await timeout(randomDelay);
const chatMessages = getChatMessagesMemoized(chatId);
const lastMessage = last(chatMessages);
const isCompletedStream = !!lastMessage?.isCompletedStream;
const lastReasoningMessage = last(lastMessage?.reasoning);
const lastReasoningMessageIsAutoAppended = lastReasoningMessage?.id === AUTO_THOUGHT_ID;
const lastReasoningMessageIsAutoAppended =
!lastReasoningMessage || lastReasoningMessage?.id === AUTO_THOUGHT_ID;
if (!isCompletedStream && lastReasoningMessageIsAutoAppended && lastMessage) {
const lastMessageId = lastMessage?.id!;

View File

@ -1,7 +1,7 @@
import { useMemoizedFn } from 'ahooks';
import { useBusterChatContextSelector } from '../ChatProvider';
import { useBusterWebSocket } from '@/context/BusterWebSocket';
import { BusterChat } from '@/api/asset_interfaces';
import { BusterChat, BusterChatMessage_text } from '@/api/asset_interfaces';
import {
ChatEvent_GeneratingReasoningMessage,
ChatEvent_GeneratingResponseMessage,
@ -11,6 +11,7 @@ import { updateChatToIChat } from '@/utils/chat';
import { useAutoAppendThought } from './useAutoAppendThought';
import { useAppLayoutContextSelector } from '@/context/BusterAppLayout';
import { BusterRoutes } from '@/routes';
import last from 'lodash/last';
export const useChatUpdateMessage = () => {
const busterSocket = useBusterWebSocket();
@ -40,12 +41,27 @@ export const useChatUpdateMessage = () => {
(d: ChatEvent_GeneratingResponseMessage) => {
const { message_id, response_message } = d;
const currentResponseMessages = getChatMessageMemoized(message_id)?.response_messages ?? [];
const isNewMessage = !currentResponseMessages.some(({ id }) => id === message_id);
const responseMessageId = response_message.id;
const foundResponseMessage = currentResponseMessages.find(
({ id }) => id === responseMessageId
);
const isNewMessage = !foundResponseMessage;
if (response_message.type === 'text') {
const existingMessage = (foundResponseMessage as BusterChatMessage_text)?.message || '';
const isStreaming = !!response_message.message_chunk;
if (isStreaming) {
response_message.message = existingMessage + response_message.message_chunk;
}
}
onUpdateChatMessage({
id: message_id,
response_messages: isNewMessage
? [...currentResponseMessages, response_message]
: currentResponseMessages.map((rm) => (rm.id === message_id ? response_message : rm))
: currentResponseMessages.map((rm) =>
rm.id === responseMessageId ? response_message : rm
)
});
}
);
@ -68,7 +84,7 @@ export const useChatUpdateMessage = () => {
);
const completeChatCallback = useMemoizedFn((d: BusterChat) => {
const { iChat, iChatMessages } = updateChatToIChat(d);
const { iChat, iChatMessages } = updateChatToIChat(d, false);
onBulkSetChatMessages(iChatMessages);
onUpdateChat(iChat);
});
@ -81,10 +97,9 @@ export const useChatUpdateMessage = () => {
});
const initializeChatCallback = useMemoizedFn((d: BusterChat) => {
const { iChat, iChatMessages } = updateChatToIChat(d);
const { iChat, iChatMessages } = updateChatToIChat(d, true);
onBulkSetChatMessages(iChatMessages);
onUpdateChat(iChat);
onChangePage({
route: BusterRoutes.APP_CHAT_ID,
chatId: iChat.id

View File

@ -35,9 +35,12 @@ const chatMessageUpgrader = (
);
};
export const updateChatToIChat = (chat: BusterChat) => {
export const updateChatToIChat = (chat: BusterChat, isNewChat: boolean) => {
const iChat = chatUpgrader(chat);
const iChatMessages = chatMessageUpgrader(chat.messages);
const iChatMessages = chatMessageUpgrader(chat.messages, {
isCompletedStream: !isNewChat,
messageId: chat.messages[0].id
});
return {
iChat,
iChatMessages