From f3fa7ac35a594dc262b97b7a6e19494f1b6474d0 Mon Sep 17 00:00:00 2001
From: Brian Hackett <bhackett1024@gmail.com>
Date: Mon, 20 Jan 2025 09:35:24 -0800
Subject: [PATCH] Messages WIP

---
 app/components/chat/Chat.client.tsx | 50 ++++++++++++-----------------
 1 file changed, 21 insertions(+), 29 deletions(-)

diff --git a/app/components/chat/Chat.client.tsx b/app/components/chat/Chat.client.tsx
index 3441c03d7..8d0faf046 100644
--- a/app/components/chat/Chat.client.tsx
+++ b/app/components/chat/Chat.client.tsx
@@ -120,6 +120,8 @@ export const ChatImpl = memo(
     const [uploadedFiles, setUploadedFiles] = useState<File[]>([]); // Move here
     const [imageDataList, setImageDataList] = useState<string[]>([]); // Move here
     const [searchParams, setSearchParams] = useSearchParams();
+    const [trailingMessages, setTrailingMessages] = useState<Message[]>([]);
+    const [simulationLoading, setSimulationLoading] = useState(false);
     const files = useStore(workbenchStore.files);
     const { promptId } = useSettings();
 
@@ -140,20 +142,10 @@ export const ChatImpl = memo(
           'There was an error processing your request: ' + (error.message ? error.message : 'No details were returned'),
         );
       },
-      onFinish: (message, response) => {
-        const usage = response.usage;
-
-        if (usage) {
-          console.log('Token usage:', usage);
-
-          // You can now use the usage data as needed
-        }
-
-        logger.debug('Finished streaming');
-      },
       initialMessages,
       initialInput: Cookies.get(PROMPT_COOKIE_KEY) || '',
     });
+
     useEffect(() => {
       const prompt = searchParams.get('prompt');
 
@@ -204,6 +196,7 @@ export const ChatImpl = memo(
       gNumAborts++;
       chatStore.setKey('aborted', true);
       workbenchStore.abortAllActions();
+      setSimulationLoading(false);
     };
 
     useEffect(() => {
@@ -244,14 +237,10 @@ export const ChatImpl = memo(
         message = "Error creating recording.";
       }
 
-      const recordingMessage: CreateMessage = {
+      const recordingMessage: Message = {
+        id: `create-recording-${messages.length}`,
         role: 'assistant',
-        content: [
-          {
-            type: 'text',
-            text: message,
-          },
-        ] as any, // Type assertion to bypass compiler check
+        content: message,
       };
 
       return { recordingId, recordingMessage };
@@ -267,14 +256,10 @@ export const ChatImpl = memo(
         message = "Error enhancing prompt.";
       }
 
-      const enhancedPromptMessage: CreateMessage = {
+      const enhancedPromptMessage: Message = {
+        id: `enhanced-prompt-${messages.length}`,
         role: 'assistant',
-        content: [
-          {
-            type: 'text',
-            text: message,
-          },
-        ] as any, // Type assertion to bypass compiler check
+        content: message,
       };
 
       return { enhancedPrompt, enhancedPromptMessage };
@@ -288,6 +273,8 @@ export const ChatImpl = memo(
         return;
       }
 
+      setSimulationLoading(true);
+
       /**
        * @note (delm) Usually saving files shouldn't take long but it may take longer if there
        * many unsaved files. In that case we need to block user input and show an indicator
@@ -309,7 +296,8 @@ export const ChatImpl = memo(
           return;
         }
 
-        append(recordingMessage);
+        console.log("RecordingMessage", recordingMessage);
+        setTrailingMessages([...trailingMessages, recordingMessage]);
 
         if (recordingId) {
           const info = await getEnhancedPrompt(recordingId, contentBase64);
@@ -319,7 +307,9 @@ export const ChatImpl = memo(
           }
   
           simulationEnhancedPrompt = info.enhancedPrompt;
-          append(info.enhancedPromptMessage);
+
+          console.log("EnhancedPromptMessage", info.enhancedPromptMessage);
+          setTrailingMessages([...trailingMessages, info.enhancedPromptMessage]);
         }
       }
   
@@ -329,6 +319,8 @@ export const ChatImpl = memo(
 
       runAnimation();
 
+      setSimulationLoading(false);
+
       append({
         role: 'user',
         content: [
@@ -397,7 +389,7 @@ export const ChatImpl = memo(
         input={input}
         showChat={showChat}
         chatStarted={chatStarted}
-        isStreaming={isLoading}
+        isStreaming={isLoading || simulationLoading}
         enhancingPrompt={enhancingPrompt}
         promptEnhanced={promptEnhanced}
         sendMessage={sendMessage}
@@ -411,7 +403,7 @@ export const ChatImpl = memo(
         description={description}
         importChat={importChat}
         exportChat={exportChat}
-        messages={messages.map((message, i) => {
+        messages={[...messages, ...trailingMessages].map((message, i) => {
           if (message.role === 'user') {
             return message;
           }