diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index a76a15d17..9c1b4bbf1 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -71,7 +71,7 @@ export default function ParameterComponent({
 
   const groupedEdge = useRef(null);
 
-  const setFilterEdge = useTypesStore((state) => state.setFilterEdge);
+  const setFilterEdge = useFlowStore((state) => state.setFilterEdge);
 
   let disabled =
     edges.some(
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
index 5852d6999..2e2c2bac7 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -50,7 +50,7 @@ export default function Page({
   );
   const types = useTypesStore((state) => state.types);
   const templates = useTypesStore((state) => state.templates);
-  const setFilterEdge = useTypesStore((state) => state.setFilterEdge);
+  const setFilterEdge = useFlowStore((state) => state.setFilterEdge);
   const reactFlowWrapper = useRef<HTMLDivElement>(null);
 
   const reactFlowInstance = useFlowStore((state) => state.reactFlowInstance);
diff --git a/src/frontend/src/pages/FlowPage/components/extraSidebarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/extraSidebarComponent/index.tsx
index b610663ac..f9e904581 100644
--- a/src/frontend/src/pages/FlowPage/components/extraSidebarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/extraSidebarComponent/index.tsx
@@ -29,8 +29,8 @@ import SidebarDraggableComponent from "./sideBarDraggableComponent";
 export default function ExtraSidebar(): JSX.Element {
   const data = useTypesStore((state) => state.data);
   const templates = useTypesStore((state) => state.templates);
-  const getFilterEdge = useTypesStore((state) => state.getFilterEdge);
-  const setFilterEdge = useTypesStore((state) => state.setFilterEdge);
+  const getFilterEdge = useFlowStore((state) => state.getFilterEdge);
+  const setFilterEdge = useFlowStore((state) => state.setFilterEdge);
   const uploadFlow = useFlowsManagerStore((state) => state.uploadFlow);
   const currentFlow = useFlowsManagerStore((state) => state.currentFlow);
   const hasStore = useStoreStore((state) => state.hasStore);
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 76929cba1..cf041ee0f 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -257,8 +257,13 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       flowState: undefined,
       sseData: {},
       isBuilt: false,
+      getFilterEdge: [],
     });
   },
+  setFilterEdge: (newState) => {
+    set({ getFilterEdge: newState });
+  },
+  getFilterEdge: []
 }));
 
 export default useFlowStore;
diff --git a/src/frontend/src/stores/typesStore.ts b/src/frontend/src/stores/typesStore.ts
index e2ab3dfc4..90f2c5b9b 100644
--- a/src/frontend/src/stores/typesStore.ts
+++ b/src/frontend/src/stores/typesStore.ts
@@ -9,7 +9,6 @@ export const useTypesStore = create<TypesStoreType>((set, get) => ({
   types: {},
   templates: {},
   data: {},
-  getFilterEdge: [],
   getTypes: () => {
     return new Promise<void>(async (resolve, reject) => {
       getAll()
@@ -44,7 +43,5 @@ export const useTypesStore = create<TypesStoreType>((set, get) => ({
     let newChange = typeof change === "function" ? change(get().data) : change;
     set({ data: newChange });
   },
-  setFilterEdge: (newState) => {
-    set({ getFilterEdge: newState });
-  },
+
 }));
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index 50d2fd45b..9b76a5347 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -50,4 +50,6 @@ export type FlowStoreType = {
   isBuilt: boolean;
   setIsBuilt: (isBuilt: boolean) => void;
   cleanFlow: () => void;
+  setFilterEdge: (newState) => void;
+  getFilterEdge: any[];
 };
diff --git a/src/frontend/src/types/zustand/types/index.ts b/src/frontend/src/types/zustand/types/index.ts
index f18d8b22e..133afbda1 100644
--- a/src/frontend/src/types/zustand/types/index.ts
+++ b/src/frontend/src/types/zustand/types/index.ts
@@ -8,6 +8,4 @@ export type TypesStoreType = {
   data: APIDataType;
   setData: (newState: {}) => void;
   getTypes: () => Promise<void>;
-  setFilterEdge: (newState) => void;
-  getFilterEdge: any[];
 };