diff --git a/src/backend/base/langflow/custom/custom_component/base_component.py b/src/backend/base/langflow/custom/custom_component/base_component.py
index 5d1d63117..2b63b8ad8 100644
--- a/src/backend/base/langflow/custom/custom_component/base_component.py
+++ b/src/backend/base/langflow/custom/custom_component/base_component.py
@@ -1,4 +1,5 @@
 import operator
+import re
 from typing import Any, ClassVar
 from uuid import UUID
 
@@ -93,7 +94,15 @@ class BaseComponent:
         if not self._code:
             return {}
 
-        cc_class = eval_custom_component_code(self._code)
+        try:
+            cc_class = eval_custom_component_code(self._code)
+
+        except AttributeError as e:
+            pattern = r"module '.*?' has no attribute '.*?'"
+            if re.search(pattern, str(e)):
+                raise ImportError(e) from e
+            raise
+
         component_instance = cc_class(_code=self._code)
         return self.get_template_config(component_instance)
 
diff --git a/src/backend/base/langflow/schema/__init__.py b/src/backend/base/langflow/schema/__init__.py
index e84ff7dfc..619ab3ec0 100644
--- a/src/backend/base/langflow/schema/__init__.py
+++ b/src/backend/base/langflow/schema/__init__.py
@@ -1,4 +1,5 @@
 from .data import Data
 from .dotdict import dotdict
+from .message import Message
 
-__all__ = ["Data", "dotdict"]
+__all__ = ["Data", "dotdict", "Message"]