From 67814608697e0b06f654d22311872455bb04cf87 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Sun, 3 Mar 2024 14:40:43 -0300
Subject: [PATCH] Refactor buildVertices process one layer at a time

---
 src/frontend/src/stores/flowStore.ts | 23 ++++++---
 src/frontend/src/types/api/index.ts  |  2 +
 src/frontend/src/utils/buildUtils.ts | 77 +++++++++++++++-------------
 3 files changed, 61 insertions(+), 41 deletions(-)

diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index a373123e8..9bdd5c9a0 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -223,12 +223,10 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       selection.nodes.some((node) => node.data.type === "ChatInput") &&
       checkChatInput(get().nodes)
     ) {
-      useAlertStore
-        .getState()
-        .setErrorData({
-          title: "Error pasting components",
-          list: ["You can only have one ChatInput component in the flow"],
-        });
+      useAlertStore.getState().setErrorData({
+        title: "Error pasting components",
+        list: ["You can only have one ChatInput component in the flow"],
+      });
       return;
     }
     let minimumX = Infinity;
@@ -451,10 +449,23 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       if (vertexBuildData && vertexBuildData.inactive_vertices) {
         get().removeFromVerticesBuild(vertexBuildData.inactive_vertices);
       }
+      get().verticesBuild &&
+        get().updateVerticesBuild({
+          verticesIds: [
+            ...get().verticesBuild!.verticesIds,
+            vertexBuildData.id,
+          ],
+          verticesLayers: [
+            ...get().verticesBuild!.verticesLayers,
+            vertexBuildData.next_vertices_ids,
+          ],
+          runId: vertexBuildData.run_id,
+        });
       get().addDataToFlowPool(
         { ...vertexBuildData, buildId },
         vertexBuildData.id
       );
+
       useFlowStore.getState().updateBuildStatus([vertexBuildData.id], status);
     }
     await buildVertices({
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index 6663692dc..06f30fb2c 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -140,7 +140,9 @@ export type VerticesOrderTypeAPI = {
 
 export type VertexBuildTypeAPI = {
   id: string;
+  next_vertices_ids: Array<string>;
   inactive_vertices: Array<string> | null;
+  run_id: string;
   valid: boolean;
   params: string;
   data: VertexDataTypeAPI;
diff --git a/src/frontend/src/utils/buildUtils.ts b/src/frontend/src/utils/buildUtils.ts
index 77fafd029..ce3a542ab 100644
--- a/src/frontend/src/utils/buildUtils.ts
+++ b/src/frontend/src/utils/buildUtils.ts
@@ -125,50 +125,57 @@ export async function buildVertices({
 
   useFlowStore.getState().updateBuildStatus(verticesIds, BuildStatus.TO_BUILD);
   useFlowStore.getState().setIsBuilding(true);
-
+  let currentLayerIndex = 0; // Start with the first layer
   // Set each vertex state to building
   const buildResults: Array<boolean> = [];
   console.log(verticesLayers);
-  for (const layer of verticesLayers) {
-    if (onBuildStart) onBuildStart(layer);
-    for (const id of layer) {
-      // Check if id is in the list of inactive nodes
-      if (!verticesIds.includes(id) && onBuildUpdate) {
-        // If it is, skip building and set the state to inactive
-        onBuildUpdate(getInactiveVertexData(id), BuildStatus.INACTIVE, runId);
-        buildResults.push(false);
-        continue;
-      }
-      await buildVertex({
-        flowId,
-        id,
-        input_value,
-        onBuildUpdate: (data: VertexBuildTypeAPI, status: BuildStatus) => {
-          if (onBuildUpdate) onBuildUpdate(data, status, runId);
-        },
-        onBuildError,
-        verticesIds,
-        buildResults,
-        stopBuild: () => {
-          stop = true;
-        },
-      });
-      if (stop) {
-        break;
-      }
-    }
+  while (currentLayerIndex < verticesLayers.length) {
+    const currentLayer = verticesLayers[currentLayerIndex];
+    if (onBuildStart) onBuildStart(currentLayer);
+    await Promise.all(
+      currentLayer.map(async (vertexId) => {
+        // Check if id is in the list of inactive nodes
+        if (!verticesIds.includes(vertexId) && onBuildUpdate) {
+          // If it is, skip building and set the state to inactive
+          onBuildUpdate(
+            getInactiveVertexData(vertexId),
+            BuildStatus.INACTIVE,
+            runId
+          );
+          buildResults.push(false);
+          return;
+        }
+        await buildVertex({
+          flowId,
+          id: vertexId,
+          input_value,
+          onBuildUpdate: (data: VertexBuildTypeAPI, status: BuildStatus) => {
+            if (onBuildUpdate) onBuildUpdate(data, status, runId);
+          },
+          onBuildError,
+          verticesIds,
+          buildResults,
+          stopBuild: () => {
+            stop = true;
+          },
+        });
+        if (stop) {
+          return;
+        }
+      })
+    );
+
     if (stop) {
       break;
     }
-  }
 
-  if (onBuildComplete) {
-    const allNodesValid = buildResults.every((result) => result);
-    onBuildComplete(allNodesValid);
-    useFlowStore.getState().setIsBuilding(false);
+    if (onBuildComplete) {
+      const allNodesValid = buildResults.every((result) => result);
+      onBuildComplete(allNodesValid);
+      useFlowStore.getState().setIsBuilding(false);
+    }
   }
 }
-
 async function buildVertex({
   flowId,
   id,