diff --git a/src/backend/langflow/api/v1/validate.py b/src/backend/langflow/api/v1/validate.py
index 658e5b274..431945602 100644
--- a/src/backend/langflow/api/v1/validate.py
+++ b/src/backend/langflow/api/v1/validate.py
@@ -31,18 +31,32 @@ def post_validate_code(code: Code):
 def post_validate_prompt(prompt: ValidatePromptRequest):
     try:
         input_variables = validate_prompt(prompt.template)
+        # Reinitialize custom_fields
+        old_custom_fields = prompt.frontend_node.custom_fields.copy()
+        prompt.frontend_node.custom_fields = []
+        # Add new variables to the template
         for variable in input_variables:
             try:
                 template_field = TemplateField(
                     name=variable, field_type="str", show=True, advanced=False
                 )
 
-                prompt.frontend_node.template[variable] = template_field.dict()
+                prompt.frontend_node.template[variable] = template_field.to_dict()
                 prompt.frontend_node.custom_fields.append(variable)
+
             except Exception as exc:
                 logger.exception(exc)
                 raise HTTPException(status_code=500, detail=str(exc)) from exc
 
+        # Remove variables that are not in the template anymore
+        for variable in old_custom_fields:
+            if variable not in input_variables:
+                try:
+                    prompt.frontend_node.template.pop(variable, None)
+                except Exception as exc:
+                    logger.exception(exc)
+                    raise HTTPException(status_code=500, detail=str(exc)) from exc
+
         return PromptValidationResponse(
             input_variables=input_variables,
             frontend_node=prompt.frontend_node,
diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 9b7590213..6b12d8b35 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -260,6 +260,9 @@ export default function ParameterComponent({
           </div>
         ) : left === true && type === "prompt" ? (
           <PromptAreaComponent
+            setNodeClass={(nodeClass) => {
+              data.node = nodeClass;
+            }}
             nodeClass={data.node}
             disabled={disabled}
             value={data.node.template[name].value ?? ""}
diff --git a/src/frontend/src/components/promptComponent/index.tsx b/src/frontend/src/components/promptComponent/index.tsx
index fa18bd001..9e85868ed 100644
--- a/src/frontend/src/components/promptComponent/index.tsx
+++ b/src/frontend/src/components/promptComponent/index.tsx
@@ -7,8 +7,8 @@ import { INPUT_STYLE } from "../../constants";
 import { ExternalLink } from "lucide-react";
 
 export default function PromptAreaComponent({
-  nodeClass,
   setNodeClass,
+  nodeClass,
   value,
   onChange,
   disabled,
diff --git a/src/frontend/src/utils.ts b/src/frontend/src/utils.ts
index 834094ef0..46b9c51cd 100644
--- a/src/frontend/src/utils.ts
+++ b/src/frontend/src/utils.ts
@@ -830,10 +830,16 @@ export function groupByFamily(data, baseClasses) {
 
   Object.keys(data).map((d) => {
     Object.keys(data[d]).map((n) => {
-      if (
-        data[d][n].base_classes.some((r) => baseClasses.split("\n").includes(r))
-      ) {
-        arrOfParent.push(d);
+      try {
+        if (
+          data[d][n].base_classes.some((r) =>
+            baseClasses.split("\n").includes(r)
+          )
+        ) {
+          arrOfParent.push(d);
+        }
+      } catch (e) {
+        console.log(e);
       }
     });
   });
@@ -844,16 +850,20 @@ export function groupByFamily(data, baseClasses) {
 
   Object.keys(data).map((d) => {
     Object.keys(data[d]).map((n) => {
-      baseClasses.split("\n").forEach((tol) => {
-        data[d][n].base_classes.forEach((data) => {
-          if (tol == data) {
-            arrOfType.push({
-              family: d,
-              type: data,
-            });
-          }
+      try {
+        baseClasses.split("\n").forEach((tol) => {
+          data[d][n].base_classes.forEach((data) => {
+            if (tol == data) {
+              arrOfType.push({
+                family: d,
+                type: data,
+              });
+            }
+          });
         });
-      });
+      } catch (e) {
+        console.log(e);
+      }
     });
   });