From 7bdf5641707f6af18d5b33970ac3e4760bfb2b18 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 12 Oct 2023 19:23:45 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20chore(constants.py):=20clean=20u?=
 =?UTF-8?q?p=20import=20statements=20and=20remove=20unnecessary=20whitespa?=
 =?UTF-8?q?ce=20to=20improve=20code=20readability=20=F0=9F=90=9B=20fix(cus?=
 =?UTF-8?q?tom=5Fcomponent.py):=20set=20the=20type=20of=20argument=20to=20?=
 =?UTF-8?q?"Data"=20if=20it=20is=20not=20specified=20to=20ensure=20consist?=
 =?UTF-8?q?ency=20and=20prevent=20errors?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../langflow/interface/custom/constants.py    | 24 ++++++++++++++++---
 .../interface/custom/custom_component.py      |  4 +++-
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/backend/langflow/interface/custom/constants.py b/src/backend/langflow/interface/custom/constants.py
index a3eb64336..6df72500c 100644
--- a/src/backend/langflow/interface/custom/constants.py
+++ b/src/backend/langflow/interface/custom/constants.py
@@ -1,15 +1,33 @@
 DEFAULT_CUSTOM_COMPONENT_CODE = """from langflow import CustomComponent
 
-from langflow.field_typing import *
+from langflow.field_typing import (
+    Tool,
+    PromptTemplate,
+    Chain,
+    BaseChatMemory,
+    BaseLLM,
+    BaseLoader,
+    BaseMemory,
+    BaseOutputParser,
+    BaseRetriever,
+    VectorStore,
+    Embeddings,
+    TextSplitter,
+    Document,
+    AgentExecutor,
+    NestedDict,
+    Data,
+)
+
 
 class Component(CustomComponent):
     display_name: str = "Custom Component"
     description: str = "Create any custom component you want!"
 
     def build_config(self):
-        return { "param": { "display_name": "Parameter" } }
+        return {"param": {"display_name": "Parameter"}}
 
     def build(self, param: Data) -> Data:
-        return params
+        return param
 
 """
diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py
index 42b68ed40..4b889e77b 100644
--- a/src/backend/langflow/interface/custom/custom_component.py
+++ b/src/backend/langflow/interface/custom/custom_component.py
@@ -108,7 +108,9 @@ class CustomComponent(Component, extra=Extra.allow):
                         ),
                     },
                 )
-        # TODO arg type should be Data if it is None
+            elif not arg.get("type"):
+                # Set the type to Data
+                arg["type"] = "Data"
         return args
 
     @property