From 43c4fe7dfcb3e15d0c75758564a4d303f309cd72 Mon Sep 17 00:00:00 2001
From: Gabriel Almeida <gabriel@logspace.ai>
Date: Sat, 1 Apr 2023 18:58:03 -0300
Subject: [PATCH] feat: get_signature can return the FrontendNode directly now

---
 src/backend/langflow/custom/customs.py | 12 +++++-------
 src/backend/langflow/interface/base.py |  6 ++++--
 src/backend/langflow/template/base.py  |  2 --
 src/backend/langflow/template/nodes.py |  2 +-
 4 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/src/backend/langflow/custom/customs.py b/src/backend/langflow/custom/customs.py
index b820cd53c..0467bac41 100644
--- a/src/backend/langflow/custom/customs.py
+++ b/src/backend/langflow/custom/customs.py
@@ -1,14 +1,12 @@
 from langflow.template import nodes
 
 CUSTOM_NODES = {
-    "prompts": {
-        **nodes.ZeroShotPromptNode().to_dict(),
-    },
-    "tools": {**nodes.PythonFunctionNode().to_dict(), **nodes.ToolNode().to_dict()},
+    "prompts": {"ZeroShotPrompt": nodes.ZeroShotPromptNode()},
+    "tools": {"PythonFunction": nodes.PythonFunctionNode(), "Tool": nodes.ToolNode()},
     "agents": {
-        **nodes.JsonAgentNode().to_dict(),
-        **nodes.CSVAgentNode().to_dict(),
-        **nodes.InitializeAgentNode().to_dict(),
+        "JsonAgent": nodes.JsonAgentNode(),
+        "CSVAgent": nodes.CSVAgentNode(),
+        "InitializeAgent": nodes.InitializeAgentNode(),
     },
 }
 
diff --git a/src/backend/langflow/interface/base.py b/src/backend/langflow/interface/base.py
index eb3432d12..ad8ccfc6a 100644
--- a/src/backend/langflow/interface/base.py
+++ b/src/backend/langflow/interface/base.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel
 
@@ -20,7 +20,7 @@ class LangChainTypeCreator(BaseModel, ABC):
         return self.type_dict
 
     @abstractmethod
-    def get_signature(self, name: str) -> Optional[Dict[Any, Any]]:
+    def get_signature(self, name: str) -> Union[Optional[Dict[Any, Any]], FrontendNode]:
         pass
 
     @abstractmethod
@@ -42,6 +42,8 @@ class LangChainTypeCreator(BaseModel, ABC):
         signature = self.get_signature(name)
         if signature is None:
             raise ValueError(f"{name} not found")
+        if isinstance(signature, FrontendNode):
+            return signature
         fields = [
             TemplateField(
                 name=key,
diff --git a/src/backend/langflow/template/base.py b/src/backend/langflow/template/base.py
index 8ed6cd7e3..87b5b6a8f 100644
--- a/src/backend/langflow/template/base.py
+++ b/src/backend/langflow/template/base.py
@@ -54,8 +54,6 @@ class TemplateFieldCreator(BaseModel, ABC):
         if "List" in _type:
             _type = _type.replace("List[", "")[:-1]
             self.is_list = True
-        else:
-            self.is_list = False
 
         # Replace 'Mapping' with 'dict'
         if "Mapping" in _type:
diff --git a/src/backend/langflow/template/nodes.py b/src/backend/langflow/template/nodes.py
index 46d39db9d..6b855e5c4 100644
--- a/src/backend/langflow/template/nodes.py
+++ b/src/backend/langflow/template/nodes.py
@@ -153,7 +153,7 @@ class JsonAgentNode(FrontendNode):
 
 
 class InitializeAgentNode(FrontendNode):
-    name: str = "InializeAgent"
+    name: str = "InitializeAgent"
     template: Template = Template(
         type_name="initailize_agent",
         fields=[