From 213d657b006695cf274fa16ec57e0c52c9d79718 Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <62335616+lucaseduoli@users.noreply.github.com>
Date: Tue, 3 Dec 2024 18:33:18 -0300
Subject: [PATCH] fix: auto fill chat input on new sessions (#5034)

* Show chatview only if messages are fetched

* Get value from chatInput node to put in new sessions and make New Flow not flash on the screen
---
 .../components/chatView/newChatView.tsx       | 92 ++++++++++---------
 src/frontend/src/modals/IOModal/newModal.tsx  | 36 ++++----
 2 files changed, 69 insertions(+), 59 deletions(-)

diff --git a/src/frontend/src/modals/IOModal/components/chatView/newChatView.tsx b/src/frontend/src/modals/IOModal/components/chatView/newChatView.tsx
index fdcefd209..95a51b4e6 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/newChatView.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/newChatView.tsx
@@ -1,7 +1,6 @@
 import LangflowLogo from "@/assets/LangflowLogo.svg?react";
 import ChainLogo from "@/assets/logo.svg?react";
 import { TextEffectPerChar } from "@/components/ui/textAnimation";
-import { TextShimmer } from "@/components/ui/TextShimmer";
 import { ENABLE_NEW_LOGO } from "@/customization/feature-flags";
 import { track } from "@/customization/utils/analytics";
 import { useMessagesStore } from "@/stores/messagesStore";
@@ -14,7 +13,6 @@ import FlowRunningSqueleton from "../flowRunningSqueleton";
 import useDragAndDrop from "./chatInput/hooks/use-drag-and-drop";
 import { useFileHandler } from "./chatInput/hooks/use-file-handler";
 import ChatInput from "./chatInput/newChatInput";
-import LogoIcon from "./chatMessage/components/chatLogoIcon";
 import ChatMessage from "./chatMessage/newChatMessage";
 
 export default function ChatView({
@@ -27,11 +25,17 @@ export default function ChatView({
   focusChat,
   closeChat,
 }: chatViewProps): JSX.Element {
-  const { flowPool, inputs, CleanFlowPool } = useFlowStore();
+  const flowPool = useFlowStore((state) => state.flowPool);
+  const inputs = useFlowStore((state) => state.inputs);
   const currentFlowId = useFlowsManagerStore((state) => state.currentFlowId);
   const messagesRef = useRef<HTMLDivElement | null>(null);
-  const [chatHistory, setChatHistory] = useState<ChatMessageType[]>([]);
+  const [chatHistory, setChatHistory] = useState<ChatMessageType[] | undefined>(
+    undefined,
+  );
   const messages = useMessagesStore((state) => state.messages);
+  const nodes = useFlowStore((state) => state.nodes);
+  const chatInput = inputs.find((input) => input.type === "ChatInput");
+  const chatInputNode = nodes.find((node) => node.id === chatInput?.id);
 
   const inputTypes = inputs.map((obj) => obj.type);
   const updateFlowPool = useFlowStore((state) => state.updateFlowPool);
@@ -79,6 +83,9 @@ export default function ChatView({
       return new Date(a.timestamp).getTime() - new Date(b.timestamp).getTime();
     });
 
+    if (finalChatHistory.length === 0 && !lockChat && chatInputNode)
+      setChatValue(chatInputNode.data.node.template["input_value"].value ?? "");
+
     setChatHistory(finalChatHistory);
   }, [flowPool, messages, visibleSession]);
   useEffect(() => {
@@ -135,45 +142,46 @@ export default function ChatView({
       onDrop={onDrop}
     >
       <div ref={messagesRef} className="chat-message-div">
-        {lockChat || chatHistory?.length > 0 ? (
-          chatHistory.map((chat, index) => (
-            <ChatMessage
-              setLockChat={setLockChat}
-              lockChat={lockChat}
-              chat={chat}
-              lastMessage={chatHistory.length - 1 === index ? true : false}
-              key={`${chat.id}-${index}`}
-              updateChat={updateChat}
-              closeChat={closeChat}
-            />
-          ))
-        ) : (
-          <div className="flex h-full w-full flex-col items-center justify-center">
-            <div className="flex flex-col items-center justify-center gap-4 p-8">
-              {ENABLE_NEW_LOGO ? (
-                <LangflowLogo
-                  title="Langflow logo"
-                  className="h-10 w-10 scale-[1.5]"
-                />
-              ) : (
-                <ChainLogo
-                  title="Langflow logo"
-                  className="h-10 w-10 scale-[1.5]"
-                />
-              )}
-              <div className="flex flex-col items-center justify-center">
-                <h3 className="mt-2 pb-2 text-2xl font-semibold text-primary">
-                  New chat
-                </h3>
-                <p className="text-lg text-muted-foreground">
-                  <TextEffectPerChar>
-                    Test your flow with a chat prompt
-                  </TextEffectPerChar>
-                </p>
+        {chatHistory &&
+          (lockChat || chatHistory?.length > 0 ? (
+            chatHistory.map((chat, index) => (
+              <ChatMessage
+                setLockChat={setLockChat}
+                lockChat={lockChat}
+                chat={chat}
+                lastMessage={chatHistory.length - 1 === index ? true : false}
+                key={`${chat.id}-${index}`}
+                updateChat={updateChat}
+                closeChat={closeChat}
+              />
+            ))
+          ) : (
+            <div className="flex h-full w-full flex-col items-center justify-center">
+              <div className="flex flex-col items-center justify-center gap-4 p-8">
+                {ENABLE_NEW_LOGO ? (
+                  <LangflowLogo
+                    title="Langflow logo"
+                    className="h-10 w-10 scale-[1.5]"
+                  />
+                ) : (
+                  <ChainLogo
+                    title="Langflow logo"
+                    className="h-10 w-10 scale-[1.5]"
+                  />
+                )}
+                <div className="flex flex-col items-center justify-center">
+                  <h3 className="mt-2 pb-2 text-2xl font-semibold text-primary">
+                    New chat
+                  </h3>
+                  <p className="text-lg text-muted-foreground">
+                    <TextEffectPerChar>
+                      Test your flow with a chat prompt
+                    </TextEffectPerChar>
+                  </p>
+                </div>
               </div>
             </div>
-          </div>
-        )}
+          ))}
         <div
           className={
             lockChat ? "w-5/6 max-w-[768px] py-4 word-break-break-word" : ""
@@ -181,7 +189,7 @@ export default function ChatView({
           ref={ref}
         >
           {lockChat &&
-            !(chatHistory[chatHistory.length - 1]?.category === "error") &&
+            !(chatHistory?.[chatHistory.length - 1]?.category === "error") &&
             flowRunningSkeletonMemo}
         </div>
       </div>
diff --git a/src/frontend/src/modals/IOModal/newModal.tsx b/src/frontend/src/modals/IOModal/newModal.tsx
index f4985051e..1ab1304be 100644
--- a/src/frontend/src/modals/IOModal/newModal.tsx
+++ b/src/frontend/src/modals/IOModal/newModal.tsx
@@ -132,7 +132,7 @@ export default function IOModal({
   );
   const flowPool = useFlowStore((state) => state.flowPool);
   const [sessionId, setSessionId] = useState<string>(currentFlowId);
-  useGetMessagesQuery(
+  const { isFetched: messagesFetched } = useGetMessagesQuery(
     {
       mode: "union",
       id: currentFlowId,
@@ -475,22 +475,24 @@ export default function IOModal({
                       : "",
                   )}
                 >
-                  <ChatView
-                    focusChat={sessionId}
-                    sendMessage={sendMessage}
-                    chatValue={chatValue}
-                    setChatValue={setChatValue}
-                    lockChat={lockChat}
-                    setLockChat={setLockChat}
-                    visibleSession={visibleSession}
-                    closeChat={
-                      !canvasOpen
-                        ? undefined
-                        : () => {
-                            setOpen(false);
-                          }
-                    }
-                  />
+                  {messagesFetched && (
+                    <ChatView
+                      focusChat={sessionId}
+                      sendMessage={sendMessage}
+                      chatValue={chatValue}
+                      setChatValue={setChatValue}
+                      lockChat={lockChat}
+                      setLockChat={setLockChat}
+                      visibleSession={visibleSession}
+                      closeChat={
+                        !canvasOpen
+                          ? undefined
+                          : () => {
+                              setOpen(false);
+                            }
+                      }
+                    />
+                  )}
                 </div>
               ) : (
                 <span className="flex h-full w-full items-center justify-center font-thin text-muted-foreground">