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