diff --git a/src/frontend/src/customNodes/genericNode/components/OutputComponent/index.tsx b/src/frontend/src/customNodes/genericNode/components/OutputComponent/index.tsx
index 69ca3965c..f3c959e5e 100644
--- a/src/frontend/src/customNodes/genericNode/components/OutputComponent/index.tsx
+++ b/src/frontend/src/customNodes/genericNode/components/OutputComponent/index.tsx
@@ -18,6 +18,7 @@ export default function OutputComponent({
   frozen = false,
   nodeId,
   idx,
+  name,
 }: outputComponentType) {
   const setNode = useFlowStore((state) => state.setNode);
   const updateNodeInternals = useUpdateNodeInternals();
@@ -27,33 +28,36 @@ export default function OutputComponent({
   }
 
   return (
-    <DropdownMenu>
-      <DropdownMenuTrigger>
-        <span
-          className={cn(frozen ? " text-ice" : "", "flex items-center gap-1")}
-        >
-          {selected}
-          <ForwardedIconComponent name="ChevronDown" className="h-4 w-4" />
-        </span>
-      </DropdownMenuTrigger>
-      <DropdownMenuContent>
-        {types.map((type) => (
-          <DropdownMenuItem
-            onSelect={() => {
-              // TODO: UDPDATE SET NODE TO NEW NODE FORM
-              setNode(nodeId, (node) => {
-                const newNode = cloneDeep(node);
-                (newNode.data as NodeDataType).node!.outputs![idx].selected =
-                  type;
-                return newNode;
-              });
-              updateNodeInternals(nodeId);
-            }}
+    <div className="flex gap-2">
+      <span>{name}</span>
+      <DropdownMenu>
+        <DropdownMenuTrigger>
+          <span
+            className={cn(frozen ? " text-ice" : "", "flex items-center gap-1")}
           >
-            {type}
-          </DropdownMenuItem>
-        ))}
-      </DropdownMenuContent>
-    </DropdownMenu>
+            {selected}
+            <ForwardedIconComponent name="ChevronDown" className="h-4 w-4" />
+          </span>
+        </DropdownMenuTrigger>
+        <DropdownMenuContent>
+          {types.map((type) => (
+            <DropdownMenuItem
+              onSelect={() => {
+                // TODO: UDPDATE SET NODE TO NEW NODE FORM
+                setNode(nodeId, (node) => {
+                  const newNode = cloneDeep(node);
+                  (newNode.data as NodeDataType).node!.outputs![idx].selected =
+                    type;
+                  return newNode;
+                });
+                updateNodeInternals(nodeId);
+              }}
+            >
+              {type}
+            </DropdownMenuItem>
+          ))}
+        </DropdownMenuContent>
+      </DropdownMenu>
+    </div>
   );
 }
diff --git a/src/frontend/src/customNodes/genericNode/components/parameterComponent/index.tsx b/src/frontend/src/customNodes/genericNode/components/parameterComponent/index.tsx
index 80d85afdc..151d40595 100644
--- a/src/frontend/src/customNodes/genericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/customNodes/genericNode/components/parameterComponent/index.tsx
@@ -59,6 +59,7 @@ export default function ParameterComponent({
   proxy,
   showNode,
   index,
+  outputName,
 }: ParameterComponentType): JSX.Element {
   const infoHtml = useRef<HTMLDivElement & ReactNode>(null);
   const nodes = useFlowStore((state) => state.nodes);
@@ -79,7 +80,7 @@ export default function ParameterComponent({
     debouncedHandleUpdateValues,
     setNode,
     isLoading,
-    setIsLoading
+    setIsLoading,
   );
 
   const { handleNodeClass: handleNodeClassHook } = useHandleNodeClass(
@@ -87,7 +88,7 @@ export default function ParameterComponent({
     name,
     takeSnapshot,
     setNode,
-    updateNodeInternals
+    updateNodeInternals,
   );
 
   const { handleRefreshButtonPress: handleRefreshButtonPressHook } =
@@ -96,7 +97,7 @@ export default function ParameterComponent({
   let disabled =
     edges.some(
       (edge) =>
-        edge.targetHandle === scapedJSONStringfy(proxy ? { ...id, proxy } : id)
+        edge.targetHandle === scapedJSONStringfy(proxy ? { ...id, proxy } : id),
     ) ?? false;
 
   const handleRefreshButtonPress = async (name, data) => {
@@ -107,7 +108,7 @@ export default function ParameterComponent({
 
   const handleOnNewValue = async (
     newValue: string | string[] | boolean | Object[],
-    skipSnapshot: boolean | undefined = false
+    skipSnapshot: boolean | undefined = false,
   ): Promise<void> => {
     handleOnNewValueHook(newValue, skipSnapshot);
   };
@@ -137,6 +138,7 @@ export default function ParameterComponent({
         selected={title}
         nodeId={data.id}
         frozen={data.node?.frozen}
+        name={outputName ?? type ?? title}
       />
     ) : (
       <span>{title}</span>
@@ -188,14 +190,14 @@ export default function ParameterComponent({
               className={classNames(
                 left ? "my-12 -ml-0.5 " : " my-12 -mr-0.5 ",
                 "h-3 w-3 rounded-full border-2 bg-background",
-                !showNode ? "mt-0" : ""
+                !showNode ? "mt-0" : "",
               )}
               style={{
                 borderColor: color ?? nodeColors.unknown,
               }}
               onClick={() => {
                 setFilterEdge(
-                  groupByFamily(myData, tooltipTitle!, left, nodes!)
+                  groupByFamily(myData, tooltipTitle!, left, nodes!),
                 );
               }}
             ></Handle>
@@ -280,12 +282,12 @@ export default function ParameterComponent({
                   }
                   className={classNames(
                     left ? "-ml-0.5" : "-mr-0.5",
-                    "h-3 w-3 rounded-full border-2 bg-background"
+                    "h-3 w-3 rounded-full border-2 bg-background",
                   )}
                   style={{ borderColor: color ?? nodeColors.unknown }}
                   onClick={() => {
                     setFilterEdge(
-                      groupByFamily(myData, tooltipTitle!, left, nodes!)
+                      groupByFamily(myData, tooltipTitle!, left, nodes!),
                     );
                   }}
                 />
diff --git a/src/frontend/src/customNodes/genericNode/index.tsx b/src/frontend/src/customNodes/genericNode/index.tsx
index a78b78088..2163af716 100644
--- a/src/frontend/src/customNodes/genericNode/index.tsx
+++ b/src/frontend/src/customNodes/genericNode/index.tsx
@@ -55,14 +55,14 @@ export default function GenericNode({
   const [nodeName, setNodeName] = useState(data.node!.display_name);
   const [inputDescription, setInputDescription] = useState(false);
   const [nodeDescription, setNodeDescription] = useState(
-    data.node?.description!
+    data.node?.description!,
   );
   const [isOutdated, setIsOutdated] = useState(false);
   const buildStatus = useFlowStore(
-    (state) => state.flowBuildStatus[data.id]?.status
+    (state) => state.flowBuildStatus[data.id]?.status,
   );
   const lastRunTime = useFlowStore(
-    (state) => state.flowBuildStatus[data.id]?.timestamp
+    (state) => state.flowBuildStatus[data.id]?.timestamp,
   );
   const [validationStatus, setValidationStatus] =
     useState<validationStatusType | null>(null);
@@ -115,7 +115,7 @@ export default function GenericNode({
 
       updateNodeInternals(data.id);
     },
-    [data.id, data.node, setNode, setIsOutdated]
+    [data.id, data.node, setNode, setIsOutdated],
   );
 
   if (!data.node!.template) {
@@ -252,60 +252,10 @@ export default function GenericNode({
     );
   };
 
-  const buildParameterComponent = ({
-    data,
-    conditionalPath,
-    showNode,
-    left,
-  }: {
-    data: NodeDataType;
-    conditionalPath: string | null;
-    showNode: boolean;
-    left: boolean;
-  }) => {
-    return (
-      <ParameterComponent
-        key={scapedJSONStringfy({
-          baseClasses: data.node!.base_classes,
-          id: data.id,
-          dataType: data.type,
-          conditionalPath: conditionalPath,
-        })}
-        data={data}
-        color={
-          (data.node?.output_types && data.node.output_types.length > 0
-            ? nodeColors[data.node.output_types[0]] ??
-              nodeColors[types[data.node.output_types[0]]]
-            : nodeColors[types[data.type]]) ?? nodeColors.unknown
-        }
-        title={
-          data.node?.output_types && data.node.output_types.length > 0
-            ? data.node.output_types.join(" | ")
-            : data.type
-        }
-        conditionPath={conditionalPath}
-        tooltipTitle={data.node?.base_classes.join("\n")}
-        id={{
-          baseClasses: data.node!.base_classes,
-          id: data.id,
-          dataType: data.type,
-          // First parameter component should be true
-          // Second should be false
-          conditionalPath: conditionalPath,
-        }}
-        // Type should be base_classes if it's not a conditional node
-        // else it should be true in the first parameter component
-        type={data.node?.base_classes.join("|")}
-        left={left}
-        showNode={showNode}
-      />
-    );
-  };
-
   const isDark = useDarkStore((state) => state.dark);
   const renderIconStatus = (
     buildStatus: BuildStatus | undefined,
-    validationStatus: validationStatusType | null
+    validationStatus: validationStatusType | null,
   ) => {
     if (buildStatus === BuildStatus.BUILDING) {
       return <Loading className="text-medium-indigo" />;
@@ -346,7 +296,7 @@ export default function GenericNode({
   };
   const getSpecificClassFromBuildStatus = (
     buildStatus: BuildStatus | undefined,
-    validationStatus: validationStatusType | null
+    validationStatus: validationStatusType | null,
   ) => {
     let isInvalid = validationStatus && !validationStatus.valid;
 
@@ -370,11 +320,11 @@ export default function GenericNode({
     selected: boolean,
     showNode: boolean,
     buildStatus: BuildStatus | undefined,
-    validationStatus: validationStatusType | null
+    validationStatus: validationStatusType | null,
   ) => {
     const specificClassFromBuildStatus = getSpecificClassFromBuildStatus(
       buildStatus,
-      validationStatus
+      validationStatus,
     );
 
     const baseBorderClass = getBaseBorderClass(selected);
@@ -383,7 +333,7 @@ export default function GenericNode({
       baseBorderClass,
       nodeSizeClass,
       "generic-node-div",
-      specificClassFromBuildStatus
+      specificClassFromBuildStatus,
     );
     return names;
   };
@@ -443,7 +393,7 @@ export default function GenericNode({
           selected,
           showNode,
           buildStatus,
-          validationStatus
+          validationStatus,
         )}
       >
         {data.node?.beta && showNode && (
@@ -588,7 +538,7 @@ export default function GenericNode({
                             }
                             title={getFieldTitle(
                               data.node?.template!,
-                              templateField
+                              templateField,
                             )}
                             info={data.node?.template[templateField].info}
                             name={templateField}
@@ -616,9 +566,9 @@ export default function GenericNode({
                             proxy={data.node?.template[templateField].proxy}
                             showNode={showNode}
                           />
-                        )
+                        ),
                     )}
-                  <ParameterComponent
+                  {/* <ParameterComponent
                     index={0}
                     key={scapedJSONStringfy({
                       baseClasses: data.node!.base_classes,
@@ -643,7 +593,7 @@ export default function GenericNode({
                     type={data.node?.base_classes.join("|")}
                     left={false}
                     showNode={showNode}
-                  />
+                  /> */}
                 </>
               )}
             </div>
@@ -775,7 +725,7 @@ export default function GenericNode({
                       !data.node?.description) &&
                       nameEditable
                       ? "font-light italic"
-                      : ""
+                      : "",
                   )}
                   onDoubleClick={(e) => {
                     setInputDescription(true);
@@ -837,13 +787,13 @@ export default function GenericNode({
                         }
                         title={getFieldTitle(
                           data.node?.template!,
-                          templateField
+                          templateField,
                         )}
                         info={data.node?.template[templateField].info}
                         name={templateField}
                         tooltipTitle={
                           data.node?.template[templateField].input_types?.join(
-                            "\n"
+                            "\n",
                           ) ?? data.node?.template[templateField].type
                         }
                         required={data.node!.template[templateField].required}
@@ -870,7 +820,7 @@ export default function GenericNode({
               <div
                 className={classNames(
                   Object.keys(data.node!.template).length < 1 ? "hidden" : "",
-                  "flex-max-width justify-center"
+                  "flex-max-width justify-center",
                 )}
               >
                 {" "}
@@ -881,7 +831,7 @@ export default function GenericNode({
                   <ParameterComponent
                     index={idx}
                     key={scapedJSONStringfy({
-                      baseClasses: [output.selected ?? output.types[0]],
+                      output_types: [output.selected ?? output.types[0]],
                       id: data.id,
                       dataType: data.type,
                     })}
@@ -892,17 +842,18 @@ export default function GenericNode({
                       nodeColors[types[data.type]] ??
                       nodeColors.unknown
                     }
-                    title={output.name}
+                    title={output.selected ?? output.types[0]}
                     tooltipTitle={output.selected ?? output.types[0]}
                     id={{
-                      baseClasses: [output.selected ?? output.types[0]],
+                      output_types: [output.selected ?? output.types[0]],
                       id: data.id,
                       dataType: data.type,
-                      idx: idx,
+                      name: output.name,
                     }}
                     type={output.types.join("|")}
                     left={false}
                     showNode={showNode}
+                    outputName={output.name}
                   />
                 ))}
             </>
diff --git a/src/frontend/src/stores/flowsManagerStore.ts b/src/frontend/src/stores/flowsManagerStore.ts
index 6a179ce84..12007199e 100644
--- a/src/frontend/src/stores/flowsManagerStore.ts
+++ b/src/frontend/src/stores/flowsManagerStore.ts
@@ -84,14 +84,14 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
       readFlowsFromDatabase()
         .then((dbData) => {
           if (dbData) {
-            const { data, flows } = processFlows(dbData, false);
+            const { data, flows } = processFlows(dbData);
             const examples = flows.filter(
-              (flow) => flow.folder_id === starterFolderId
+              (flow) => flow.folder_id === starterFolderId,
             );
             get().setExamples(examples);
 
             const flowsWithoutStarterFolder = flows.filter(
-              (flow) => flow.folder_id !== starterFolderId
+              (flow) => flow.folder_id !== starterFolderId,
             );
 
             get().setFlows(flowsWithoutStarterFolder);
@@ -119,7 +119,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
     if (get().currentFlow) {
       get().saveFlow(
         { ...get().currentFlow!, data: { nodes, edges, viewport } },
-        true
+        true,
       );
     }
   },
@@ -145,7 +145,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
                   return updatedFlow;
                 }
                 return flow;
-              })
+              }),
             );
             //update tabs state
 
@@ -194,7 +194,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
     flow?: FlowType,
     override?: boolean,
     position?: XYPosition,
-    fromDragAndDrop?: boolean
+    fromDragAndDrop?: boolean,
   ): Promise<string | undefined> => {
     let flowData = flow
       ? processDataFromFlow(flow)
@@ -209,7 +209,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
         const newFlow = createNewFlow(
           flowData!,
           flow!,
-          folder_id || my_collection_id!
+          folder_id || my_collection_id!,
         );
         const { id } = await saveFlowToDatabase(newFlow);
         newFlow.id = id;
@@ -232,7 +232,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
       const newFlow = createNewFlow(
         flowData!,
         flow!,
-        folder_id || my_collection_id!
+        folder_id || my_collection_id!,
       );
 
       const newName = addVersionToDuplicates(newFlow, get().flows);
@@ -268,7 +268,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
         .getState()
         .paste(
           { nodes: flow!.data!.nodes, edges: flow!.data!.edges },
-          position ?? { x: 10, y: 10 }
+          position ?? { x: 10, y: 10 },
         );
     }
   },
@@ -278,7 +278,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
         multipleDeleteFlowsComponents(id)
           .then(() => {
             const { data, flows } = processFlows(
-              get().flows.filter((flow) => !id.includes(flow.id))
+              get().flows.filter((flow) => !id.includes(flow.id)),
             );
             get().setFlows(flows);
             set({ isLoading: false });
@@ -298,7 +298,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
           deleteFlowFromDatabase(id)
             .then(() => {
               const { data, flows } = processFlows(
-                get().flows.filter((flow) => flow.id !== id)
+                get().flows.filter((flow) => flow.id !== id),
               );
               get().setFlows(flows);
               set({ isLoading: false });
@@ -320,7 +320,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
     return new Promise<void>((resolve) => {
       let componentFlow = get().flows.find(
         (componentFlow) =>
-          componentFlow.is_component && componentFlow.name === key
+          componentFlow.is_component && componentFlow.name === key,
       );
 
       if (componentFlow) {
@@ -368,7 +368,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
               fileData,
               undefined,
               position,
-              true
+              true,
             );
             resolve(id);
           }
@@ -409,7 +409,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
     return get().addFlow(
       true,
       createFlowComponent(component, useDarkStore.getState().version),
-      override
+      override,
     );
   },
   takeSnapshot: () => {
@@ -430,7 +430,7 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
     if (pastLength > 0) {
       past[currentFlowId] = past[currentFlowId].slice(
         pastLength - defaultOptions.maxHistorySize + 1,
-        pastLength
+        pastLength,
       );
 
       past[currentFlowId].push(newState);
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index 054ed7543..5f04e67d6 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -70,6 +70,7 @@ export type ParameterComponentType = {
   showNode?: boolean;
   index: number;
   onCloseModal?: (close: boolean) => void;
+  outputName?: string;
 };
 export type InputListComponentType = {
   value: string[];
@@ -118,6 +119,7 @@ export type outputComponentType = {
   nodeId: string;
   frozen?: boolean;
   idx: number;
+  name: string;
 };
 
 export type PromptAreaComponentType = {
@@ -528,7 +530,7 @@ export type nodeToolbarPropsType = {
   updateNodeCode?: (
     newNodeClass: APIClassType,
     code: string,
-    name: string
+    name: string,
   ) => void;
   setShowState: (show: boolean | SetStateAction<boolean>) => void;
   isOutdated?: boolean;
@@ -578,7 +580,7 @@ export type chatMessagePropsType = {
   updateChat: (
     chat: ChatMessageType,
     message: string,
-    stream_url?: string
+    stream_url?: string,
   ) => void;
 };
 
@@ -670,12 +672,12 @@ export type codeTabsPropsType = {
       value: string,
       node: NodeType,
       template: InputFieldType,
-      tweak: tweakType
+      tweak: tweakType,
     ) => string;
     buildTweakObject?: (
       tw: string,
       changes: string | string[] | boolean | number | Object[] | Object,
-      template: InputFieldType
+      template: InputFieldType,
     ) => Promise<string | void>;
   };
   activeTweaks?: boolean;
diff --git a/src/frontend/src/types/flow/index.ts b/src/frontend/src/types/flow/index.ts
index 5c3fb10a1..4fa51eda5 100644
--- a/src/frontend/src/types/flow/index.ts
+++ b/src/frontend/src/types/flow/index.ts
@@ -58,9 +58,9 @@ export type TweaksType = Array<
 export type sourceHandleType = {
   dataType: string;
   id: string;
-  baseClasses: string[];
+  output_types: string[];
   conditionalPath?: string | null;
-  idx: number;
+  name: string;
 };
 //left side
 export type targetHandleType = {
diff --git a/src/frontend/src/utils/reactflowUtils.ts b/src/frontend/src/utils/reactflowUtils.ts
index 3c3d2f2c0..261a39b6c 100644
--- a/src/frontend/src/utils/reactflowUtils.ts
+++ b/src/frontend/src/utils/reactflowUtils.ts
@@ -45,8 +45,7 @@ export function checkChatInput(nodes: Node[]) {
   return nodes.some((node) => node.data.type === "ChatInput");
 }
 
-export function cleanEdges(nodes: Node[], edges: Edge[]) {
-  console.log("cleanEdges");
+export function cleanEdges(nodes: NodeType[], edges: Edge[]) {
   let newEdges = cloneDeep(edges);
   edges.forEach((edge) => {
     // check if the source and target node still exists
@@ -76,12 +75,16 @@ export function cleanEdges(nodes: Node[], edges: Edge[]) {
       }
     }
     if (sourceHandle) {
-      const index = scapeJSONParse(sourceHandle).idx ?? 0;
+      const name = scapeJSONParse(sourceHandle).name;
+      const output = sourceNode.data.node!.outputs?.find(
+        (output) => output.name === name,
+      );
+      const outputTypes = [output?.selected ?? ""];
       const id: sourceHandleType = {
         id: sourceNode.data.id,
-        baseClasses: [sourceNode.data.node.outputs[index].selected],
+        name: name,
+        output_types: outputTypes,
         dataType: sourceNode.data.type,
-        idx: index,
       };
       if (scapedJSONStringfy(id) !== sourceHandle) {
         newEdges = newEdges.filter((e) => e.id !== edge.id);
@@ -102,18 +105,18 @@ export function unselectAllNodes({ updateNodes, data }: unselectAllNodesType) {
 export function isValidConnection(
   { source, target, sourceHandle, targetHandle }: Connection,
   nodes: Node[],
-  edges: Edge[]
+  edges: Edge[],
 ) {
   const targetHandleObject: targetHandleType = scapeJSONParse(targetHandle!);
   const sourceHandleObject: sourceHandleType = scapeJSONParse(sourceHandle!);
   if (
     targetHandleObject.inputTypes?.some(
-      (n) => n === sourceHandleObject.dataType
+      (n) => n === sourceHandleObject.dataType,
     ) ||
-    sourceHandleObject.baseClasses.some(
+    sourceHandleObject.output_types.some(
       (t) =>
         targetHandleObject.inputTypes?.some((n) => n === t) ||
-        t === targetHandleObject.type
+        t === targetHandleObject.type,
     )
   ) {
     let targetNode = nodes.find((node) => node.id === target!)?.data?.node;
@@ -146,7 +149,7 @@ export function removeApiKeys(flow: FlowType): FlowType {
 
 export function updateTemplate(
   reference: APITemplateType,
-  objectToUpdate: APITemplateType
+  objectToUpdate: APITemplateType,
 ): APITemplateType {
   let clonedObject: APITemplateType = cloneDeep(reference);
 
@@ -206,7 +209,7 @@ export const processDataFromFlow = (flow: FlowType, refreshIds = true) => {
 
 export function updateIds(
   { edges, nodes }: { edges: Edge[]; nodes: Node[] },
-  selection?: { edges: Edge[]; nodes: Node[] }
+  selection?: { edges: Edge[]; nodes: Node[] },
 ) {
   let idsMap = {};
   const selectionIds = selection?.nodes.map((n) => n.id);
@@ -234,7 +237,7 @@ export function updateIds(
       edge.source = idsMap[edge.source];
       edge.target = idsMap[edge.target];
       const sourceHandleObject: sourceHandleType = scapeJSONParse(
-        edge.sourceHandle!
+        edge.sourceHandle!,
       );
       edge.sourceHandle = scapedJSONStringfy({
         ...sourceHandleObject,
@@ -244,7 +247,7 @@ export function updateIds(
         edge.data.sourceHandle.id = edge.source;
       }
       const targetHandleObject: targetHandleType = scapeJSONParse(
-        edge.targetHandle!
+        edge.targetHandle!,
       );
       edge.targetHandle = scapedJSONStringfy({
         ...targetHandleObject,
@@ -290,11 +293,11 @@ export function validateNode(node: NodeType, edges: Edge[]): Array<string> {
           (scapeJSONParse(edge.targetHandle!) as targetHandleType).fieldName ===
             t &&
           (scapeJSONParse(edge.targetHandle!) as targetHandleType).id ===
-            node.id
+            node.id,
       )
     ) {
       errors.push(
-        `${displayName || type} is missing ${getFieldTitle(template, t)}.`
+        `${displayName || type} is missing ${getFieldTitle(template, t)}.`,
       );
     } else if (
       template[t].type === "dict" &&
@@ -308,15 +311,15 @@ export function validateNode(node: NodeType, edges: Edge[]): Array<string> {
         errors.push(
           `${displayName || type} (${getFieldTitle(
             template,
-            t
-          )}) contains duplicate keys with the same values.`
+            t,
+          )}) contains duplicate keys with the same values.`,
         );
       if (hasEmptyKey(template[t].value))
         errors.push(
           `${displayName || type} (${getFieldTitle(
             template,
-            t
-          )}) field must not be empty.`
+            t,
+          )}) field must not be empty.`,
         );
     }
     return errors;
@@ -325,7 +328,7 @@ export function validateNode(node: NodeType, edges: Edge[]): Array<string> {
 
 export function validateNodes(
   nodes: Node[],
-  edges: Edge[]
+  edges: Edge[],
 ): // this returns an array of tuples with the node id and the errors
 Array<{ id: string; errors: Array<string> }> {
   if (nodes.length === 0) {
@@ -346,19 +349,16 @@ export function updateEdges(edges: Edge[]) {
   if (edges)
     edges.forEach((edge) => {
       const targetHandleObject: targetHandleType = scapeJSONParse(
-        edge.targetHandle!
+        edge.targetHandle!,
       );
       edge.className = "stroke-gray-900 stroke-connection";
     });
 }
 
 export function addVersionToDuplicates(flow: FlowType, flows: FlowType[]) {
-  console.log("flow", flow);
-  console.log("flows", flows);
   const existingNames = flows
     .filter((f) => f.folder_id === flow.folder_id)
     .map((item) => item.name);
-  console.log("existingNames", existingNames);
   let newName = flow.name;
   let count = 1;
 
@@ -374,6 +374,7 @@ export function updateEdgesHandleIds({
   edges,
   nodes,
 }: updateEdgesHandleIdsType): Edge[] {
+  console.log("updateEdgesHandleIds");
   let newEdges = cloneDeep(edges);
   newEdges.forEach((edge) => {
     const sourceNodeId = edge.source;
@@ -396,11 +397,14 @@ export function updateEdgesHandleIds({
       };
     }
     if (source && sourceNode) {
+      const output_types =
+        sourceNode.data.node!.output_types ??
+        sourceNode.data.node!.base_classes;
       newSource = {
         id: sourceNode.data.id,
-        baseClasses: sourceNode.data.node!.base_classes,
+        output_types,
         dataType: sourceNode.data.type,
-        idx: 0,
+        name: output_types.join(" | "),
       };
     }
     edge.sourceHandle = scapedJSONStringfy(newSource!);
@@ -415,51 +419,70 @@ export function updateEdgesHandleIds({
 }
 
 export function updateNewOutput({ nodes, edges }: updateEdgesHandleIdsType) {
-  console.log("updateNewOutput");
   let newEdges = cloneDeep(edges);
   let newNodes = cloneDeep(nodes);
   newEdges.forEach((edge) => {
     if (edge.sourceHandle && edge.targetHandle) {
       let newSourceHandle: sourceHandleType = scapeJSONParse(edge.sourceHandle);
       let newTargetHandle: targetHandleType = scapeJSONParse(edge.targetHandle);
+      const id = newSourceHandle.id;
+      const sourceNodeIndex = newNodes.findIndex((node) => node.id === id);
+      let sourceNode: NodeType | undefined = undefined;
+      if (sourceNodeIndex !== -1) {
+        sourceNode = newNodes[sourceNodeIndex];
+      }
+
       let intersection;
+      //@ts-ignore
+      if (newSourceHandle.baseClasses) {
+        if (!newSourceHandle.output_types) {
+          if (sourceNode?.data.node!.output_types) {
+            newSourceHandle.output_types = sourceNode?.data.node!.output_types;
+          } else {
+            //@ts-ignore
+            newSourceHandle.output_types = newSourceHandle.baseClasses;
+          }
+        }
+        //@ts-ignore
+        delete newSourceHandle.baseClasses;
+      }
       if (newTargetHandle.inputTypes && newTargetHandle.inputTypes.length > 0) {
         //conjuction subtraction
-        intersection = newSourceHandle.baseClasses.filter((type) =>
-          newTargetHandle.inputTypes!.includes(type)
+        intersection = newSourceHandle.output_types.filter((type) =>
+          newTargetHandle.inputTypes!.includes(type),
         );
       } else {
-        intersection = newSourceHandle.baseClasses.filter(
-          (type) => type === newTargetHandle.type
+        intersection = newSourceHandle.output_types.filter(
+          (type) => type === newTargetHandle.type,
         );
       }
       const selected = intersection[0];
-      newSourceHandle.baseClasses = [selected];
-      const id = newSourceHandle.id;
-      newSourceHandle.idx = 0;
-      const sourceNodeIndex = newNodes.findIndex((node) => node.id === id);
-      if (sourceNodeIndex > -1) {
-        const sourceNode = newNodes[sourceNodeIndex];
+      newSourceHandle.name = newSourceHandle.output_types.join(" | ");
+      newSourceHandle.output_types = [selected];
+      if (sourceNode) {
+        if (!sourceNode.data.node?.outputs) {
+          sourceNode.data.node!.outputs = [];
+        }
+        const types =
+          sourceNode.data.node!.output_types ??
+          sourceNode.data.node!.base_classes;
         if (
-          !sourceNode.data.node?.outputs ||
-          sourceNode.data.node!.outputs!.length === 0
+          !sourceNode.data.node!.outputs.some(
+            (output) => output.selected === selected,
+          )
         ) {
-          const outputTypes = sourceNode.data.node!.output_types;
-          // create a new output field for each output type
-          sourceNode.data.node!.outputs = [
-            {
-              types: outputTypes ?? [],
-              selected: selected,
-              name: outputTypes?.join(" | ") ?? "",
-            },
-          ];
+          sourceNode.data.node!.outputs.push({
+            types,
+            selected: selected,
+            name: types.join(" | "),
+          });
         }
       }
+
       edge.sourceHandle = scapedJSONStringfy(newSourceHandle);
       edge.data.sourceHandle = newSourceHandle;
     }
   });
-
   return { nodes: newNodes, edges: newEdges };
 }
 
@@ -468,7 +491,7 @@ export function handleKeyDown(
     | React.KeyboardEvent<HTMLInputElement>
     | React.KeyboardEvent<HTMLTextAreaElement>,
   inputValue: string | string[] | null,
-  block: string
+  block: string,
 ) {
   //condition to fix bug control+backspace on Windows/Linux
   if (
@@ -493,7 +516,7 @@ export function handleKeyDown(
 }
 
 export function handleOnlyIntegerInput(
-  event: React.KeyboardEvent<HTMLInputElement>
+  event: React.KeyboardEvent<HTMLInputElement>,
 ) {
   if (
     event.key === "." ||
@@ -509,7 +532,7 @@ export function handleOnlyIntegerInput(
 
 export function getConnectedNodes(
   edge: Edge,
-  nodes: Array<NodeType>
+  nodes: Array<NodeType>,
 ): Array<NodeType> {
   const sourceId = edge.source;
   const targetId = edge.target;
@@ -610,7 +633,7 @@ export function checkOldEdgesHandles(edges: Edge[]): boolean {
       !edge.sourceHandle ||
       !edge.targetHandle ||
       !edge.sourceHandle.includes("{") ||
-      !edge.targetHandle.includes("{")
+      !edge.targetHandle.includes("{"),
   );
 }
 
@@ -637,7 +660,7 @@ export function customStringify(obj: any): string {
 
   const keys = Object.keys(obj).sort();
   const keyValuePairs = keys.map(
-    (key) => `"${key}":${customStringify(obj[key])}`
+    (key) => `"${key}":${customStringify(obj[key])}`,
   );
   return `{${keyValuePairs.join(",")}}`;
 }
@@ -666,7 +689,7 @@ export function getHandleId(
   source: string,
   sourceHandle: string,
   target: string,
-  targetHandle: string
+  targetHandle: string,
 ) {
   return (
     "reactflow__edge-" + source + sourceHandle + "-" + target + targetHandle
@@ -677,7 +700,7 @@ export function generateFlow(
   selection: OnSelectionChangeParams,
   nodes: Node[],
   edges: Edge[],
-  name: string
+  name: string,
 ): generateFlowType {
   const newFlowData = { nodes, edges, viewport: { zoom: 1, x: 0, y: 0 } };
   const uid = new ShortUniqueId({ length: 5 });
@@ -686,7 +709,7 @@ export function generateFlow(
   newFlowData.edges = selection.edges.filter(
     (edge) =>
       selection.nodes.some((node) => node.id === edge.target) &&
-      selection.nodes.some((node) => node.id === edge.source)
+      selection.nodes.some((node) => node.id === edge.source),
   );
   newFlowData.nodes = selection.nodes;
 
@@ -707,7 +730,7 @@ export function generateFlow(
       (edge) =>
         (selection.nodes.some((node) => node.id === edge.target) ||
           selection.nodes.some((node) => node.id === edge.source)) &&
-        newFlowData.edges.every((e) => e.id !== edge.id)
+        newFlowData.edges.every((e) => e.id !== edge.id),
     ),
   };
 }
@@ -718,13 +741,13 @@ export function reconnectEdges(groupNode: NodeType, excludedEdges: Edge[]) {
   const { nodes, edges } = groupNode.data.node!.flow!.data!;
   const lastNode = findLastNode(groupNode.data.node!.flow!.data!);
   newEdges = newEdges.filter(
-    (e) => !(nodes.some((n) => n.id === e.source) && e.source !== lastNode?.id)
+    (e) => !(nodes.some((n) => n.id === e.source) && e.source !== lastNode?.id),
   );
   newEdges.forEach((edge) => {
     if (lastNode && edge.source === lastNode.id) {
       edge.source = groupNode.id;
       let newSourceHandle: sourceHandleType = scapeJSONParse(
-        edge.sourceHandle!
+        edge.sourceHandle!,
       );
       newSourceHandle.id = groupNode.id;
       edge.sourceHandle = scapedJSONStringfy(newSourceHandle);
@@ -751,7 +774,7 @@ export function reconnectEdges(groupNode: NodeType, excludedEdges: Edge[]) {
 export function filterFlow(
   selection: OnSelectionChangeParams,
   setNodes: (update: Node[] | ((oldState: Node[]) => Node[])) => void,
-  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void
+  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void,
 ) {
   setNodes((nodes) => nodes.filter((node) => !selection.nodes.includes(node)));
   setEdges((edges) => edges.filter((edge) => !selection.edges.includes(edge)));
@@ -789,7 +812,7 @@ export function updateFlowPosition(NewPosition: XYPosition, flow: FlowType) {
 export function concatFlows(
   flow: FlowType,
   setNodes: (update: Node[] | ((oldState: Node[]) => Node[])) => void,
-  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void
+  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void,
 ) {
   const { nodes, edges } = flow.data!;
   setNodes((old) => [...old, ...nodes]);
@@ -798,7 +821,7 @@ export function concatFlows(
 
 export function validateSelection(
   selection: OnSelectionChangeParams,
-  edges: Edge[]
+  edges: Edge[],
 ): Array<string> {
   const clonedSelection = cloneDeep(selection);
   const clonedEdges = cloneDeep(edges);
@@ -812,7 +835,7 @@ export function validateSelection(
   let nodesSet = new Set(clonedSelection.nodes.map((n) => n.id));
   // then filter the edges that are connected to the nodes in the set
   let connectedEdges = clonedSelection.edges.filter(
-    (e) => nodesSet.has(e.source) && nodesSet.has(e.target)
+    (e) => nodesSet.has(e.source) && nodesSet.has(e.target),
   );
   // add the edges to the selection
   clonedSelection.edges = connectedEdges;
@@ -826,17 +849,17 @@ export function validateSelection(
     clonedSelection.nodes.some(
       (node) =>
         isInputNode(node.data as NodeDataType) ||
-        isOutputNode(node.data as NodeDataType)
+        isOutputNode(node.data as NodeDataType),
     )
   ) {
     errorsArray.push(
-      "Please select only nodes that are not input or output nodes"
+      "Please select only nodes that are not input or output nodes",
     );
   }
   //check if there are two or more nodes with free outputs
   if (
     clonedSelection.nodes.filter(
-      (n) => !clonedSelection.edges.some((e) => e.source === n.id)
+      (n) => !clonedSelection.edges.some((e) => e.source === n.id),
     ).length > 1
   ) {
     errorsArray.push("Please select only one node with free outputs");
@@ -847,7 +870,7 @@ export function validateSelection(
     clonedSelection.nodes.some(
       (node) =>
         !clonedSelection.edges.some((edge) => edge.target === node.id) &&
-        !clonedSelection.edges.some((edge) => edge.source === node.id)
+        !clonedSelection.edges.some((edge) => edge.source === node.id),
     )
   ) {
     errorsArray.push("Please select only nodes that are connected");
@@ -904,8 +927,8 @@ export function mergeNodeTemplates({
               nodeTemplate[key].display_name
                 ? nodeTemplate[key].display_name
                 : nodeTemplate[key].name
-                ? toTitleCase(nodeTemplate[key].name)
-                : toTitleCase(key);
+                  ? toTitleCase(nodeTemplate[key].name)
+                  : toTitleCase(key);
           }
         }
       });
@@ -916,7 +939,7 @@ function isHandleConnected(
   edges: Edge[],
   key: string,
   field: InputFieldType,
-  nodeId: string
+  nodeId: string,
 ) {
   /*
 		this function receives a flow and a handleId and check if there is a connection with this handle
@@ -932,7 +955,7 @@ function isHandleConnected(
             id: nodeId,
             proxy: { id: field.proxy!.id, field: field.proxy!.field },
             inputTypes: field.input_types,
-          } as targetHandleType)
+          } as targetHandleType),
       )
     ) {
       return true;
@@ -947,7 +970,7 @@ function isHandleConnected(
             fieldName: key,
             id: nodeId,
             inputTypes: field.input_types,
-          } as targetHandleType)
+          } as targetHandleType),
       )
     ) {
       return true;
@@ -970,7 +993,7 @@ export function generateNodeTemplate(Flow: FlowType) {
 
 export function generateNodeFromFlow(
   flow: FlowType,
-  getNodeId: (type: string) => string
+  getNodeId: (type: string) => string,
 ): NodeType {
   const { nodes } = flow.data!;
   const outputNode = cloneDeep(findLastNode(flow.data!));
@@ -1001,7 +1024,7 @@ export function generateNodeFromFlow(
 export function connectedInputNodesOnHandle(
   nodeId: string,
   handleId: string,
-  { nodes, edges }: { nodes: NodeType[]; edges: Edge[] }
+  { nodes, edges }: { nodes: NodeType[]; edges: Edge[] },
 ) {
   const connectedNodes: Array<{ name: string; id: string; isGroup: boolean }> =
     [];
@@ -1038,7 +1061,7 @@ export function connectedInputNodesOnHandle(
 
 export function updateProxyIdsOnTemplate(
   template: APITemplateType,
-  idsMap: { [key: string]: string }
+  idsMap: { [key: string]: string },
 ) {
   Object.keys(template).forEach((key) => {
     if (template[key].proxy && idsMap[template[key].proxy!.id]) {
@@ -1049,7 +1072,7 @@ export function updateProxyIdsOnTemplate(
 
 export function updateEdgesIds(
   edges: Edge[],
-  idsMap: { [key: string]: string }
+  idsMap: { [key: string]: string },
 ) {
   edges.forEach((edge) => {
     let targetHandle: targetHandleType = edge.data.targetHandle;
@@ -1090,7 +1113,7 @@ export function expandGroupNode(
   nodes: Node[],
   edges: Edge[],
   setNodes: (update: Node[] | ((oldState: Node[]) => Node[])) => void,
-  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void
+  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void,
 ) {
   const idsMap = updateIds(flow!.data!);
   updateProxyIdsOnTemplate(template, idsMap);
@@ -1133,7 +1156,7 @@ export function expandGroupNode(
       const lastNode = cloneDeep(findLastNode(flow!.data!));
       newEdge.source = lastNode!.id;
       let newSourceHandle: sourceHandleType = scapeJSONParse(
-        newEdge.sourceHandle!
+        newEdge.sourceHandle!,
       );
       newSourceHandle.id = lastNode!.id;
       newEdge.data.sourceHandle = newSourceHandle;
@@ -1190,7 +1213,7 @@ export function expandGroupNode(
 
 export function getGroupStatus(
   flow: FlowType,
-  ssData: { [key: string]: { valid: boolean; params: string } }
+  ssData: { [key: string]: { valid: boolean; params: string } },
 ) {
   let status = { valid: true, params: SUCCESS_BUILD };
   const { nodes } = flow.data!;
@@ -1209,7 +1232,7 @@ export function getGroupStatus(
 
 export function createFlowComponent(
   nodeData: NodeDataType,
-  version: string
+  version: string,
 ): FlowType {
   const flowNode: FlowType = {
     data: {
@@ -1245,7 +1268,7 @@ export function downloadNode(NodeFLow: FlowType) {
 
 export function updateComponentNameAndType(
   data: any,
-  component: NodeDataType
+  component: NodeDataType,
 ) {}
 
 export function removeFileNameFromComponents(flow: FlowType) {
@@ -1319,7 +1342,7 @@ export function extractFieldsFromComponenents(data: APIObjectType) {
 export function downloadFlow(
   flow: FlowType,
   flowName: string,
-  flowDescription?: string
+  flowDescription?: string,
 ) {
   let clonedFlow = cloneDeep(flow);
   removeFileNameFromComponents(clonedFlow);
@@ -1329,7 +1352,7 @@ export function downloadFlow(
       ...clonedFlow,
       name: flowName,
       description: flowDescription,
-    })
+    }),
   )}`;
 
   // create a link element and set its properties
@@ -1344,7 +1367,7 @@ export function downloadFlow(
 export function downloadFlows() {
   downloadFlowsFromDatabase().then((flows) => {
     const jsonString = `data:text/json;chatset=utf-8,${encodeURIComponent(
-      JSON.stringify(flows)
+      JSON.stringify(flows),
     )}`;
 
     // create a link element and set its properties
@@ -1368,7 +1391,7 @@ export function getRandomDescription(): string {
 export const createNewFlow = (
   flowData: ReactFlowJsonObject,
   flow: FlowType,
-  folderId: string
+  folderId: string,
 ) => {
   return {
     description: flow?.description ?? getRandomDescription(),