diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 21ee312b4..9b7590213 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -229,7 +229,9 @@ export default function ParameterComponent({
           </div>
         ) : left === true && type === "code" ? (
           <CodeAreaComponent
-          setNodeClass={(nodeClass)=>{data.node = nodeClass}}
+            setNodeClass={(nodeClass) => {
+              data.node = nodeClass;
+            }}
             nodeClass={data.node}
             disabled={disabled}
             value={data.node.template[name].value ?? ""}
@@ -258,6 +260,7 @@ export default function ParameterComponent({
           </div>
         ) : left === true && type === "prompt" ? (
           <PromptAreaComponent
+            nodeClass={data.node}
             disabled={disabled}
             value={data.node.template[name].value ?? ""}
             onChange={handleOnNewValue}
diff --git a/src/frontend/src/components/promptComponent/index.tsx b/src/frontend/src/components/promptComponent/index.tsx
index 9752c0c10..abc16c8db 100644
--- a/src/frontend/src/components/promptComponent/index.tsx
+++ b/src/frontend/src/components/promptComponent/index.tsx
@@ -7,6 +7,7 @@ import { INPUT_STYLE } from "../../constants";
 import { ExternalLink } from "lucide-react";
 
 export default function PromptAreaComponent({
+  nodeClass,
   value,
   onChange,
   disabled,
@@ -44,6 +45,7 @@ export default function PromptAreaComponent({
                   setMyValue(t);
                   onChange(t);
                 }}
+                nodeClass={nodeClass}
               />
             );
           }}
diff --git a/src/frontend/src/modals/genericModal/index.tsx b/src/frontend/src/modals/genericModal/index.tsx
index 6f89e24a6..8f7dcefc8 100644
--- a/src/frontend/src/modals/genericModal/index.tsx
+++ b/src/frontend/src/modals/genericModal/index.tsx
@@ -1,9 +1,8 @@
-import { Fragment, useContext, useRef, useState } from "react";
+import { useContext, useRef, useState } from "react";
 import { PopUpContext } from "../../contexts/popUpContext";
 import { darkContext } from "../../contexts/darkContext";
 import { checkPrompt } from "../../controllers/API";
 import { alertContext } from "../../contexts/alertContext";
-import { TypeModal } from "../../utils";
 import {
   Dialog,
   DialogContent,
@@ -17,6 +16,7 @@ import { Button } from "../../components/ui/button";
 import { Textarea } from "../../components/ui/textarea";
 import { PROMPT_DIALOG_SUBTITLE, TEXT_DIALOG_SUBTITLE } from "../../constants";
 import { FileText } from "lucide-react";
+import { APIClassType } from "../../types/api";
 
 export default function GenericModal({
   value,
@@ -24,12 +24,14 @@ export default function GenericModal({
   buttonText,
   modalTitle,
   type,
+  nodeClass,
 }: {
   setValue: (value: string) => void;
   value: string;
   buttonText: string;
   modalTitle: string;
   type: number;
+  nodeClass: APIClassType;
 }) {
   const [myButtonText] = useState(buttonText);
   const [myModalTitle] = useState(modalTitle);
@@ -97,9 +99,13 @@ export default function GenericModal({
                   setModalOpen(false);
                   break;
                 case 2:
-                  checkPrompt(myValue)
+                  checkPrompt(myValue, nodeClass)
                     .then((apiReturn) => {
                       if (apiReturn.data) {
+                        if (apiReturn.data) {
+                          setNodeClass(data);
+                          setModalOpen(false);
+                        }
                         let inputVariables = apiReturn.data.input_variables;
                         if (inputVariables.length === 0) {
                           setErrorData({
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index b72a7e003..c1cc3756c 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -1,13 +1,7 @@
-import {
-  ComponentType,
-  ForwardRefExoticComponent,
-  ReactElement,
-  ReactNode,
-  SVGProps,
-} from "react";
+import { ReactElement, ReactNode } from "react";
 import { NodeDataType } from "../flow/index";
 import { typesContextType } from "../typesContext";
-import { APIClassType, APITemplateType } from "../api";
+import { APIClassType } from "../api";
 export type InputComponentType = {
   value: string;
   disabled?: boolean;
@@ -51,6 +45,7 @@ export type InputListComponentType = {
 };
 
 export type TextAreaComponentType = {
+  nodeClass?: APIClassType;
   disabled: boolean;
   onChange: (value: string[] | string) => void;
   value: string;