From ee997a6c983fcd56e7542ca84dcdc4f785485223 Mon Sep 17 00:00:00 2001
From: Nate Kelley <nate@buster.so>
Date: Wed, 5 Mar 2025 23:11:41 -0700
Subject: [PATCH] reasoning message update

---
 web/src/components/ui/avatar/AvatarBase.tsx   |  2 +-
 .../ui/streaming/StreamingMessage_File.tsx    |  2 +-
 .../NewChatProvider/useChatStreamMessage.ts   |  1 -
 .../ChatResponseMessages.tsx                  | 32 ++++------
 .../ChatResponseReasoning.tsx                 | 61 ++++---------------
 5 files changed, 25 insertions(+), 73 deletions(-)

diff --git a/web/src/components/ui/avatar/AvatarBase.tsx b/web/src/components/ui/avatar/AvatarBase.tsx
index 5f3248864..81917eea9 100644
--- a/web/src/components/ui/avatar/AvatarBase.tsx
+++ b/web/src/components/ui/avatar/AvatarBase.tsx
@@ -35,7 +35,7 @@ const AvatarFallback = React.forwardRef<
   <AvatarPrimitive.Fallback
     ref={ref}
     className={cn(
-      'bg-gray-light/60 text-background flex h-full w-full items-center justify-center rounded-full text-base',
+      'bg-gray-light/60 text-background flex h-full w-full items-center justify-center rounded-full text-xs',
       className
     )}
     {...props}
diff --git a/web/src/components/ui/streaming/StreamingMessage_File.tsx b/web/src/components/ui/streaming/StreamingMessage_File.tsx
index 2f19e4bbc..1f2187bfe 100644
--- a/web/src/components/ui/streaming/StreamingMessage_File.tsx
+++ b/web/src/components/ui/streaming/StreamingMessage_File.tsx
@@ -45,7 +45,7 @@ StreamingMessage_File.displayName = 'StreamingMessage_File';
 const StreamHeader: React.FC<{ file_name: string; version_number: number }> = React.memo(
   ({ file_name, version_number }) => {
     return (
-      <div className="file-header bg-item-hover border-border flex h-8 w-full items-center space-x-1.5 overflow-hidden border-b px-2.5">
+      <div className="file-header bg-item-select border-border flex h-8 w-full items-center space-x-1.5 overflow-hidden border-b px-2.5">
         <Text truncate>{file_name}</Text>
         <VersionPill version_number={version_number} />
       </div>
diff --git a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
index 7f3503f36..ad8f12f8a 100644
--- a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
+++ b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts
@@ -182,7 +182,6 @@ export const useChatStreamMessage = () => {
             draft[chat_id]?.messages?.[message_id]?.response_messages?.[responseMessageId];
           if (!responseMessage) return;
           const messageText = responseMessage as BusterChatMessageReasoning_text;
-
           Object.assign(messageText, {
             ...existingResponseMessageText,
             ...response_message,
diff --git a/web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessages.tsx b/web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessages.tsx
index f20f5ee3a..977d90e40 100644
--- a/web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessages.tsx
+++ b/web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessages.tsx
@@ -20,13 +20,19 @@ export const ChatResponseMessages: React.FC<ChatResponseMessagesProps> = React.m
       messageId,
       (x) => x?.reasoning_message_ids?.[x.reasoning_message_ids.length - 1]
     );
-
-    const lastMessageIndex = responseMessageIds.length - 1;
-    const showDefaultMessage = responseMessageIds.length === 0;
+    const finalReasoningMessage = useMessageIndividual(
+      messageId,
+      (x) => x?.final_reasoning_message
+    );
 
     return (
-      <MessageContainer className="flex w-full flex-col overflow-hidden">
-        {showDefaultMessage && <DefaultFirstMessage />}
+      <MessageContainer className="flex w-full flex-col space-y-3 overflow-hidden">
+        <ChatResponseReasoning
+          reasoningMessageId={lastReasoningMessageId}
+          finalReasoningMessage={finalReasoningMessage}
+          isCompletedStream={isCompletedStream}
+          messageId={messageId}
+        />
 
         {responseMessageIds.map((responseMessageId, index) => (
           <React.Fragment key={responseMessageId}>
@@ -35,14 +41,6 @@ export const ChatResponseMessages: React.FC<ChatResponseMessagesProps> = React.m
               messageId={messageId}
               isCompletedStream={isCompletedStream}
             />
-
-            {index === lastMessageIndex && lastReasoningMessageId && (
-              <ChatResponseReasoning
-                reasoningMessageId={lastReasoningMessageId}
-                isCompletedStream={isCompletedStream}
-                messageId={messageId}
-              />
-            )}
           </React.Fragment>
         ))}
       </MessageContainer>
@@ -51,11 +49,3 @@ export const ChatResponseMessages: React.FC<ChatResponseMessagesProps> = React.m
 );
 
 ChatResponseMessages.displayName = 'ChatResponseMessages';
-
-const DefaultFirstMessage: React.FC = () => {
-  return (
-    <div>
-      <ShimmerText text="Thinking..." />
-    </div>
-  );
-};
diff --git a/web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseReasoning.tsx b/web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseReasoning.tsx
index 551a7a8e6..ef47ed1fd 100644
--- a/web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseReasoning.tsx
+++ b/web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseReasoning.tsx
@@ -3,11 +3,9 @@ import { ShimmerText } from '@/components/ui/typography/ShimmerText';
 import { useMemoizedFn } from 'ahooks';
 import { motion } from 'framer-motion';
 import { AnimatePresence } from 'framer-motion';
-import { Stars } from '@/components/ui/icons';
 import { Text } from '@/components/ui/typography';
 import { useChatLayoutContextSelector } from '../../../ChatLayoutContext';
 import { useMessageIndividual } from '@/context/Chats';
-import { cn } from '@/lib/classMerge';
 
 const animations = {
   initial: { opacity: 0 },
@@ -16,21 +14,23 @@ const animations = {
 };
 
 export const ChatResponseReasoning: React.FC<{
-  reasoningMessageId: string;
+  reasoningMessageId: string | undefined;
+  finalReasoningMessage: string | undefined | null;
   isCompletedStream: boolean;
   messageId: string;
 }> = React.memo(({ reasoningMessageId, isCompletedStream, messageId }) => {
   const lastMessageTitle = useMessageIndividual(
     messageId,
-    (x) => x?.reasoning_messages?.[reasoningMessageId]?.title
+    (x) => x?.reasoning_messages?.[reasoningMessageId ?? '']?.title
   );
   const finalReasoningMessage = useMessageIndividual(messageId, (x) => x?.final_reasoning_message);
   const onSetSelectedFile = useChatLayoutContextSelector((x) => x.onSetSelectedFile);
   const selectedFileType = useChatLayoutContextSelector((x) => x.selectedFileType);
-  const isReasonginFileSelected = selectedFileType === 'reasoning';
+  const isReasonginFileSelected = selectedFileType === 'reasoning' && isCompletedStream;
 
   const text: string = useMemo(() => {
     if (finalReasoningMessage) return finalReasoningMessage;
+    if (lastMessageTitle) return lastMessageTitle;
     return lastMessageTitle || 'Thinking...';
   }, [lastMessageTitle, finalReasoningMessage]);
 
@@ -41,56 +41,19 @@ export const ChatResponseReasoning: React.FC<{
     });
   });
 
+  console.log(isReasonginFileSelected, isCompletedStream);
+
   return (
     <AnimatePresence initial={!isCompletedStream} mode="wait">
       <motion.div {...animations} key={text} className="mb-3.5 w-fit" onClick={onClickReasoning}>
-        <ShimmerTextWithIcon
-          text={text ?? ''}
-          isCompletedStream={isCompletedStream}
-          isSelected={isReasonginFileSelected}
-        />
+        {isReasonginFileSelected ? (
+          <Text className="cursor-pointer hover:underline">{text}</Text>
+        ) : (
+          <ShimmerText text={text ?? ''} />
+        )}
       </motion.div>
     </AnimatePresence>
   );
 });
 
 ChatResponseReasoning.displayName = 'ChatThoughts';
-
-const ShimmerTextWithIcon = React.memo(
-  ({
-    text,
-    isCompletedStream,
-    isSelected
-  }: {
-    text: string;
-    isCompletedStream: boolean;
-    isSelected: boolean;
-  }) => {
-    if (isCompletedStream) {
-      return (
-        <div
-          className={cn(
-            'text-icon-color hover:text-foreground',
-            'cursor-pointer',
-            'flex w-fit items-center gap-1',
-            isSelected && 'text-foreground'
-          )}>
-          <div>
-            <Stars />
-          </div>
-          <Text variant="inherit">{text}</Text>
-        </div>
-      );
-    }
-
-    return (
-      <div className={'text-icon-color flex cursor-pointer items-center gap-1'}>
-        <div className={'text-icon-color'}>
-          <Stars />
-        </div>
-        <ShimmerText text={text} />
-      </div>
-    );
-  }
-);
-ShimmerTextWithIcon.displayName = 'ShimmerTextWithIcon';