diff --git a/src/frontend/src/components/chatComponent/buildTrigger/index.tsx b/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
index 084afc455..8ea6f9048 100644
--- a/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
+++ b/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
@@ -52,9 +52,7 @@ export default function BuildTrigger({
     }
   }
 
-  const checkInputAndOutput = useFlowStore(
-    (state) => state.checkInputAndOutput
-  );
+  const hasIO = useFlowStore((state) => state.hasIO);
 
   async function enforceMinimumLoadingTime(
     startTime: number,
@@ -80,11 +78,7 @@ export default function BuildTrigger({
       leaveTo="translate-y-96"
     >
       <div
-        className={
-          checkInputAndOutput()
-            ? "fixed bottom-20 right-4"
-            : "fixed bottom-4 right-4"
-        }
+        className={hasIO ? "fixed bottom-20 right-4" : "fixed bottom-4 right-4"}
       >
         <div
           className={`${eventClick} round-button-form`}
diff --git a/src/frontend/src/components/chatComponent/index.tsx b/src/frontend/src/components/chatComponent/index.tsx
index fb3ad5001..1f6b997f9 100644
--- a/src/frontend/src/components/chatComponent/index.tsx
+++ b/src/frontend/src/components/chatComponent/index.tsx
@@ -12,27 +12,14 @@ export default function Chat({ flow }: ChatType): JSX.Element {
   const [open, setOpen] = useState(false);
   const flowState = useFlowStore((state) => state.flowState);
   const nodes = useFlowStore((state) => state.nodes);
-  const checkInputAndOutput = useFlowStore(
-    (state) => state.checkInputAndOutput
-  );
-  const getOutputs = useFlowStore((state) => state.getOutputs);
-  const getInputs = useFlowStore((state) => state.getInputs);
-  const [showTrigger, setShowTrigger] = useState(checkInputAndOutput());
-  useEffect(() => {
-    const haveIO = checkInputAndOutput();
-    setShowTrigger(haveIO);
-    if (haveIO) {
-      getOutputs();
-      getInputs();
-    }
-  }, [nodes]);
+  const hasIO = useFlowStore((state) => state.hasIO);
 
   useEffect(() => {
     const handleKeyDown = (event: KeyboardEvent) => {
       if (
         (event.key === "K" || event.key === "k") &&
         (event.metaKey || event.ctrlKey) &&
-        checkInputAndOutput()
+        useFlowStore.getState().hasIO
       ) {
         event.preventDefault();
         setOpen((oldState) => !oldState);
@@ -50,7 +37,7 @@ export default function Chat({ flow }: ChatType): JSX.Element {
     <>
       <div className="flex flex-col">
         <BuildTrigger open={open} flow={flow} />
-        {showTrigger && (
+        {hasIO && (
           <BaseModal open={open} setOpen={setOpen}>
             <BaseModal.Trigger asChild>
               <div onClick={() => setOpen(true)}>
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 329e02055..4127de910 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -21,11 +21,10 @@ import {
   cleanEdges,
   getHandleId,
   getNodeId,
-  isInputNode,
-  isOutputNode,
   scapeJSONParse,
   scapedJSONStringfy,
 } from "../utils/reactflowUtils";
+import { getInputsAndOutputs } from "../utils/storeUtils";
 import useAlertStore from "./alertStore";
 import useFlowsManagerStore from "./flowsManagerStore";
 
@@ -36,6 +35,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
   edges: [],
   isBuilding: false,
   isPending: false,
+  hasIO: false,
   reactFlowInstance: null,
   lastCopiedSelection: null,
   flowPool: {},
@@ -61,10 +61,18 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     set({ isPending });
   },
   resetFlow: ({ nodes, edges, viewport }) => {
+    let newEdges = cleanEdges(nodes, edges);
+    const { inputs, outputs } = getInputsAndOutputs(nodes);
+
     set({
       nodes,
-      edges,
+      edges: newEdges,
       flowState: undefined,
+      inputIds: inputs.map((input) => input.id),
+      outputIds: outputs.map((output) => output.id),
+      inputTypes: inputs.map((input) => input.type),
+      outputTypes: outputs.map((output) => output.type),
+      hasIO: inputs.length > 0 && outputs.length > 0,
     });
     get().reactFlowInstance!.setViewport(viewport);
   },
@@ -97,11 +105,17 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
   setNodes: (change) => {
     let newChange = typeof change === "function" ? change(get().nodes) : change;
     let newEdges = cleanEdges(newChange, get().edges);
+    const { inputs, outputs } = getInputsAndOutputs(newChange);
 
     set({
       edges: newEdges,
       nodes: newChange,
       flowState: undefined,
+      inputIds: inputs.map((input) => input.id),
+      outputIds: outputs.map((output) => output.id),
+      inputTypes: inputs.map((input) => input.type),
+      outputTypes: outputs.map((output) => output.type),
+      hasIO: inputs.length > 0 && outputs.length > 0,
     });
 
     const flowsManager = useFlowsManagerStore.getState();
@@ -144,21 +158,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       })
     );
   },
-  checkInputAndOutput: () => {
-    let has_input = false;
-    let has_output = false;
-    const nodes = get().nodes;
-    nodes.forEach((node) => {
-      const nodeData: NodeDataType = node.data as NodeDataType;
-      if (isInputNode(nodeData)) {
-        has_input = true;
-      }
-      if (isOutputNode(nodeData)) {
-        has_output = true;
-      }
-    });
-    return has_input && has_output;
-  },
   getNode: (id: string) => {
     return get().nodes.find((node) => node.id === id);
   },
@@ -369,36 +368,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       viewport: get().reactFlowInstance?.getViewport()!,
     };
   },
-  getOutputs(): { type: string; id: string }[] {
-    let outputs: { type: string; id: string }[] = [];
-    const nodes = get().nodes;
-    nodes.forEach((node) => {
-      const nodeData: NodeDataType = node.data as NodeDataType;
-      if (isOutputNode(nodeData)) {
-        outputs.push({ type: nodeData.type, id: nodeData.id });
-      }
-    });
-    set({
-      outputIds: outputs.map((output) => output.id),
-      outputTypes: outputs.map((output) => output.type),
-    });
-    return outputs;
-  },
-  getInputs(): { type: string; id: string }[] {
-    let inputs: { type: string; id: string }[] = [];
-    const nodes = get().nodes;
-    nodes.forEach((node) => {
-      const nodeData: NodeDataType = node.data as NodeDataType;
-      if (isOutputNode(nodeData)) {
-        inputs.push({ type: nodeData.type, id: nodeData.id });
-      }
-    });
-    set({
-      inputIds: inputs.map((input) => input.id),
-      inputTypes: inputs.map((input) => input.type),
-    });
-    return inputs;
-  },
 }));
 
 export default useFlowStore;
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index 5b5ac495f..69e59c08d 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -38,6 +38,7 @@ export type FlowStoreType = {
   inputTypes: string[];
   outputTypes: string[];
   inputIds: string[];
+  hasIO: boolean;
   outputIds: string[];
   setFlowPool: (flowPool: FlowPoolType) => void;
   addDataToFlowPool: (data: any, nodeId: string) => void;
@@ -84,8 +85,5 @@ export type FlowStoreType = {
   onConnect: (connection: Connection) => void;
   unselectAll: () => void;
   buildFlow: (nodeId?: string) => Promise<void>;
-  checkInputAndOutput: () => boolean;
   getFlow: () => { nodes: Node[]; edges: Edge[]; viewport: Viewport };
-  getOutputs: () => { type: string; id: string }[];
-  getInputs: () => { type: string; id: string }[];
 };