diff --git a/src/frontend/src/components/promptComponent/index.tsx b/src/frontend/src/components/promptComponent/index.tsx
index abc16c8db..fa18bd001 100644
--- a/src/frontend/src/components/promptComponent/index.tsx
+++ b/src/frontend/src/components/promptComponent/index.tsx
@@ -8,6 +8,7 @@ import { ExternalLink } from "lucide-react";
 
 export default function PromptAreaComponent({
   nodeClass,
+  setNodeClass,
   value,
   onChange,
   disabled,
@@ -46,6 +47,7 @@ export default function PromptAreaComponent({
                   onChange(t);
                 }}
                 nodeClass={nodeClass}
+                setNodeClass={setNodeClass}
               />
             );
           }}
diff --git a/src/frontend/src/controllers/API/index.ts b/src/frontend/src/controllers/API/index.ts
index 5c5c43a4d..9f212895e 100644
--- a/src/frontend/src/controllers/API/index.ts
+++ b/src/frontend/src/controllers/API/index.ts
@@ -1,10 +1,10 @@
 import {
   BuildStatusTypeAPI,
-  PromptTypeAPI,
   errorsTypeAPI,
   InitTypeAPI,
   UploadFileTypeAPI,
   APIClassType,
+  PromptTypeAPI,
 } from "./../../types/api/index";
 import { APIObjectType, sendAllProps } from "../../types/api/index";
 import axios, { AxiosResponse } from "axios";
@@ -54,9 +54,10 @@ export async function postValidateCode(
  * Checks the prompt for the code block by sending it to an API endpoint.
  *
  * @param {string} template - The template string of the prompt to check.
+ * @param {APIClassType} frontend_node - The frontend node to check.
  * @returns {Promise<AxiosResponse<PromptTypeAPI>>} A promise that resolves to an AxiosResponse containing the validation results.
  */
-export async function checkPrompt(
+export async function postValidatePrompt(
   template: string,
   frontend_node: APIClassType
 ): Promise<AxiosResponse<PromptTypeAPI>> {
diff --git a/src/frontend/src/modals/genericModal/index.tsx b/src/frontend/src/modals/genericModal/index.tsx
index 8f7dcefc8..215ba0585 100644
--- a/src/frontend/src/modals/genericModal/index.tsx
+++ b/src/frontend/src/modals/genericModal/index.tsx
@@ -1,7 +1,7 @@
 import { useContext, useRef, useState } from "react";
 import { PopUpContext } from "../../contexts/popUpContext";
 import { darkContext } from "../../contexts/darkContext";
-import { checkPrompt } from "../../controllers/API";
+import { postValidatePrompt } from "../../controllers/API";
 import { alertContext } from "../../contexts/alertContext";
 import {
   Dialog,
@@ -25,13 +25,15 @@ export default function GenericModal({
   modalTitle,
   type,
   nodeClass,
+  setNodeClass,
 }: {
   setValue: (value: string) => void;
   value: string;
   buttonText: string;
   modalTitle: string;
   type: number;
-  nodeClass: APIClassType;
+  nodeClass?: APIClassType;
+  setNodeClass?: (Class: APIClassType) => void;
 }) {
   const [myButtonText] = useState(buttonText);
   const [myModalTitle] = useState(modalTitle);
@@ -99,13 +101,12 @@ export default function GenericModal({
                   setModalOpen(false);
                   break;
                 case 2:
-                  checkPrompt(myValue, nodeClass)
+                  postValidatePrompt(myValue, nodeClass)
                     .then((apiReturn) => {
                       if (apiReturn.data) {
-                        if (apiReturn.data) {
-                          setNodeClass(data);
-                          setModalOpen(false);
-                        }
+                        setNodeClass(apiReturn.data.frontend_node);
+                        setModalOpen(false);
+
                         let inputVariables = apiReturn.data.input_variables;
                         if (inputVariables.length === 0) {
                           setErrorData({
diff --git a/src/frontend/src/modals/promptModal/index.tsx b/src/frontend/src/modals/promptModal/index.tsx
index 8a3ce4185..2ab1c140a 100644
--- a/src/frontend/src/modals/promptModal/index.tsx
+++ b/src/frontend/src/modals/promptModal/index.tsx
@@ -3,7 +3,7 @@ import { XMarkIcon, DocumentTextIcon } from "@heroicons/react/24/outline";
 import { Fragment, useContext, useRef, useState } from "react";
 import { PopUpContext } from "../../contexts/popUpContext";
 import { darkContext } from "../../contexts/darkContext";
-import { checkPrompt } from "../../controllers/API";
+import { postValidatePrompt } from "../../controllers/API";
 import { alertContext } from "../../contexts/alertContext";
 export default function PromptAreaModal({
   value,
@@ -107,7 +107,7 @@ export default function PromptAreaModal({
                       type="button"
                       className="inline-flex w-full justify-center rounded-md border border-transparent bg-indigo-600 px-4 py-2 text-base font-medium text-white shadow-sm hover:bg-ring focus:outline-none focus:ring-1 focus:ring-indigo-500 focus:ring-offset-1 sm:ml-3 sm:w-auto sm:text-sm"
                       onClick={() => {
-                        checkPrompt(myValue)
+                        postValidatePrompt(myValue)
                           .then((apiReturn) => {
                             if (apiReturn.data) {
                               let inputVariables =
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index 92827cc4e..d47663d18 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -14,6 +14,7 @@ export type APIClassType = {
   display_name: string;
   [key: string]: Array<string> | string | APITemplateType;
 };
+
 export type TemplateVariableType = {
   type: string;
   required: boolean;
@@ -38,7 +39,10 @@ export type errorsTypeAPI = {
   function: { errors: Array<string> };
   imports: { errors: Array<string> };
 };
-export type PromptTypeAPI = { input_variables: Array<string> };
+export type PromptTypeAPI = {
+  input_variables: Array<string>;
+  frontend_node: APIClassType;
+};
 
 export type BuildStatusTypeAPI = {
   built: boolean;
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index c1cc3756c..090d3a49f 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -46,6 +46,7 @@ export type InputListComponentType = {
 
 export type TextAreaComponentType = {
   nodeClass?: APIClassType;
+  setNodeClass?: (value: APIClassType) => void;
   disabled: boolean;
   onChange: (value: string[] | string) => void;
   value: string;