From fbc926a4360be84c19a12eec768b5f2dc151b81e Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Tue, 23 Jan 2024 16:07:18 -0300
Subject: [PATCH] Add flowPool and buildFlow functionality

---
 .../src/CustomNodes/GenericNode/index.tsx         | 15 +++++++++++----
 src/frontend/src/types/zustand/flow/index.ts      |  3 +++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index ea9f49845..61ba66967 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -30,6 +30,8 @@ export default function GenericNode({
 }): JSX.Element {
   const types = useTypesStore((state) => state.types);
   const deleteNode = useFlowStore((state) => state.deleteNode);
+  const flowPool = useFlowStore((state) => state.flowPool);
+  const buildFlow = useFlowStore((state) => state.buildFlow);
   const setNode = useFlowStore((state) => state.setNode);
   const name = nodeIconsLucide[data.type] ? data.type : types[data.type];
   const [inputName, setInputName] = useState(false);
@@ -92,16 +94,18 @@ export default function GenericNode({
     setNodeName(data.node!.display_name);
   }, [data.node!.display_name]);
 
-  // New useEffect to watch for changes in sseData and update validation status
   useEffect(() => {
-    const relevantData = sseData[data.id];
+    const relevantData =
+      flowPool[data.id] && flowPool[data.id]?.length > 0
+        ? flowPool[data.id][flowPool[data.id].length - 1]
+        : null;
     if (relevantData) {
       // Extract validation information from relevantData and update the validationStatus state
       setValidationStatus(relevantData);
     } else {
       setValidationStatus(null);
     }
-  }, [sseData, data.id]);
+  }, [flowPool, data.id]);
 
   const showNode = data.showNode ?? true;
 
@@ -317,7 +321,10 @@ export default function GenericNode({
             </div>
 
             {showNode && (
-              <div className="round-button-div">
+              <div
+                className="round-button-div"
+                onClick={() => buildFlow(data.id)}
+              >
                 <div>
                   <Tooltip
                     title={
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index 92d520934..1dcb72fda 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -24,6 +24,8 @@ export type FlowPoolObjectType = {
   valid: boolean;
   params: any;
   data: { artifacts: any; results: any | ChatOutputType | chatInputType };
+  duration: string;
+  progress: number;
   id: string;
 };
 
@@ -85,4 +87,5 @@ export type FlowStoreType = {
   getFilterEdge: any[];
   onConnect: (connection: Connection) => void;
   unselectAll: () => void;
+  buildFlow: (nodeId?: string) => void;
 };