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