diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py
index 06a4f4d7c..353298cbd 100644
--- a/src/backend/langflow/interface/custom/custom_component.py
+++ b/src/backend/langflow/interface/custom/custom_component.py
@@ -30,6 +30,8 @@ class CustomComponent(Component, extra=Extra.allow):
         return self.field_config
 
     def _class_template_validation(self, code: str) -> bool:
+        TYPE_HINT_LIST = ["Optional", "Prompt", "PromptTemplate", "LLMChain"]
+
         if not code:
             raise HTTPException(
                 status_code=400,
@@ -40,14 +42,14 @@ class CustomComponent(Component, extra=Extra.allow):
             )
 
         reader = DirectoryReader("", False)
-        if reader.is_type_hint_used_but_not_imported("Optional", code):
-            raise HTTPException(
-                status_code=400,
-                detail={
+
+        for type_hint in TYPE_HINT_LIST:
+            if reader.is_type_hint_used_but_not_imported(type_hint, code):
+                error_detail = {
                     "error": "Type hint Error",
-                    "traceback": "Name type hint 'Optional' is not defined",
-                },
-            )
+                    "traceback": f"Type hint '{type_hint}' is used but not imported in the code.",
+                }
+                raise HTTPException(status_code=400, detail=error_detail)
 
     def is_check_valid(self) -> bool:
         return self._class_template_validation(self.code)
diff --git a/src/backend/langflow/interface/custom/directory_reader.py b/src/backend/langflow/interface/custom/directory_reader.py
index 2f448e319..b3c8a1883 100644
--- a/src/backend/langflow/interface/custom/directory_reader.py
+++ b/src/backend/langflow/interface/custom/directory_reader.py
@@ -196,7 +196,7 @@ class DirectoryReader:
         elif not self.validate_build(file_content):
             return False, "Missing build function"
         elif self.is_type_hint_used_but_not_imported("Optional", file_content):
-            return False, "Name type hint 'Optional' is not defined"
+            return False, "Type hint 'Optional' is used but not imported in the code."
         else:
             if self.compress_code_field:
                 file_content = str(StringCompressor(file_content).compress_string())