From ec7e76cc7b3f8339293a2df04c69dd1f4519b02a Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Wed, 22 Jan 2025 22:59:11 -0300
Subject: [PATCH] fix: Incorporate DEFAULT_IMPORT_STRING for dynamic class
 creation and clean up unused function (#5882)

---
 .../base/langflow/field_typing/constants.py   | 48 +++++++++++++++++++
 src/backend/base/langflow/utils/validate.py   | 13 +++--
 2 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/src/backend/base/langflow/field_typing/constants.py b/src/backend/base/langflow/field_typing/constants.py
index 3444b659a..c9f016bd3 100644
--- a/src/backend/base/langflow/field_typing/constants.py
+++ b/src/backend/base/langflow/field_typing/constants.py
@@ -82,3 +82,51 @@ CUSTOM_COMPONENT_SUPPORTED_TYPES = {
     "Retriever": Retriever,
     "DataFrame": DataFrame,
 }
+
+DEFAULT_IMPORT_STRING = """from langchain.agents.agent import AgentExecutor
+from langchain.chains.base import Chain
+from langchain.memory.chat_memory import BaseChatMemory
+from langchain_core.chat_history import BaseChatMessageHistory
+from langchain_core.document_loaders import BaseLoader
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models import BaseLanguageModel, BaseLLM
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.memory import BaseMemory
+from langchain_core.output_parsers import BaseLLMOutputParser, BaseOutputParser
+from langchain_core.prompts import BasePromptTemplate, ChatPromptTemplate, PromptTemplate
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.tools import BaseTool, Tool
+from langchain_core.vectorstores import VectorStore, VectorStoreRetriever
+from langchain_text_splitters import TextSplitter
+
+from langflow.io import (
+    BoolInput,
+    CodeInput,
+    DataFrameInput,
+    DataInput,
+    DefaultPromptField,
+    DictInput,
+    DropdownInput,
+    FileInput,
+    FloatInput,
+    HandleInput,
+    IntInput,
+    LinkInput,
+    MessageInput,
+    MessageTextInput,
+    MultilineInput,
+    MultilineSecretInput,
+    MultiselectInput,
+    NestedDictInput,
+    Output,
+    PromptInput,
+    SecretStrInput,
+    SliderInput,
+    StrInput,
+    TableInput,
+)
+from langflow.schema.data import Data
+from langflow.schema.dataframe import DataFrame
+from langflow.schema.message import Message
+"""
diff --git a/src/backend/base/langflow/utils/validate.py b/src/backend/base/langflow/utils/validate.py
index dee5eb61f..7d5646285 100644
--- a/src/backend/base/langflow/utils/validate.py
+++ b/src/backend/base/langflow/utils/validate.py
@@ -9,7 +9,7 @@ from langchain_core._api.deprecation import LangChainDeprecationWarning
 from loguru import logger
 from pydantic import ValidationError
 
-from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES
+from langflow.field_typing.constants import CUSTOM_COMPONENT_SUPPORTED_TYPES, DEFAULT_IMPORT_STRING
 
 
 def add_type_ignores() -> None:
@@ -189,8 +189,10 @@ def create_class(code, class_name):
         "from langflow.interface.custom.custom_component import CustomComponent",
         "from langflow.custom import CustomComponent",
     )
+    # Add DEFAULT_IMPORT_STRING
+    code = DEFAULT_IMPORT_STRING + "\n" + code
     module = ast.parse(code)
-    exec_globals = prepare_global_scope(code, module)
+    exec_globals = prepare_global_scope(module)
 
     class_code = extract_class_code(module, class_name)
     compiled_class = compile_class_code(class_code)
@@ -215,11 +217,10 @@ def create_type_ignore_class():
     return TypeIgnore
 
 
-def prepare_global_scope(code, module):
+def prepare_global_scope(module):
     """Prepares the global scope with necessary imports from the provided code module.
 
     Args:
-        code: The Python code
         module: AST parsed module
 
     Returns:
@@ -229,7 +230,6 @@ def prepare_global_scope(code, module):
         ModuleNotFoundError: If a module is not found in the code
     """
     exec_globals = globals().copy()
-    exec_globals.update(get_default_imports(code))
     for node in module.body:
         if isinstance(node, ast.Import):
             for alias in node.names:
@@ -309,13 +309,12 @@ def build_class_constructor(compiled_class, exec_globals, class_name):
             if isinstance(module, type(importlib)):
                 globals()[module_name] = module
 
-        exec_globals[class_name]
-
         return exec_globals[class_name]
 
     return build_custom_class()
 
 
+# TODO: Remove this function
 def get_default_imports(code_string):
     """Returns a dictionary of default imports for the dynamic class constructor."""
     default_imports = {