From 60945ce883eb78a69f8bda27beceebe91e0deb50 Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Wed, 30 Aug 2023 18:49:41 -0300
Subject: [PATCH] added check to onCloseWebsocket event

---
 src/frontend/src/controllers/API/index.ts   |   2 +-
 src/frontend/src/modals/formModal/index.tsx | 175 ++++++++++----------
 2 files changed, 93 insertions(+), 84 deletions(-)

diff --git a/src/frontend/src/controllers/API/index.ts b/src/frontend/src/controllers/API/index.ts
index 3e6c78801..780d784e5 100644
--- a/src/frontend/src/controllers/API/index.ts
+++ b/src/frontend/src/controllers/API/index.ts
@@ -312,7 +312,7 @@ export async function getHealth() {
  */
 export async function getBuildStatus(
   flowId: string
-): Promise<BuildStatusTypeAPI> {
+): Promise<AxiosResponse<BuildStatusTypeAPI>> {
   return await api.get(`${BASE_URL_API}build/${flowId}/status`);
 }
 
diff --git a/src/frontend/src/modals/formModal/index.tsx b/src/frontend/src/modals/formModal/index.tsx
index 041132d0d..0a9bf7483 100644
--- a/src/frontend/src/modals/formModal/index.tsx
+++ b/src/frontend/src/modals/formModal/index.tsx
@@ -8,7 +8,7 @@ import { classNames } from "../../utils/utils";
 import ChatInput from "./chatInput";
 import ChatMessage from "./chatMessage";
 
-import _ from "lodash";
+import _, { set } from "lodash";
 import AccordionComponent from "../../components/AccordionComponent";
 import IconComponent from "../../components/genericIconComponent";
 import ToggleShadComponent from "../../components/toggleShadComponent";
@@ -27,6 +27,7 @@ import { AuthContext } from "../../contexts/authContext";
 import { TabsContext } from "../../contexts/tabsContext";
 import { TabsState } from "../../types/tabs";
 import { validateNodes } from "../../utils/reactflowUtils";
+import { getBuildStatus } from "../../controllers/API";
 
 export default function FormModal({
   flow,
@@ -155,9 +156,21 @@ export default function FormModal({
 
   function handleOnClose(event: CloseEvent): void {
     if (isOpen.current) {
+      getBuildStatus(flow.id).then((response) => {
+        if (response.data.built) {
+          connectWS();
+        }
+        else {
+          setErrorData({
+            title: "Please build the flow again before using the chat."
+          })
+        }
+      }).catch((error) => {
+        setErrorData({title:error.data?.detail?error.data.detail:error.message})
+
+      });
       setErrorData({ title: event.reason });
       setTimeout(() => {
-        connectWS();
         setLockChat(false);
       }, 1000);
     }
@@ -173,9 +186,8 @@ export default function FormModal({
     const host = isDevelopment ? "localhost:7860" : window.location.host;
     const chatEndpoint = `/api/v1/chat/${chatId}`;
 
-    return `${
-      isDevelopment ? "ws" : webSocketProtocol
-    }://${host}${chatEndpoint}?token=${encodeURIComponent(accessToken!)}`;
+    return `${isDevelopment ? "ws" : webSocketProtocol
+      }://${host}${chatEndpoint}?token=${encodeURIComponent(accessToken!)}`;
   }
 
   function handleWsMessage(data: any) {
@@ -197,20 +209,20 @@ export default function FormModal({
               newChatHistory.push(
                 chatItem.files
                   ? {
-                      isSend: !chatItem.is_bot,
-                      message: chatItem.message,
-                      template: chatItem.template,
-                      thought: chatItem.intermediate_steps,
-                      files: chatItem.files,
-                      chatKey: chatItem.chatKey,
-                    }
+                    isSend: !chatItem.is_bot,
+                    message: chatItem.message,
+                    template: chatItem.template,
+                    thought: chatItem.intermediate_steps,
+                    files: chatItem.files,
+                    chatKey: chatItem.chatKey,
+                  }
                   : {
-                      isSend: !chatItem.is_bot,
-                      message: chatItem.message,
-                      template: chatItem.template,
-                      thought: chatItem.intermediate_steps,
-                      chatKey: chatItem.chatKey,
-                    }
+                    isSend: !chatItem.is_bot,
+                    message: chatItem.message,
+                    template: chatItem.template,
+                    thought: chatItem.intermediate_steps,
+                    chatKey: chatItem.chatKey,
+                  }
               );
             }
           }
@@ -260,7 +272,6 @@ export default function FormModal({
       };
       newWs.onmessage = (event) => {
         const data = JSON.parse(event.data);
-        console.log("Received data:", data);
         handleWsMessage(data);
         //get chat history
       };
@@ -268,7 +279,6 @@ export default function FormModal({
         handleOnClose(event);
       };
       newWs.onerror = (ev) => {
-        console.log(ev, "error");
         if (flow.id === "") {
           connectWS();
         } else {
@@ -294,7 +304,6 @@ export default function FormModal({
   useEffect(() => {
     connectWS();
     return () => {
-      console.log("unmount");
       console.log(ws);
       if (ws.current) {
         ws.current.close();
@@ -433,73 +442,73 @@ export default function FormModal({
 
               {tabsState[id.current]?.formKeysData?.input_keys
                 ? Object.keys(
-                    tabsState[id.current].formKeysData.input_keys!
-                  ).map((key, index) => (
-                    <div className="file-component-accordion-div" key={index}>
-                      <AccordionComponent
-                        trigger={
-                          <div className="file-component-badge-div">
-                            <Badge variant="gray" size="md">
-                              {key}
-                            </Badge>
+                  tabsState[id.current].formKeysData.input_keys!
+                ).map((key, index) => (
+                  <div className="file-component-accordion-div" key={index}>
+                    <AccordionComponent
+                      trigger={
+                        <div className="file-component-badge-div">
+                          <Badge variant="gray" size="md">
+                            {key}
+                          </Badge>
 
-                            <div
-                              className="-mb-1"
-                              onClick={(event) => {
-                                event.stopPropagation();
-                              }}
-                            >
-                              <ToggleShadComponent
-                                enabled={chatKey === key}
-                                setEnabled={(value) =>
-                                  handleOnCheckedChange(value, key)
-                                }
-                                size="small"
-                                disabled={tabsState[
-                                  id.current
-                                ].formKeysData.handle_keys!.some(
-                                  (t) => t === key
-                                )}
-                              />
-                            </div>
+                          <div
+                            className="-mb-1"
+                            onClick={(event) => {
+                              event.stopPropagation();
+                            }}
+                          >
+                            <ToggleShadComponent
+                              enabled={chatKey === key}
+                              setEnabled={(value) =>
+                                handleOnCheckedChange(value, key)
+                              }
+                              size="small"
+                              disabled={tabsState[
+                                id.current
+                              ].formKeysData.handle_keys!.some(
+                                (t) => t === key
+                              )}
+                            />
                           </div>
-                        }
-                        key={index}
-                        keyValue={key}
-                      >
-                        <div className="file-component-tab-column">
-                          {tabsState[id.current].formKeysData.handle_keys!.some(
-                            (t) => t === key
-                          ) && (
+                        </div>
+                      }
+                      key={index}
+                      keyValue={key}
+                    >
+                      <div className="file-component-tab-column">
+                        {tabsState[id.current].formKeysData.handle_keys!.some(
+                          (t) => t === key
+                        ) && (
                             <div className="font-normal text-muted-foreground ">
                               Source: Component
                             </div>
                           )}
-                          <Textarea
-                            className="custom-scroll"
-                            value={
-                              tabsState[id.current].formKeysData.input_keys![
-                                key
-                              ]
-                            }
-                            onChange={(e) => {
-                              //@ts-ignore
-                              setTabsState((old: TabsState) => {
-                                let newTabsState = _.cloneDeep(old);
-                                newTabsState[
-                                  id.current
-                                ].formKeysData.input_keys![key] =
-                                  e.target.value;
-                                return newTabsState;
-                              });
-                            }}
-                            disabled={chatKey === key}
-                            placeholder="Enter text..."
-                          ></Textarea>
-                        </div>
-                      </AccordionComponent>
-                    </div>
-                  ))
+                        <Textarea
+                          className="custom-scroll"
+                          value={
+                            tabsState[id.current].formKeysData.input_keys![
+                            key
+                            ]
+                          }
+                          onChange={(e) => {
+                            //@ts-ignore
+                            setTabsState((old: TabsState) => {
+                              let newTabsState = _.cloneDeep(old);
+                              newTabsState[
+                                id.current
+                              ].formKeysData.input_keys![key] =
+                                e.target.value;
+                              return newTabsState;
+                            });
+                          }}
+                          disabled={chatKey === key}
+                          placeholder="Enter text..."
+                        ></Textarea>
+                      </div>
+                    </AccordionComponent>
+                  </div>
+                ))
                 : null}
               {tabsState[id.current].formKeysData.memory_keys!.map(
                 (key, index) => (
@@ -513,7 +522,7 @@ export default function FormModal({
                           <div className="-mb-1">
                             <ToggleShadComponent
                               enabled={chatKey === key}
-                              setEnabled={() => {}}
+                              setEnabled={() => { }}
                               size="small"
                               disabled={true}
                             />