From 0ef54c5ad9d0a6ecfaa0f924cf9503b172c88963 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Tue, 28 Jan 2025 09:18:40 -0300
Subject: [PATCH] fix: Ensure importability of modules not included in
 __init__.py files (#5965)

* fix: ensure that modules not included in __init__.py files are importable

* test: add test for module import in custom component
---
 src/backend/base/langflow/utils/validate.py  |  8 +++++++-
 src/backend/tests/unit/test_validate_code.py | 20 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/backend/base/langflow/utils/validate.py b/src/backend/base/langflow/utils/validate.py
index 7d5646285..46d1c49de 100644
--- a/src/backend/base/langflow/utils/validate.py
+++ b/src/backend/base/langflow/utils/validate.py
@@ -244,7 +244,13 @@ def prepare_global_scope(module):
                     warnings.simplefilter("ignore", LangChainDeprecationWarning)
                     imported_module = importlib.import_module(node.module)
                     for alias in node.names:
-                        exec_globals[alias.name] = getattr(imported_module, alias.name)
+                        try:
+                            # First try getting it as an attribute
+                            exec_globals[alias.name] = getattr(imported_module, alias.name)
+                        except AttributeError:
+                            # If that fails, try importing the full module path
+                            full_module_path = f"{node.module}.{alias.name}"
+                            exec_globals[alias.name] = importlib.import_module(full_module_path)
             except ModuleNotFoundError as e:
                 msg = f"Module {node.module} not found. Please install it and try again"
                 raise ModuleNotFoundError(msg) from e
diff --git a/src/backend/tests/unit/test_validate_code.py b/src/backend/tests/unit/test_validate_code.py
index 4470b3efe..0e9394efa 100644
--- a/src/backend/tests/unit/test_validate_code.py
+++ b/src/backend/tests/unit/test_validate_code.py
@@ -127,6 +127,26 @@ class MyComponent(CustomComponent):
     assert result.value == "test"
 
 
+def test_create_class_module_import():
+    code = """
+from langflow.custom import CustomComponent
+from PIL import ImageDraw
+
+class ExternalClass:
+    def __init__(self, value):
+        self.value = value
+
+class MyComponent(CustomComponent):
+    def build(self):
+        return ExternalClass("test")
+"""
+    class_name = "MyComponent"
+    created_class = create_class(code, class_name)
+    instance = created_class()
+    result = instance.build()
+    assert result.value == "test"
+
+
 def test_create_class_with_multiple_external_classes():
     code = """
 from langflow.custom import CustomComponent