diff --git a/src/frontend/src/App.tsx b/src/frontend/src/App.tsx
index 7e4d80f4d..9b7ec8b75 100644
--- a/src/frontend/src/App.tsx
+++ b/src/frontend/src/App.tsx
@@ -16,17 +16,19 @@ import {
   FETCH_ERROR_MESSAGE,
 } from "./constants/constants";
 import { AuthContext } from "./contexts/authContext";
-import { locationContext } from "./contexts/locationContext";
 import { getHealth, getRepoStars, getVersion } from "./controllers/API";
 import Router from "./routes";
 import useAlertStore from "./stores/alertStore";
 import { useTypesStore } from "./stores/typesStore";
 import { useDarkStore } from "./stores/darkStore";
 import useFlowsManagerStore from "./stores/flowsManagerStore";
+import { useLocationStore } from "./stores/locationStore";
 
 export default function App() {
-  let { setCurrent, setShowSideBar, setIsStackedOpen } =
-    useContext(locationContext);
+  const setCurrent = useLocationStore((state) => state.setCurrent);
+  const setShowSideBar = useLocationStore((state) => state.setShowSideBar);
+  const setIsStackedOpen = useLocationStore((state) => state.setIsStackedOpen);
+
   let location = useLocation();
   useEffect(() => {
     setCurrent(location.pathname.replace(/\/$/g, "").split("/"));
diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index 339eb73ac..f4a284327 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -6,7 +6,6 @@ import IconComponent from "../../components/genericIconComponent";
 import InputComponent from "../../components/inputComponent";
 import { Textarea } from "../../components/ui/textarea";
 import { priorityFields } from "../../constants/constants";
-import { useSSE } from "../../contexts/SSEContext";
 import NodeToolbarComponent from "../../pages/FlowPage/components/nodeToolbarComponent";
 import useFlowStore from "../../stores/flowStore";
 import { validationStatusType } from "../../types/components";
@@ -80,7 +79,8 @@ export default function GenericNode({
   }, [data, data.node]);
 
   // State for outline color
-  const { sseData, isBuilding } = useSSE();
+  const sseData = useFlowStore((state) => state.sseData);
+  const isBuilding = useFlowStore((state) => state.isBuilding);
 
   useEffect(() => {
     setNodeDescription(data.node!.description);
diff --git a/src/frontend/src/components/chatComponent/buildTrigger/index.tsx b/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
index 1dcc327ee..3a1a47d9a 100644
--- a/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
+++ b/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
@@ -1,18 +1,16 @@
 import { Transition } from "@headlessui/react";
-import { useContext, useState } from "react";
+import { useState } from "react";
 import Loading from "../../../components/ui/loading";
-import { useSSE } from "../../../contexts/SSEContext";
 import { postBuildInit } from "../../../controllers/API";
 import { FlowType } from "../../../types/flow";
 
 import useAlertStore from "../../../stores/alertStore";
 import useFlowStore from "../../../stores/flowStore";
-import useFlowsManagerStore from "../../../stores/flowsManagerStore";
 import { parsedDataType } from "../../../types/components";
-import { FlowState } from "../../../types/tabs";
 import { validateNodes } from "../../../utils/reactflowUtils";
 import RadialProgressComponent from "../../RadialProgress";
 import IconComponent from "../../genericIconComponent";
+import { stat } from "fs";
 
 export default function BuildTrigger({
   open,
@@ -24,17 +22,17 @@ export default function BuildTrigger({
   setIsBuilt: any;
   isBuilt: boolean;
 }): JSX.Element {
-  const { updateSSEData, isBuilding, setIsBuilding, sseData } = useSSE();
+  const updateSSEData = useFlowStore((state) => state.updateSSEData);
+  const isBuilding = useFlowStore((state) => state.isBuilding);
+  const setIsBuilding = useFlowStore((state) => state.setIsBuilding);
   const nodes = useFlowStore((state) => state.nodes);
   const edges = useFlowStore((state) => state.edges);
   const setErrorData = useAlertStore((state) => state.setErrorData);
   const setSuccessData = useAlertStore((state) => state.setSuccessData);
-  const setCurrentFlowState = useFlowsManagerStore(
-    (state) => state.setCurrentFlowState
+  const setFlowState = useFlowStore(
+    (state) => state.setFlowState
   );
-  const saveFlow = useFlowsManagerStore((state) => state.saveFlow);
 
-  const [isIconTouched, setIsIconTouched] = useState(false);
   const eventClick = isBuilding ? "pointer-events-none" : "";
   const [progress, setProgress] = useState(0);
 
@@ -79,7 +77,6 @@ export default function BuildTrigger({
   }
   async function streamNodeData(flow: FlowType) {
     // Step 1: Make a POST request to send the flow data and receive a unique session ID
-    const id = saveFlow(flow, true);
     const response = await postBuildInit(flow);
     const { flowId } = response.data;
     // Step 2: Use the session ID to establish an SSE connection using EventSource
@@ -102,7 +99,7 @@ export default function BuildTrigger({
           // If the event is a log, log it
           setSuccessData({ title: parsedData.log });
         } else if (parsedData.input_keys !== undefined) {
-          setCurrentFlowState(parsedData);
+          setFlowState(parsedData);
         } else {
           // Otherwise, process the data
           const isValid = processStreamResult(parsedData);
@@ -148,14 +145,6 @@ export default function BuildTrigger({
     }
   }
 
-  const handleMouseEnter = () => {
-    setIsIconTouched(true);
-  };
-
-  const handleMouseLeave = () => {
-    setIsIconTouched(false);
-  };
-
   return (
     <Transition
       show={!open}
@@ -173,8 +162,6 @@ export default function BuildTrigger({
           onClick={() => {
             handleBuild(flow);
           }}
-          onMouseEnter={handleMouseEnter}
-          onMouseLeave={handleMouseLeave}
         >
           <button>
             <div className="round-button-div">
diff --git a/src/frontend/src/components/chatComponent/index.tsx b/src/frontend/src/components/chatComponent/index.tsx
index 80d6fb991..0ee71c997 100644
--- a/src/frontend/src/components/chatComponent/index.tsx
+++ b/src/frontend/src/components/chatComponent/index.tsx
@@ -13,11 +13,9 @@ import useFlowsManagerStore from "../../stores/flowsManagerStore";
 
 export default function Chat({ flow }: ChatType): JSX.Element {
   const [open, setOpen] = useState(false);
-  const [canOpen, setCanOpen] = useState(false);
   const isBuilt = useFlowStore((state) => state.isBuilt);
   const setIsBuilt = useFlowStore((state) => state.setIsBuilt);
-  const currentFlowState = useFlowsManagerStore((state) => state.currentFlowState);
-
+  const flowState = useFlowStore((state) => state.flowState);
   useEffect(() => {
     const handleKeyDown = (event: KeyboardEvent) => {
       if (
@@ -58,18 +56,8 @@ export default function Chat({ flow }: ChatType): JSX.Element {
     ) {
       setIsBuilt(false);
     }
-    if (
-      currentFlowState &&
-      currentFlowState.formKeysData &&
-      currentFlowState.formKeysData.input_keys !== null
-    ) {
-      setCanOpen(true);
-    } else {
-      setCanOpen(false);
-    }
-
     prevNodesRef.current = currentNodes;
-  }, [currentFlowState, flow.id]);
+  }, [flowState, flow.id]);
 
   return (
     <>
@@ -81,9 +69,8 @@ export default function Chat({ flow }: ChatType): JSX.Element {
           isBuilt={isBuilt}
         />
         {isBuilt &&
-          currentFlowState &&
-          currentFlowState.formKeysData &&
-          canOpen && (
+          flowState &&
+          !!flowState?.input_keys && (
             <FormModal
               key={flow.id}
               flow={flow}
@@ -92,7 +79,7 @@ export default function Chat({ flow }: ChatType): JSX.Element {
             />
           )}
         <ChatTrigger
-          canOpen={canOpen}
+          canOpen={!!flowState?.input_keys}
           open={open}
           setOpen={setOpen}
           isBuilt={isBuilt}
diff --git a/src/frontend/src/components/ui/accordion.tsx b/src/frontend/src/components/ui/accordion.tsx
index 27dbba93b..c77130ea9 100644
--- a/src/frontend/src/components/ui/accordion.tsx
+++ b/src/frontend/src/components/ui/accordion.tsx
@@ -24,16 +24,16 @@ const AccordionTrigger = React.forwardRef<
   React.ComponentPropsWithoutRef<typeof AccordionPrimitive.Trigger>
 >(({ className, children, ...props }, ref) => (
   <AccordionPrimitive.Header className="flex">
-    <AccordionPrimitive.Trigger
-      ref={ref}
-      className={cn(
-        "flex flex-1 items-center justify-between py-4 text-sm font-medium transition-all [&[data-state=open]>svg]:rotate-180",
-        className
-      )}
-      {...props}
-    >
-      {children}
-      <ChevronDownIcon className="h-4 w-4 text-muted-foreground transition-transform duration-200" />
+    <AccordionPrimitive.Trigger asChild ref={ref} {...props}>
+      <div
+        className={cn(
+          "cursor-pointer flex flex-1 items-center justify-between py-4 text-sm font-medium transition-all [&[data-state=open]>svg]:rotate-180",
+          className
+        )}
+      >
+        {children}
+        <ChevronDownIcon className="h-4 w-4 text-muted-foreground transition-transform duration-200" />
+      </div>
     </AccordionPrimitive.Trigger>
   </AccordionPrimitive.Header>
 ));
diff --git a/src/frontend/src/contexts/SSEContext.tsx b/src/frontend/src/contexts/SSEContext.tsx
deleted file mode 100644
index 9db32310b..000000000
--- a/src/frontend/src/contexts/SSEContext.tsx
+++ /dev/null
@@ -1,34 +0,0 @@
-import { createContext, useCallback, useContext, useState } from "react";
-
-const initialValue = {
-  updateSSEData: ({}) => {},
-  sseData: {},
-  isBuilding: false,
-  setIsBuilding: (isBuilding: boolean) => {},
-};
-
-const SSEContext = createContext(initialValue);
-
-export function useSSE() {
-  return useContext(SSEContext);
-}
-
-export function SSEProvider({ children }) {
-  const [sseData, setSSEData] = useState({});
-  const [isBuilding, setIsBuilding] = useState(false);
-
-  const updateSSEData = useCallback((newData: any) => {
-    setSSEData((prevData) => ({
-      ...prevData,
-      ...newData,
-    }));
-  }, []);
-
-  return (
-    <SSEContext.Provider
-      value={{ sseData, updateSSEData, isBuilding, setIsBuilding }}
-    >
-      {children}
-    </SSEContext.Provider>
-  );
-}
diff --git a/src/frontend/src/contexts/index.tsx b/src/frontend/src/contexts/index.tsx
index 7090b7c40..f6660e7fe 100644
--- a/src/frontend/src/contexts/index.tsx
+++ b/src/frontend/src/contexts/index.tsx
@@ -3,9 +3,7 @@ import { BrowserRouter } from "react-router-dom";
 import { ReactFlowProvider } from "reactflow";
 import { TooltipProvider } from "../components/ui/tooltip";
 import { ApiInterceptor } from "../controllers/API/api";
-import { SSEProvider } from "./SSEContext";
 import { AuthProvider } from "./authContext";
-import { LocationProvider } from "./locationContext";
 
 export default function ContextWrapper({ children }: { children: ReactNode }) {
   //element to wrap all context
@@ -15,12 +13,8 @@ export default function ContextWrapper({ children }: { children: ReactNode }) {
         <AuthProvider>
           <TooltipProvider>
             <ReactFlowProvider>
-              <LocationProvider>
                 <ApiInterceptor />
-                <SSEProvider>
                   {children}
-                </SSEProvider>
-              </LocationProvider>
             </ReactFlowProvider>
           </TooltipProvider>
         </AuthProvider>
diff --git a/src/frontend/src/contexts/locationContext.tsx b/src/frontend/src/contexts/locationContext.tsx
deleted file mode 100644
index b0a9ef8c5..000000000
--- a/src/frontend/src/contexts/locationContext.tsx
+++ /dev/null
@@ -1,50 +0,0 @@
-import { createContext, ReactNode, useState } from "react";
-import { locationContextType } from "../types/typesContext";
-
-//initial value for location context
-const initialValue = {
-  //actual
-  current: window.location.pathname.replace(/\/$/g, "").split("/"),
-  isStackedOpen:
-    window.innerWidth > 1024 && window.location.pathname.split("/")[1]
-      ? true
-      : false,
-  setCurrent: () => {},
-  setIsStackedOpen: () => {},
-  showSideBar: window.location.pathname.split("/")[1] ? true : false,
-  setShowSideBar: () => {},
-  extraNavigation: { title: "" },
-  setExtraNavigation: () => {},
-  extraComponent: <></>,
-  setExtraComponent: () => {},
-};
-
-export const locationContext = createContext<locationContextType>(initialValue);
-
-export function LocationProvider({ children }: { children: ReactNode }) {
-  const [current, setCurrent] = useState(initialValue.current);
-  const [isStackedOpen, setIsStackedOpen] = useState(
-    initialValue.isStackedOpen
-  );
-  const [showSideBar, setShowSideBar] = useState(initialValue.showSideBar);
-  const [extraNavigation, setExtraNavigation] = useState({ title: "" });
-  const [extraComponent, setExtraComponent] = useState(<></>);
-  return (
-    <locationContext.Provider
-      value={{
-        isStackedOpen,
-        setIsStackedOpen,
-        current,
-        setCurrent,
-        showSideBar,
-        setShowSideBar,
-        extraNavigation,
-        setExtraNavigation,
-        extraComponent,
-        setExtraComponent,
-      }}
-    >
-      {children}
-    </locationContext.Provider>
-  );
-}
diff --git a/src/frontend/src/contexts/typesContext.tsx b/src/frontend/src/contexts/typesContext.tsx
deleted file mode 100644
index 16deecf00..000000000
--- a/src/frontend/src/contexts/typesContext.tsx
+++ /dev/null
@@ -1,105 +0,0 @@
-import _ from "lodash";
-import {
-  createContext,
-  ReactNode,
-  useContext,
-  useEffect,
-  useState,
-} from "react";
-import { getAll, getHealth } from "../controllers/API";
-import useAlertStore from "../stores/alertStore";
-import { APIKindType } from "../types/api";
-import { typesContextType } from "../types/typesContext";
-import { AuthContext } from "./authContext";
-
-//context to share types adn functions from nodes to flow
-
-const initialValue: typesContextType = {
-  types: {},
-  setTypes: () => {},
-  templates: {},
-  setTemplates: () => {},
-  data: {},
-  setData: () => {},
-  getTypes: () => {},
-  setFetchError: () => {},
-  fetchError: false,
-  setFilterEdge: (filter) => {},
-  getFilterEdge: [],
-};
-
-export const typesContext = createContext<typesContextType>(initialValue);
-
-export function TypesProvider({ children }: { children: ReactNode }) {
-  const [types, setTypes] = useState({});
-  const [templates, setTemplates] = useState({});
-  const [data, setData] = useState({});
-  const [fetchError, setFetchError] = useState(false);
-  const setLoading = useAlertStore((state) => state.setLoading);
-  const [getFilterEdge, setFilterEdge] = useState([]);
-
-  async function getTypes(): Promise<void> {
-    // We will keep a flag to handle the case where the component is unmounted before the API call resolves.
-    let isMounted = true;
-    try {
-      const result = await getAll();
-      // Make sure to only update the state if the component is still mounted.
-      if (isMounted && result?.status === 200) {
-        setLoading(false);
-        let { data } = _.cloneDeep(result);
-        setData((old) => ({ ...old, ...data }));
-        setTemplates(
-          Object.keys(data).reduce((acc, curr) => {
-            Object.keys(data[curr]).forEach((c: keyof APIKindType) => {
-              //prevent wrong overwriting of the component template by a group of the same type
-              if (!data[curr][c].flow) acc[c] = data[curr][c];
-            });
-            return acc;
-          }, {})
-        );
-        // Set the types by reducing over the keys of the result data and updating the accumulator.
-        setTypes(
-          // Reverse the keys so the tool world does not overlap
-          Object.keys(data)
-            .reverse()
-            .reduce((acc, curr) => {
-              Object.keys(data[curr]).forEach((c: keyof APIKindType) => {
-                acc[c] = curr;
-                // Add the base classes to the accumulator as well.
-                data[curr][c].base_classes?.forEach((b) => {
-                  acc[b] = curr;
-                });
-              });
-              return acc;
-            }, {})
-        );
-      }
-    } catch (error) {
-      console.error("An error has occurred while fetching types.");
-      console.log(error);
-      await getHealth().catch((e) => {
-        setFetchError(true);
-      });
-    }
-  }
-
-  return (
-    <typesContext.Provider
-      value={{
-        types,
-        setTypes,
-        setTemplates,
-        templates,
-        data,
-        setData,
-        getTypes,
-        fetchError,
-        setFetchError,
-        setFilterEdge,
-        getFilterEdge,
-      }}
-    >
-      {children}
-    </typesContext.Provider>
-  );
-}
diff --git a/src/frontend/src/modals/ApiModal/index.tsx b/src/frontend/src/modals/ApiModal/index.tsx
index a76269107..861194ffe 100644
--- a/src/frontend/src/modals/ApiModal/index.tsx
+++ b/src/frontend/src/modals/ApiModal/index.tsx
@@ -31,6 +31,7 @@ import {
   tabsArray,
 } from "../../utils/utils";
 import BaseModal from "../baseModal";
+import useFlowStore from "../../stores/flowStore";
 
 const ApiModal = forwardRef(
   (
@@ -49,18 +50,18 @@ const ApiModal = forwardRef(
     const tweak = useRef<tweakType>([]);
     const tweaksList = useRef<string[]>([]);
     const [getTweak, setTweak] = useState<tweakType>([]);
-    const currentFlowState = useFlowsManagerStore(
-      (state) => state.currentFlowState
+    const flowState = useFlowStore(
+      (state) => state.flowState
     );
     const pythonApiCode = getPythonApiCode(
       flow,
       autoLogin,
       tweak.current,
-      currentFlowState
+      flowState
     );
-    const curl_code = getCurlCode(flow, autoLogin, tweak.current, currentFlowState);
-    const pythonCode = getPythonCode(flow, tweak.current, currentFlowState);
-    const widgetCode = getWidgetCode(flow, autoLogin, currentFlowState);
+    const curl_code = getCurlCode(flow, autoLogin, tweak.current, flowState);
+    const pythonCode = getPythonCode(flow, tweak.current, flowState);
+    const widgetCode = getWidgetCode(flow, autoLogin, flowState);
     const tweaksCode = buildTweaks(flow);
     const codesArray = [
       curl_code,
@@ -171,11 +172,11 @@ const ApiModal = forwardRef(
         flow,
         autoLogin,
         tweak.current,
-        currentFlowState
+        flowState
       );
-      const curl_code = getCurlCode(flow, autoLogin, tweak.current, currentFlowState);
-      const pythonCode = getPythonCode(flow, tweak.current, currentFlowState);
-      const widgetCode = getWidgetCode(flow, autoLogin, currentFlowState);
+      const curl_code = getCurlCode(flow, autoLogin, tweak.current, flowState);
+      const pythonCode = getPythonCode(flow, tweak.current, flowState);
+      const widgetCode = getWidgetCode(flow, autoLogin, flowState);
 
       tabs![0].code = curl_code;
       tabs![1].code = pythonApiCode;
diff --git a/src/frontend/src/modals/formModal/index.tsx b/src/frontend/src/modals/formModal/index.tsx
index b14d7a9d0..68e32a325 100644
--- a/src/frontend/src/modals/formModal/index.tsx
+++ b/src/frontend/src/modals/formModal/index.tsx
@@ -40,20 +40,19 @@ export default function FormModal({
 }): JSX.Element {
   const nodes = useFlowStore((state) => state.nodes);
   const edges = useFlowStore((state) => state.edges);
-  const currentFlowState = useFlowsManagerStore(
-    (state) => state.currentFlowState
+  const flowState = useFlowStore(
+    (state) => state.flowState
   );
-  const setCurrentFlowState = useFlowsManagerStore(
-    (state) => state.setCurrentFlowState
+  const setFlowState = useFlowStore(
+    (state) => state.setFlowState
   );
   const [chatValue, setChatValue] = useState(() => {
     try {
-      const formKeysData = currentFlowState?.formKeysData;
-      if (!formKeysData) {
-        throw new Error("formKeysData is undefined");
+      if (!flowState) {
+        throw new Error("flowState is undefined");
       }
-      const inputKeys = formKeysData.input_keys;
-      const handleKeys = formKeysData.handle_keys;
+      const inputKeys = flowState.input_keys;
+      const handleKeys = flowState.handle_keys;
 
       const keyToUse = Object.keys(inputKeys!).find(
         (key) => !handleKeys?.some((j) => j === key) && inputKeys![key] === ""
@@ -68,7 +67,7 @@ export default function FormModal({
   });
 
   const [chatHistory, setChatHistory] = useState<ChatMessageType[]>([]);
-  const template = useRef(currentFlowState?.formKeysData.template ?? undefined);
+  const template = useRef(flowState?.template ?? undefined);
   const { accessToken } = useContext(AuthContext);
   const setErrorData = useAlertStore((state) => state.setErrorData);
   const ws = useRef<WebSocket | null>(null);
@@ -77,11 +76,11 @@ export default function FormModal({
   const messagesRef = useRef<HTMLDivElement | null>(null);
 
   const [chatKey, setChatKey] = useState(() => {
-    if (currentFlowState?.formKeysData?.input_keys) {
-      return Object.keys(currentFlowState.formKeysData.input_keys!).find(
+    if (flowState?.input_keys) {
+      return Object.keys(flowState.input_keys!).find(
         (key) =>
-          !currentFlowState.formKeysData.handle_keys!.some((j) => j === key) &&
-          currentFlowState.formKeysData.input_keys![key] === ""
+          !flowState.handle_keys!.some((j) => j === key) &&
+          flowState.input_keys![key] === ""
       );
     }
     // TODO: return a sensible default
@@ -387,7 +386,7 @@ export default function FormModal({
     let nodeValidationErrors = validateNodes(nodes, edges);
     if (nodeValidationErrors.length === 0) {
       setLockChat(true);
-      let inputs = currentFlowState?.formKeysData.input_keys;
+      let inputs = flowState?.input_keys;
       setChatValue("");
       const message = inputs;
       addChatHistory(message!, true, chatKey!, template.current);
@@ -399,10 +398,10 @@ export default function FormModal({
         description: flow.description,
         chatKey: chatKey!,
       });
-      if (currentFlowState && chatKey) {
-        setCurrentFlowState((old: FlowState | undefined) => {
+      if (flowState && chatKey) {
+        setFlowState((old: FlowState | undefined) => {
           let newFlowState = cloneDeep(old!);
-          newFlowState.formKeysData.input_keys![chatKey] = "";
+          newFlowState.input_keys![chatKey] = "";
           return newFlowState;
         });
       }
@@ -415,7 +414,7 @@ export default function FormModal({
   }
   function clearChat(): void {
     setChatHistory([]);
-    template.current = currentFlowState?.formKeysData.template;
+    template.current = flowState?.template;
     ws.current?.send(JSON.stringify({ clear_history: true }));
     if (lockChat) setLockChat(false);
   }
@@ -423,7 +422,7 @@ export default function FormModal({
   function handleOnCheckedChange(checked: boolean, i: string) {
     if (checked === true) {
       setChatKey(i);
-      setChatValue(currentFlowState?.formKeysData.input_keys![i] ?? "");
+      setChatValue(flowState?.input_keys![i] ?? "");
     } else {
       setChatKey(null!);
       setChatValue("");
@@ -432,7 +431,7 @@ export default function FormModal({
   return (
     <Dialog open={open} onOpenChange={setOpen}>
       <DialogTrigger hidden></DialogTrigger>
-      {currentFlowState && currentFlowState.formKeysData && (
+      {flowState && flowState && (
         <DialogContent className="min-w-[80vw]">
           <DialogHeader>
             <DialogTitle className="flex items-center">
@@ -468,8 +467,8 @@ export default function FormModal({
                 </div>
               </div>
 
-              {currentFlowState?.formKeysData?.input_keys
-                ? Object.keys(currentFlowState?.formKeysData.input_keys!).map(
+              {flowState?.input_keys
+                ? Object.keys(flowState?.input_keys!).map(
                     (key, index) => (
                       <div className="file-component-accordion-div" key={index}>
                         <AccordionComponent
@@ -491,7 +490,7 @@ export default function FormModal({
                                     handleOnCheckedChange(value, key)
                                   }
                                   size="small"
-                                  disabled={currentFlowState.formKeysData.handle_keys!.some(
+                                  disabled={flowState.handle_keys!.some(
                                     (t) => t === key
                                   )}
                                 />
@@ -502,7 +501,7 @@ export default function FormModal({
                           keyValue={key}
                         >
                           <div className="file-component-tab-column">
-                            {currentFlowState?.formKeysData.handle_keys!.some(
+                            {flowState?.handle_keys!.some(
                               (t) => t === key
                             ) && (
                               <div className="font-normal text-muted-foreground ">
@@ -512,14 +511,14 @@ export default function FormModal({
                             <Textarea
                               className="custom-scroll"
                               value={
-                                currentFlowState?.formKeysData.input_keys![key]
+                                flowState?.input_keys![key]
                               }
                               onChange={(e) => {
-                                if (currentFlowState) {
-                                  setCurrentFlowState(
+                                if (flowState) {
+                                  setFlowState(
                                     (old: FlowState | undefined) => {
                                       let newFlowState = cloneDeep(old!);
-                                      newFlowState.formKeysData.input_keys![
+                                      newFlowState.input_keys![
                                         key
                                       ] = e.target.value;
                                       return newFlowState;
@@ -536,7 +535,7 @@ export default function FormModal({
                     )
                   )
                 : null}
-              {currentFlowState?.formKeysData.memory_keys!.map((key, index) => (
+              {flowState?.memory_keys!.map((key, index) => (
                 <div className="file-component-accordion-div" key={index}>
                   <AccordionComponent
                     trigger={
@@ -628,11 +627,11 @@ export default function FormModal({
                       sendMessage={sendMessage}
                       setChatValue={(value) => {
                         setChatValue(value);
-                        if (currentFlowState && chatKey) {
-                          setCurrentFlowState(
+                        if (flowState && chatKey) {
+                          setFlowState(
                             (old: FlowState | undefined) => {
                               let newFlowState = cloneDeep(old!);
-                              newFlowState.formKeysData.input_keys![
+                              newFlowState.input_keys![
                                 chatKey
                               ] = value;
                               return newFlowState;
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
index d4b2b8e3b..cd2a5214c 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -1,12 +1,5 @@
 import _ from "lodash";
-import {
-  MouseEvent,
-  useCallback,
-  useContext,
-  useEffect,
-  useRef,
-  useState,
-} from "react";
+import { MouseEvent, useCallback, useEffect, useRef, useState } from "react";
 import ReactFlow, {
   Background,
   Connection,
@@ -18,13 +11,10 @@ import ReactFlow, {
   SelectionDragHandler,
   addEdge,
   updateEdge,
-  useEdgesState,
-  useNodesState,
 } from "reactflow";
 import GenericNode from "../../../../CustomNodes/GenericNode";
 import Chat from "../../../../components/chatComponent";
 import Loading from "../../../../components/ui/loading";
-import { locationContext } from "../../../../contexts/locationContext";
 import useAlertStore from "../../../../stores/alertStore";
 import useFlowStore from "../../../../stores/flowStore";
 import useFlowsManagerStore from "../../../../stores/flowsManagerStore";
@@ -43,6 +33,7 @@ import { cn, getRandomName, isWrappedWithClass } from "../../../../utils/utils";
 import ConnectionLineComponent from "../ConnectionLineComponent";
 import SelectionMenu from "../SelectionMenuComponent";
 import ExtraSidebar from "../extraSidebarComponent";
+import { useLocationStore } from "../../../../stores/locationStore";
 
 const nodeTypes = {
   genericNode: GenericNode,
@@ -64,11 +55,6 @@ export default function Page({
   const setFilterEdge = useTypesStore((state) => state.setFilterEdge);
   const reactFlowWrapper = useRef<HTMLDivElement>(null);
 
-  const [lastCopiedSelection, setLastCopiedSelection] = useState<{
-    nodes: any;
-    edges: any;
-  } | null>(null);
-
   const reactFlowInstance = useFlowStore((state) => state.reactFlowInstance);
   const setReactFlowInstance = useFlowStore(
     (state) => state.setReactFlowInstance
@@ -86,6 +72,13 @@ export default function Page({
   const redo = useFlowsManagerStore((state) => state.redo);
   const takeSnapshot = useFlowsManagerStore((state) => state.takeSnapshot);
   const paste = useFlowStore((state) => state.paste);
+  const resetFlow = useFlowStore((state) => state.resetFlow);
+  const lastCopiedSelection = useFlowStore(
+    (state) => state.lastCopiedSelection
+  );
+  const setLastCopiedSelection = useFlowStore(
+    (state) => state.setLastCopiedSelection
+  );
 
   const position = useRef({ x: 0, y: 0 });
   const [lastSelection, setLastSelection] =
@@ -163,7 +156,12 @@ export default function Page({
 
   const [selectionMenuVisible, setSelectionMenuVisible] = useState(false);
 
-  const { setExtraComponent, setExtraNavigation } = useContext(locationContext);
+  const setExtraComponent = useLocationStore(
+    (state) => state.setExtraComponent
+  );
+  const setExtraNavigation = useLocationStore(
+    (state) => state.setExtraNavigation
+  );
   const setErrorData = useAlertStore((state) => state.setErrorData);
 
   const edgeUpdateSuccessful = useRef(true);
@@ -177,13 +175,11 @@ export default function Page({
   useEffect(() => {
     setLoading(true);
     if (reactFlowInstance) {
-      useFlowStore.setState({
+      resetFlow({
         nodes: flow?.data?.nodes ?? [],
         edges: flow?.data?.edges ?? [],
-      });
-      reactFlowInstance.setViewport(
-        flow?.data?.viewport ?? { zoom: 1, x: 0, y: 0 }
-      );
+        viewport: flow?.data?.viewport ?? { zoom: 1, x: 0, y: 0 },
+      })
     }
 
     // Clear the previous timeout
@@ -224,8 +220,7 @@ export default function Page({
   const onMoveEnd: OnMove = useCallback(() => {
     // 👇 make moving the canvas undoable
     autoSaveCurrentFlow(nodes, edges, reactFlowInstance?.getViewport()!);
-  }
-  , [takeSnapshot, autoSaveCurrentFlow, nodes, edges, reactFlowInstance]);
+  }, [takeSnapshot, autoSaveCurrentFlow, nodes, edges, reactFlowInstance]);
 
   const onSelectionDragStart: SelectionDragHandler = useCallback(() => {
     // 👇 make dragging a selection undoable
@@ -290,14 +285,16 @@ export default function Page({
         const newNode: NodeType = {
           id: newId,
           type: "genericNode",
-          position: {x: 0, y:0},
+          position: { x: 0, y: 0 },
           data: {
             ...data,
             id: newId,
           },
         };
-        paste({ nodes: [newNode], edges: [] }, {x: event.clientX, y: event.clientY});
-
+        paste(
+          { nodes: [newNode], edges: [] },
+          { x: event.clientX, y: event.clientY }
+        );
       } else if (event.dataTransfer.types.some((types) => types === "Files")) {
         takeSnapshot();
         if (event.dataTransfer.files.item(0)!.type === "application/json") {
@@ -384,7 +381,6 @@ export default function Page({
     setFilterEdge([]);
   }, []);
 
-
   return (
     <div className="flex h-full overflow-hidden">
       {!view && <ExtraSidebar />}
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 273ee2533..18bd40bea 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -28,16 +28,51 @@ import useFlowsManagerStore from "./flowsManagerStore";
 
 // this is our useStore hook that we can use in our components to get parts of the store and call actions
 const useFlowStore = create<FlowStoreType>((set, get) => ({
-  reactFlowInstance: null,
-  setReactFlowInstance: (newState) => {
-    set({ reactFlowInstance: newState });
-  },
+  sseData: {},
+  flowState: undefined,
   nodes: [],
   edges: [],
-  isBuilt: false,
+  isBuilding: false,
+  isBuilt: false, 
+  reactFlowInstance: null,
+  lastCopiedSelection: null,
+
+
+  resetFlow: ({ nodes, edges, viewport }) => {
+    set({
+      nodes,
+      edges,
+      flowState: undefined,
+      sseData: {},
+      isBuilt: false,
+
+    });
+    get().reactFlowInstance!.setViewport(viewport);
+  },
+  updateSSEData: (sseData) => {
+    set((state) => ({ sseData: { ...state.sseData, ...sseData } }));
+  },
+  setIsBuilding: (isBuilding) => {
+    set({ isBuilding });
+  },
   setIsBuilt: (isBuilt) => {
     set({ isBuilt });
   },
+  setFlowState: (flowState) => {
+    const newFlowState =
+      typeof flowState === "function"
+        ? flowState(get().flowState)
+        : flowState;
+    
+    if(newFlowState !== get().flowState){
+      set(() => ({
+        flowState: newFlowState,
+      }));
+    }
+  },
+  setReactFlowInstance: (newState) => {
+    set({ reactFlowInstance: newState });
+  },
   onNodesChange: (changes: NodeChange[]) => {
     set({
       nodes: applyNodeChanges(changes, get().nodes),
@@ -52,29 +87,30 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     let newChange = typeof change === "function" ? change(get().nodes) : change;
     let newEdges = cleanEdges(newChange, get().edges);
 
-    set({ edges: newEdges });
-    set({ nodes: newChange });
+    set({ edges: newEdges, nodes: newChange, flowState: undefined, isBuilt: false, sseData: {} });
 
-    useFlowsManagerStore
-      .getState()
-      .autoSaveCurrentFlow(
+    const flowsManager = useFlowsManagerStore.getState()
+
+    flowsManager.autoSaveCurrentFlow(
         newChange,
         newEdges,
         get().reactFlowInstance?.getViewport() ?? { x: 0, y: 0, zoom: 1 }
       );
+
   },
   setEdges: (change) => {
     let newChange = typeof change === "function" ? change(get().edges) : change;
 
-    set({ edges: newChange });
+    set({ edges: newChange, flowState: undefined, isBuilt: false, sseData: {} });
 
-    useFlowsManagerStore
-      .getState()
-      .autoSaveCurrentFlow(
+    const flowsManager = useFlowsManagerStore.getState()
+
+    flowsManager.autoSaveCurrentFlow(
         get().nodes,
         newChange,
         get().reactFlowInstance?.getViewport() ?? { x: 0, y: 0, zoom: 1 }
       );
+
   },
   setNode: (id: string, change: Node | ((oldState: Node) => Node)) => {
     let newChange =
@@ -204,6 +240,9 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     });
     set({ edges: newEdges });
   },
+  setLastCopiedSelection: (newSelection) => {
+    set({ lastCopiedSelection: newSelection });
+  },
 }));
 
 export default useFlowStore;
diff --git a/src/frontend/src/stores/flowsManagerStore.ts b/src/frontend/src/stores/flowsManagerStore.ts
index b076f75e4..e89b2a00f 100644
--- a/src/frontend/src/stores/flowsManagerStore.ts
+++ b/src/frontend/src/stores/flowsManagerStore.ts
@@ -10,8 +10,7 @@ import {
 } from "../controllers/API";
 import { FlowType, NodeDataType } from "../types/flow";
 import { FlowState } from "../types/tabs";
-import { UseUndoRedoOptions } from "../types/typesContext";
-import { FlowsManagerStoreType } from "../types/zustand/flowsManager";
+import { FlowsManagerStoreType, UseUndoRedoOptions } from "../types/zustand/flowsManager";
 import {
   addVersionToDuplicates,
   createFlowComponent,
@@ -39,7 +38,6 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
   setCurrentFlowId: (currentFlowId: string) => {
     set((state) => ({
       currentFlowId,
-      currentFlowState: state.flowsState[state.currentFlowId],
       currentFlow: state.flows.find((flow) => flow.id === currentFlowId),
     }));
   },
@@ -53,23 +51,6 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
   currentFlow: undefined,
   isLoading: true,
   setIsLoading: (isLoading: boolean) => set({ isLoading }),
-  flowsState: {},
-  currentFlowState: undefined,
-  setCurrentFlowState: (
-    flowState: FlowState | ((oldState: FlowState | undefined) => FlowState)
-  ) => {
-    const newFlowState =
-      typeof flowState === "function"
-        ? flowState(get().currentFlowState)
-        : flowState;
-    set((state) => ({
-      flowsState: {
-        ...state.flowsState,
-        [state.currentFlowId]: newFlowState,
-      },
-      currentFlowState: newFlowState,
-    }));
-  },
   refreshFlows: () => {
     return new Promise<void>((resolve, reject) => {
       set({ isLoading: true });
diff --git a/src/frontend/src/stores/locationStore.tsx b/src/frontend/src/stores/locationStore.tsx
new file mode 100644
index 000000000..247304712
--- /dev/null
+++ b/src/frontend/src/stores/locationStore.tsx
@@ -0,0 +1,28 @@
+import { create } from "zustand";
+import { LocationStoreType } from "../types/zustand/location";
+
+export const useLocationStore = create<LocationStoreType>((set, get) => ({
+    current: window.location.pathname.replace(/\/$/g, "").split("/"),
+    isStackedOpen:
+      window.innerWidth > 1024 && window.location.pathname.split("/")[1]
+        ? true
+        : false,
+    setCurrent: (newState) => {
+        set({ current: newState });
+        },
+    setIsStackedOpen: (newState) => {
+        set({ isStackedOpen: newState });
+    },
+    showSideBar: window.location.pathname.split("/")[1] ? true : false,
+    setShowSideBar: (newState) => {
+        set({ showSideBar: newState });
+    },
+    extraNavigation: { title: "" },
+    setExtraNavigation: (newState) => {
+        set({ extraNavigation: newState });
+    },
+    extraComponent: <></>,
+    setExtraComponent: (newState) => {
+        set({ extraComponent: newState });
+    },
+}));
diff --git a/src/frontend/src/types/tabs/index.ts b/src/frontend/src/types/tabs/index.ts
index f5815c62c..f07bc7752 100644
--- a/src/frontend/src/types/tabs/index.ts
+++ b/src/frontend/src/types/tabs/index.ts
@@ -55,16 +55,14 @@ export type FlowsContextType = {
 };
 
 export type FlowsState = {
-  [key: string]: FlowState;
+  [key: string]: FlowState | undefined;
 };
 
 export type FlowState = {
-  formKeysData: {
-    template?: string;
-    input_keys?: Object;
-    memory_keys?: Array<string>;
-    handle_keys?: Array<string>;
-  };
+  template?: string;
+  input_keys?: Object;
+  memory_keys?: Array<string>;
+  handle_keys?: Array<string>;
 };
 
 export type errorsVarType = {
diff --git a/src/frontend/src/types/typesContext/index.ts b/src/frontend/src/types/typesContext/index.ts
deleted file mode 100644
index c7314fea0..000000000
--- a/src/frontend/src/types/typesContext/index.ts
+++ /dev/null
@@ -1,89 +0,0 @@
-import { Edge, Node } from "reactflow";
-import { AlertItemType } from "../alerts";
-
-export type alertContextType = {
-  errorData: { title: string; list?: Array<string> };
-  setErrorData: (newState: { title: string; list?: Array<string> }) => void;
-  errorOpen: boolean;
-  setErrorOpen: (newState: boolean) => void;
-  noticeData: { title: string; link?: string };
-  setNoticeData: (newState: { title: string; link?: string }) => void;
-  noticeOpen: boolean;
-  setNoticeOpen: (newState: boolean) => void;
-  successData: { title: string };
-  setSuccessData: (newState: { title: string }) => void;
-  successOpen: boolean;
-  setSuccessOpen: (newState: boolean) => void;
-  notificationCenter: boolean;
-  setNotificationCenter: (newState: boolean) => void;
-  notificationList: Array<AlertItemType>;
-  pushNotificationList: (Object: AlertItemType) => void;
-  clearNotificationList: () => void;
-  removeFromNotificationList: (index: string) => void;
-  loading: boolean;
-  setLoading: (newState: boolean) => void;
-  modalContextOpen: boolean | null;
-  setModalContextOpen: (newState: boolean) => void;
-};
-
-export type darkContextType = {
-  dark: {};
-  setDark: (newState: {}) => void;
-  stars: number;
-  setStars: (stars: number) => void;
-  gradientIndex: number;
-  setGradientIndex: (index: number) => void;
-};
-
-export type locationContextType = {
-  current: Array<string>;
-  setCurrent: (newState: Array<string>) => void;
-  isStackedOpen: boolean;
-  setIsStackedOpen: (newState: boolean) => void;
-  showSideBar: boolean;
-  setShowSideBar: (newState: boolean) => void;
-  extraNavigation: {
-    title: string;
-    options?: Array<{
-      name: string;
-      href: string;
-      icon: React.ElementType;
-      children?: Array<JSX.Element>;
-    }>;
-  };
-  setExtraNavigation: (newState: {
-    title: string;
-    options?: Array<{
-      name: string;
-      href: string;
-      icon: React.ElementType;
-      children?: Array<JSX.Element>;
-    }>;
-  }) => void;
-  extraComponent: any;
-  setExtraComponent: (newState: JSX.Element) => void;
-};
-
-export type undoRedoContextType = {
-  undo: () => void;
-  redo: () => void;
-  takeSnapshot: () => void;
-};
-
-export type UseUndoRedoOptions = {
-  maxHistorySize: number;
-  enableShortcuts: boolean;
-};
-
-export type UseUndoRedo = (options?: UseUndoRedoOptions) => {
-  undo: () => void;
-  redo: () => void;
-  takeSnapshot: () => void;
-  canUndo: boolean;
-  canRedo: boolean;
-};
-
-export type HistoryItem = {
-  nodes: Node[];
-  edges: Edge[];
-};
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index cacc7f909..a4ded1642 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -5,11 +5,20 @@ import {
   OnEdgesChange,
   OnNodesChange,
   ReactFlowInstance,
+  Viewport,
 } from "reactflow";
+import { FlowState } from "../../tabs";
 
 export type FlowStoreType = {
+  updateSSEData: (sseData: object) => void;
+  sseData: object;
+  isBuilding: boolean;
+  setIsBuilding: (isBuilding: boolean) => void;
+  resetFlow: (flow: {nodes: Node[], edges: Edge[], viewport: Viewport}) => void;
   reactFlowInstance: ReactFlowInstance | null;
   setReactFlowInstance: (newState: ReactFlowInstance) => void;
+  flowState: FlowState | undefined;
+  setFlowState: (state: FlowState | undefined | ((oldState: FlowState | undefined) => FlowState)) => void;
   nodes: Node[];
   edges: Edge[];
   onNodesChange: OnNodesChange;
@@ -24,6 +33,11 @@ export type FlowStoreType = {
     selection: { nodes: any; edges: any },
     position: { x: number; y: number; paneX?: number; paneY?: number }
   ) => void;
+  lastCopiedSelection: { nodes: any; edges: any } | null;
+  setLastCopiedSelection: (
+    newSelection: { nodes: any; edges: any } | null
+  ) => void;
   isBuilt: boolean;
   setIsBuilt: (isBuilt: boolean) => void;
+
 };
diff --git a/src/frontend/src/types/zustand/flowsManager/index.ts b/src/frontend/src/types/zustand/flowsManager/index.ts
index d130c00e6..a596ee724 100644
--- a/src/frontend/src/types/zustand/flowsManager/index.ts
+++ b/src/frontend/src/types/zustand/flowsManager/index.ts
@@ -10,9 +10,6 @@ export type FlowsManagerStoreType = {
   setCurrentFlowId: (currentFlowId: string) => void;
   isLoading: boolean;
   setIsLoading: (isLoading: boolean) => void;
-  flowsState: FlowsState;
-  currentFlowState: FlowState | undefined;
-  setCurrentFlowState: (state: FlowState | ((oldState: FlowState | undefined) => FlowState)) => void;
   refreshFlows: () => Promise<void>;
   saveFlow: (flow: FlowType, silent?: boolean) => Promise<void>;
   autoSaveCurrentFlow: (nodes: Node[], edges: Edge[], viewport: Viewport) => void;
@@ -26,3 +23,8 @@ export type FlowsManagerStoreType = {
   redo: () => void;
   takeSnapshot: () => void;
 };
+
+export type UseUndoRedoOptions = {
+  maxHistorySize: number;
+  enableShortcuts: boolean;
+};
\ No newline at end of file
diff --git a/src/frontend/src/types/zustand/location/index.ts b/src/frontend/src/types/zustand/location/index.ts
new file mode 100644
index 000000000..2d2d44766
--- /dev/null
+++ b/src/frontend/src/types/zustand/location/index.ts
@@ -0,0 +1,28 @@
+export type LocationStoreType = {
+  current: Array<string>;
+  setCurrent: (newState: Array<string>) => void;
+  isStackedOpen: boolean;
+  setIsStackedOpen: (newState: boolean) => void;
+  showSideBar: boolean;
+  setShowSideBar: (newState: boolean) => void;
+  extraNavigation: {
+    title: string;
+    options?: Array<{
+      name: string;
+      href: string;
+      icon: React.ElementType;
+      children?: Array<JSX.Element>;
+    }>;
+  };
+  setExtraNavigation: (newState: {
+    title: string;
+    options?: Array<{
+      name: string;
+      href: string;
+      icon: React.ElementType;
+      children?: Array<JSX.Element>;
+    }>;
+  }) => void;
+  extraComponent: any;
+  setExtraComponent: (newState: JSX.Element) => void;
+};
diff --git a/src/frontend/src/utils/utils.ts b/src/frontend/src/utils/utils.ts
index 006d29435..86539634c 100644
--- a/src/frontend/src/utils/utils.ts
+++ b/src/frontend/src/utils/utils.ts
@@ -217,12 +217,11 @@ export function groupByFamily(
       }));
 }
 
-export function buildInputs(currentFlowState?: FlowState): string {
-  return currentFlowState &&
-    currentFlowState.formKeysData &&
-    currentFlowState.formKeysData.input_keys &&
-    Object.keys(currentFlowState.formKeysData.input_keys!).length > 0
-    ? JSON.stringify(currentFlowState.formKeysData.input_keys)
+export function buildInputs(flowState?: FlowState): string {
+  return flowState &&
+    flowState.input_keys &&
+    Object.keys(flowState.input_keys!).length > 0
+    ? JSON.stringify(flowState.input_keys)
     : '{"input": "message"}';
 }
 
@@ -297,16 +296,15 @@ export function buildTweakObject(tweak: tweakType) {
  * @param {FlowsState} tabsState - The current tabs state.
  * @returns {string} - The chat input field
  */
-export function getChatInputField(flow: FlowType, currentFlowState?: FlowState) {
+export function getChatInputField(flow: FlowType, flowState?: FlowState) {
   let chat_input_field = "text";
 
   if (
-    currentFlowState &&
-    currentFlowState.formKeysData &&
-    currentFlowState.formKeysData.input_keys
+    flowState &&
+    flowState.input_keys
   ) {
     chat_input_field = Object.keys(
-      currentFlowState.formKeysData.input_keys!
+      flowState.input_keys!
     )[0];
   }
   return chat_input_field;
@@ -321,7 +319,7 @@ export function getPythonApiCode(
   flow: FlowType,
   isAuth: boolean,
   tweak?: any[],
-  currentFlowState?: FlowState
+  flowState?: FlowState
 ): string {
   const flowId = flow.id;
 
@@ -330,7 +328,7 @@ export function getPythonApiCode(
   //   node.data.id
   // }
   const tweaks = buildTweaks(flow);
-  const inputs = buildInputs(currentFlowState);
+  const inputs = buildInputs(flowState);
   return `import requests
 from typing import Optional
 
@@ -385,11 +383,11 @@ export function getCurlCode(
   flow: FlowType,
   isAuth: boolean,
   tweak?: any[],
-  currentFlowState?: FlowState
+  flowState?: FlowState
 ): string {
   const flowId = flow.id;
   const tweaks = buildTweaks(flow);
-  const inputs = buildInputs(currentFlowState);
+  const inputs = buildInputs(flowState);
 
   return `curl -X POST \\
   ${window.location.protocol}//${
@@ -413,11 +411,11 @@ export function getCurlCode(
 export function getPythonCode(
   flow: FlowType,
   tweak?: any[],
-  currentFlowState?: FlowState
+  flowState?: FlowState
 ): string {
   const flowName = flow.name;
   const tweaks = buildTweaks(flow);
-  const inputs = buildInputs(currentFlowState);
+  const inputs = buildInputs(flowState);
   return `from langflow import load_flow_from_json
 TWEAKS = ${
     tweak && tweak.length > 0
@@ -438,12 +436,12 @@ flow(inputs)`;
 export function getWidgetCode(
   flow: FlowType,
   isAuth: boolean,
-  currentFlowState?: FlowState
+  flowState?: FlowState
 ): string {
   const flowId = flow.id;
   const flowName = flow.name;
-  const inputs = buildInputs(currentFlowState);
-  let chat_input_field = getChatInputField(flow, currentFlowState);
+  const inputs = buildInputs(flowState);
+  let chat_input_field = getChatInputField(flow, flowState);
 
   return `<script src="https://cdn.jsdelivr.net/gh/logspace-ai/langflow-embedded-chat@main/dist/build/static/js/bundle.min.js"></script>
 
@@ -454,7 +452,7 @@ chat_input_field: Input key that you want the chat to send the user message with
   window_title="${flowName}"
   flow_id="${flowId}"
   ${
-    currentFlowState && currentFlowState.formKeysData
+    flowState
       ? `chat_inputs='${inputs}'
   chat_input_field="${chat_input_field}"
   `