🔨 refactor(langflow): improve code parsing and custom component handling

- Refactor code parsing in `code_parser.py` to handle imports, function definitions, and class attributes more robustly and cleanly.
- Add new methods in `component.py` to parse Assign and AnnAssign statements, and FunctionDef statements.
- Refactor `custom_component.py` to improve the handling of custom components, including better extraction of main class name and template configuration.
- Update `types.py` to better handle the building of custom component templates, including handling of field configurations and error handling.
- Minor formatting fix in `conftest.py` test fixture.

These changes improve the robustness and readability of the code, and provide better handling and validation of custom components.

src/backend/langflow/interface/custom/code_parser.py

@@ -61,12 +61,11 @@ class CodeParser:
         Extracts "imports" from the code.
         """
         if isinstance(node, ast.Import):
-            module = node.names[0].name
+            for alias in node.names:
-            self.data["imports"].append(module)
+                self.data["imports"].append(alias.name)
         elif isinstance(node, ast.ImportFrom):
-            module = node.module
+            for alias in node.names:
-            names = [alias.name for alias in node.names]
+                self.data["imports"].append((node.module, alias.name))
-            self.data["imports"].append((module, names))
 
     def parse_functions(self, node: ast.FunctionDef) -> None:
         """
@@ -97,7 +96,7 @@ class CodeParser:
 
         # Handle positional arguments with default values
         defaults = [None] * (len(node.args.args) - len(node.args.defaults)) + [
-            ast.unparse(default) for default in node.args.defaults
+            ast.unparse(default) if default else None for default in node.args.defaults
         ]
 
         for arg, default in zip(node.args.args, defaults):
@@ -126,10 +125,38 @@ class CodeParser:
             func["body"].append(ast.unparse(line))
         return func
 
+    def parse_assign(self, stmt):
+        """
+        Parses an Assign statement and returns a dictionary
+        with the target's name and value.
+        """
+        for target in stmt.targets:
+            if isinstance(target, ast.Name):
+                return {"name": target.id, "value": ast.unparse(stmt.value)}
+
+    def parse_ann_assign(self, stmt):
+        """
+        Parses an AnnAssign statement and returns a dictionary
+        with the target's name, value, and annotation.
+        """
+        if isinstance(stmt.target, ast.Name):
+            return {
+                "name": stmt.target.id,
+                "value": ast.unparse(stmt.value) if stmt.value else None,
+                "annotation": ast.unparse(stmt.annotation),
+            }
+
+    def parse_function_def(self, stmt):
+        """
+        Parses a FunctionDef statement and returns the parsed
+        method and a boolean indicating if it's an __init__ method.
+        """
+        method = self.parse_callable_details(stmt)
+        return (method, True) if stmt.name == "__init__" else (method, False)
+
     def parse_classes(self, node: ast.ClassDef) -> None:
         """
-        Extracts "classes" from the code, including
+        Extracts "classes" from the code, including inheritance and init methods.
-        inheritance and init methods.
         """
         class_dict = {
             "name": node.name,
@@ -140,15 +167,15 @@ class CodeParser:
         }
 
         for stmt in node.body:
-            if isinstance(stmt, ast.AnnAssign):
+            if isinstance(stmt, ast.Assign):
-                attr = {"name": stmt.target.id, "type": ast.unparse(stmt.annotation)}
+                if attr := self.parse_assign(stmt):
-                class_dict["attributes"].append(attr)
+                    class_dict["attributes"].append(attr)
-            elif isinstance(stmt, ast.Assign):
+            elif isinstance(stmt, ast.AnnAssign):
-                attr = {"name": stmt.targets[0].id, "value": ast.unparse(stmt.value)}
+                if attr := self.parse_ann_assign(stmt):
-                class_dict["attributes"].append(attr)
+                    class_dict["attributes"].append(attr)
             elif isinstance(stmt, ast.FunctionDef):
-                method = self.parse_callable_details(stmt)
+                method, is_init = self.parse_function_def(stmt)
-                if stmt.name == "__init__":
+                if is_init:
                     class_dict["init"] = method
                 else:
                     class_dict["methods"].append(method)

src/backend/langflow/interface/custom/component.py

@@ -1,3 +1,4 @@
+import ast
 from pydantic import BaseModel
 from fastapi import HTTPException
 
@@ -48,5 +49,23 @@ class Component(BaseModel):
 
         return validate.create_function(self.code, self.function_entrypoint_name)
 
+    def build_template_config(self, attributes) -> dict:
+        template_config = {}
+
+        for item in attributes:
+            item_name = item.get("name")
+
+            if item_value := item.get("value"):
+                if "langflow_display_name" in item_name:
+                    template_config["display_name"] = ast.literal_eval(item_value)
+
+                elif "langflow_description" in item_name:
+                    template_config["description"] = ast.literal_eval(item_value)
+
+                elif "langflow_field_config" in item_name:
+                    template_config["field_config"] = ast.literal_eval(item_value)
+
+        return template_config
+
     def build(self):
         raise NotImplementedError

src/backend/langflow/interface/custom/custom_component.py

@@ -1,4 +1,3 @@
-import ast
 from typing import Callable, Optional
 from fastapi import HTTPException
 from langflow.interface.custom.constants import LANGCHAIN_BASE_TYPES
@@ -28,7 +27,7 @@ class CustomComponent(Component):
                 },
             )
 
-        # TODO: build logic
+        # TODO: Create the logic to validate what the Custom Component should have as a prerequisite to be able to execute
         return True
 
     def is_check_valid(self) -> bool:
@@ -92,24 +91,35 @@ class CustomComponent(Component):
         return build_method["return_type"]
 
     @property
-    def get_template_config(self) -> dict:
+    def get_main_class_name(self):
-        extra_attributes = {}  # self.get_extra_attributes
+        tree = self.get_code_tree(self.code)
-        template_config = {}
 
-        if "field_config" in extra_attributes:
+        base_name = self.code_class_base_inheritance
-            template_config["field_config"] = ast.literal_eval(
+        method_name = self.function_entrypoint_name
-                extra_attributes["field_config"]
-            )
-        if "display_name" in extra_attributes:
-            template_config["display_name"] = ast.literal_eval(
-                extra_attributes["display_name"]
-            )
-        if "description" in extra_attributes:
-            template_config["description"] = ast.literal_eval(
-                extra_attributes["description"]
-            )
 
-        return template_config
+        classes = []
+        for item in tree.get("classes"):
+            if base_name in item["bases"]:
+                method_names = [method["name"] for method in item["methods"]]
+                if method_name in method_names:
+                    classes.append(item["name"])
+
+        # Get just the first item
+        return next(iter(classes), "")
+
+    @property
+    def build_template_config(self):
+        tree = self.get_code_tree(self.code)
+
+        attributes = [
+            main_class["attributes"]
+            for main_class in tree.get("classes")
+            if main_class["name"] == self.get_main_class_name
+        ]
+        # Get just the first item
+        attributes = next(iter(attributes), [])
+
+        return super().build_template_config(attributes)
 
     @property
     def get_function(self):

src/backend/langflow/interface/types.py

@@ -161,7 +161,18 @@ def build_langchain_template_custom_component(custom_component: CustomComponent)
 
     function_args = custom_component.get_function_entrypoint_args
     return_type = custom_component.get_function_entrypoint_return_type
-    # template_config = custom_component.get_template_config
+    template_config = custom_component.build_template_config
+
+    # Rewrite diplay_name and description values
+    if frontend_node:
+        if "display_name" in template_config:
+            frontend_node["display_name"] = template_config["display_name"]
+
+        elif "description" in template_config:
+            frontend_node["description"] = template_config["description"]
+
+    # Rewrite field configurations
+    field_config = template_config.get("field_config", {})
 
     if function_args is not None:
         # Add extra fields
@@ -174,7 +185,6 @@ def build_langchain_template_custom_component(custom_component: CustomComponent)
                     extra_field.get("default") if "default" in extra_field else ""
                 )
                 field_required = True
-                field_config = {}
 
                 # TODO: Validate type - if is possible to render into frontend
                 if "optional" in field_type.lower():
@@ -184,13 +194,14 @@ def build_langchain_template_custom_component(custom_component: CustomComponent)
                 if not field_type:
                     field_type = "str"
 
+                config = field_config.get(field_name, {})
                 frontend_node = add_new_custom_field(
                     frontend_node,
                     field_name,
                     field_type,
                     field_value,
                     field_required,
-                    field_config,
+                    config,
                 )
 
     frontend_node = add_code_field(frontend_node, custom_component.code)

tests/conftest.py

@@ -252,7 +252,8 @@ class CustomChain(CustomComponent):
 
     def build(self, prompt, llm, input: str) -> Document:
         chain = MyCustomChain(prompt=prompt, llm=llm)
-        return chain(input)'''
+        return chain(input)
+'''
 
 
 @pytest.fixture