feat: ui build in one single http request (#3020)

* feat: ui build in one single http request

* fix use session_id

* fix frozen

* [autofix.ci] apply automated fixes

* prettier

* add tests

* add tests

* fix mypy

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

src/frontend/src/controllers/API/api.tsx

@@ -16,11 +16,12 @@ const api: AxiosInstance = axios.create({
   baseURL: "",
 });
 
+const cookies = new Cookies();
 function ApiInterceptor() {
   const autoLogin = useAuthStore((state) => state.autoLogin);
   const setErrorData = useAlertStore((state) => state.setErrorData);
   let { accessToken, authenticationErrorCount } = useContext(AuthContext);
-  const cookies = new Cookies();
+
   const setSaveLoading = useFlowsManagerStore((state) => state.setSaveLoading);
   const { mutate: mutationLogout } = useLogout();
   const { mutate: mutationRenewAccessToken } = useRefreshAccessToken();
@@ -205,4 +206,83 @@ function ApiInterceptor() {
   return null;
 }
 
-export { ApiInterceptor, api };
+export type StreamingRequestParams = {
+  method: string;
+  url: string;
+  onData: (event: object) => Promise<boolean>;
+  body?: object;
+  onError?: (statusCode: number) => void;
+};
+
+async function performStreamingRequest({
+  method,
+  url,
+  onData,
+  body,
+  onError,
+}: StreamingRequestParams) {
+  let headers = {
+    "Content-Type": "application/json",
+    // this flag is fundamental to ensure server stops tasks when client disconnects
+    Connection: "close",
+  };
+  const accessToken = cookies.get(LANGFLOW_ACCESS_TOKEN);
+  if (accessToken) {
+    headers["Authorization"] = `Bearer ${accessToken}`;
+  }
+  const controller = new AbortController();
+  const params = {
+    method: method,
+    headers: headers,
+    signal: controller.signal,
+  };
+  if (body) {
+    params["body"] = JSON.stringify(body);
+  }
+  let current: string[] = [];
+  let textDecoder = new TextDecoder();
+  const response = await fetch(url, params);
+  if (!response.ok) {
+    if (onError) {
+      onError(response.status);
+    } else {
+      throw new Error("error in streaming request");
+    }
+  }
+  if (response.body === null) {
+    return;
+  }
+  for await (const chunk of response.body) {
+    const decodedChunk = await textDecoder.decode(chunk);
+    let all = decodedChunk.split("\n\n");
+    for (const string of all) {
+      if (string.endsWith("}")) {
+        const allString = current.join("") + string;
+        let data: object;
+        try {
+          data = JSON.parse(allString);
+          current = [];
+        } catch (e) {
+          current.push(string);
+          continue;
+        }
+        const shouldContinue = await onData(data);
+        if (!shouldContinue) {
+          controller.abort();
+          return;
+        }
+      } else {
+        current.push(string);
+      }
+    }
+  }
+  if (current.length > 0) {
+    const allString = current.join("");
+    if (allString) {
+      const data = JSON.parse(current.join(""));
+      await onData(data);
+    }
+  }
+}
+
+export { ApiInterceptor, api, performStreamingRequest };

src/frontend/src/modals/IOModal/components/chatView/chatMessage/index.tsx

@@ -38,7 +38,6 @@ export default function ChatMessage({
   const [chatMessage, setChatMessage] = useState(chatMessageString);
   const [isStreaming, setIsStreaming] = useState(false);
   const eventSource = useRef<EventSource | undefined>(undefined);
-  const updateFlowPool = useFlowStore((state) => state.updateFlowPool);
   const setErrorData = useAlertStore((state) => state.setErrorData);
   const chatMessageRef = useRef(chatMessage);
 

src/frontend/src/stores/flowStore.ts

@@ -29,7 +29,7 @@ import {
   targetHandleType,
 } from "../types/flow";
 import { FlowStoreType, VertexLayerElementType } from "../types/zustand/flow";
-import { buildVertices } from "../utils/buildUtils";
+import { buildFlowVerticesWithFallback } from "../utils/buildUtils";
 import {
   checkChatInput,
   checkOldComponents,
@@ -607,20 +607,8 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       );
 
       useFlowStore.getState().updateBuildStatus([vertexBuildData.id], status);
-
-      const verticesIds = get().verticesBuild?.verticesIds;
-      const newFlowBuildStatus = { ...get().flowBuildStatus };
-      // filter out the vertices that are not status
-
-      const verticesToUpdate = verticesIds?.filter(
-        (id) => newFlowBuildStatus[id]?.status !== BuildStatus.BUILT,
-      );
-
-      if (verticesToUpdate) {
-        useFlowStore.getState().updateBuildStatus(verticesToUpdate, status);
-      }
     }
-    await buildVertices({
+    await buildFlowVerticesWithFallback({
       input_value,
       files,
       flowId: currentFlow!.id,
@@ -672,8 +660,8 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
         useFlowStore.getState().updateBuildStatus(idList, BuildStatus.BUILDING);
       },
       onValidateNodes: validateSubgraph,
-      nodes: !get().onFlowPage ? get().nodes : undefined,
-      edges: !get().onFlowPage ? get().edges : undefined,
+      nodes: get().onFlowPage ? get().nodes : undefined,
+      edges: get().onFlowPage ? get().edges : undefined,
     });
     get().setIsBuilding(false);
     get().setLockChat(false);
@@ -690,7 +678,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     vertices: {
       verticesIds: string[];
       verticesLayers: VertexLayerElementType[][];
-      runId: string;
+      runId?: string;
       verticesToRun: string[];
     } | null,
   ) => {

src/frontend/src/types/zustand/flow/index.ts

@@ -147,7 +147,7 @@ export type FlowStoreType = {
     vertices: {
       verticesIds: string[];
       verticesLayers: VertexLayerElementType[][];
-      runId: string;
+      runId?: string;
       verticesToRun: string[];
     } | null,
   ) => void;
@@ -156,7 +156,7 @@ export type FlowStoreType = {
   verticesBuild: {
     verticesIds: string[];
     verticesLayers: VertexLayerElementType[][];
-    runId: string;
+    runId?: string;
     verticesToRun: string[];
   } | null;
   updateBuildStatus: (nodeId: string[], status: BuildStatus) => void;

src/frontend/src/utils/buildUtils.ts

@@ -1,3 +1,5 @@
+import { BASE_URL_API } from "@/constants/constants";
+import { performStreamingRequest } from "@/controllers/API/api";
 import { AxiosError } from "axios";
 import { Edge, Node } from "reactflow";
 import { BuildStatus } from "../constants/enums";
@@ -66,7 +68,7 @@ export async function updateVerticesOrder(
 ): Promise<{
   verticesLayers: VertexLayerElementType[][];
   verticesIds: string[];
-  runId: string;
+  runId?: string;
   verticesToRun: string[];
 }> {
   return new Promise(async (resolve, reject) => {
@@ -115,6 +117,176 @@ export async function updateVerticesOrder(
   });
 }
 
+export async function buildFlowVerticesWithFallback(
+  params: BuildVerticesParams,
+) {
+  try {
+    return await buildFlowVertices(params);
+  } catch (e: any) {
+    if (e.message === "endpoint not available") {
+      return await buildVertices(params);
+    }
+    throw e;
+  }
+}
+
+const MIN_VISUAL_BUILD_TIME_MS = 300;
+
+export async function buildFlowVertices({
+  flowId,
+  input_value,
+  files,
+  startNodeId,
+  stopNodeId,
+  onGetOrderSuccess,
+  onBuildUpdate,
+  onBuildComplete,
+  onBuildError,
+  onBuildStart,
+  onValidateNodes,
+  nodes,
+  edges,
+  setLockChat,
+}: BuildVerticesParams) {
+  let url = `${BASE_URL_API}build/${flowId}/flow?`;
+  if (startNodeId) {
+    url = `${url}&start_component_id=${startNodeId}`;
+  }
+  if (stopNodeId) {
+    url = `${url}&stop_component_id=${stopNodeId}`;
+  }
+  const postData = {};
+  if (typeof input_value !== "undefined") {
+    postData["inputs"] = { input_value: input_value };
+  }
+  if (files) {
+    postData["files"] = files;
+  }
+  if (nodes) {
+    postData["data"] = {
+      nodes,
+      edges,
+    };
+  }
+
+  const buildResults: Array<boolean> = [];
+
+  const verticesStartTimeMs: Map<string, number> = new Map();
+
+  const onEvent = async (type, data): Promise<boolean> => {
+    const onStartVertices = (ids: Array<string>) => {
+      useFlowStore.getState().updateBuildStatus(ids, BuildStatus.TO_BUILD);
+      if (onBuildStart)
+        onBuildStart(ids.map((id) => ({ id: id, reference: id })));
+      ids.forEach((id) => verticesStartTimeMs.set(id, Date.now()));
+    };
+    switch (type) {
+      case "vertices_sorted": {
+        const verticesToRun = data.to_run;
+        const verticesIds = data.ids;
+
+        onStartVertices(verticesIds);
+
+        let verticesLayers: Array<Array<VertexLayerElementType>> =
+          verticesIds.map((id: string) => {
+            return [{ id: id, reference: id }];
+          });
+
+        useFlowStore.getState().updateVerticesBuild({
+          verticesLayers,
+          verticesIds,
+          verticesToRun,
+        });
+        if (onValidateNodes) {
+          try {
+            onValidateNodes(data.to_run);
+            if (onGetOrderSuccess) onGetOrderSuccess();
+            useFlowStore.getState().setIsBuilding(true);
+            return true;
+          } catch (e) {
+            useFlowStore.getState().setIsBuilding(false);
+            setLockChat && setLockChat(false);
+            return false;
+          }
+        }
+        return true;
+      }
+      case "end_vertex": {
+        const buildData = data.build_data;
+        const startTimeMs = verticesStartTimeMs.get(buildData.id);
+        if (startTimeMs) {
+          const delta = Date.now() - startTimeMs;
+          if (delta < MIN_VISUAL_BUILD_TIME_MS) {
+            // this is a visual trick to make the build process look more natural
+            await new Promise((resolve) =>
+              setTimeout(resolve, MIN_VISUAL_BUILD_TIME_MS - delta),
+            );
+          }
+        }
+
+        if (onBuildUpdate) {
+          if (!buildData.valid) {
+            // lots is a dictionary with the key the output field name and the value the log object
+            // logs: { [key: string]: { message: any; type: string }[] };
+            const errorMessages = Object.keys(buildData.data.outputs).map(
+              (key) => {
+                const outputs = buildData.data.outputs[key];
+                if (Array.isArray(outputs)) {
+                  return outputs
+                    .filter((log) => isErrorLogType(log.message))
+                    .map((log) => log.message.errorMessage);
+                }
+                if (!isErrorLogType(outputs.message)) {
+                  return [];
+                }
+                return [outputs.message.errorMessage];
+              },
+            );
+            onBuildError!("Error Building Component", errorMessages, [
+              { id: buildData.id },
+            ]);
+            onBuildUpdate(buildData, BuildStatus.ERROR, "");
+            buildResults.push(false);
+            return false;
+          } else {
+            onBuildUpdate(buildData, BuildStatus.BUILT, "");
+            buildResults.push(true);
+          }
+        }
+        if (buildData.next_vertices_ids) {
+          onStartVertices(buildData.next_vertices_ids);
+        }
+        return true;
+      }
+      case "end": {
+        const allNodesValid = buildResults.every((result) => result);
+        onBuildComplete!(allNodesValid);
+        useFlowStore.getState().setIsBuilding(false);
+        return true;
+      }
+      default:
+        return true;
+    }
+    return true;
+  };
+  return performStreamingRequest({
+    method: "POST",
+    url,
+    body: postData,
+    onData: async (event) => {
+      const type = event["event"];
+      const data = event["data"];
+      return await onEvent(type, data);
+    },
+    onError: (statusCode) => {
+      if (statusCode === 404) {
+        throw new Error("endpoint not available");
+      }
+      throw new Error("error in streaming request");
+    },
+  });
+}
+
 export async function buildVertices({
   flowId,
   input_value,
@@ -252,6 +424,7 @@ export async function buildVertices({
     useFlowStore.getState().setIsBuilding(false);
   }
 }
+
 async function buildVertex({
   flowId,
   id,