From d93ecee87613505b371081a4f3c938cefc5c9313 Mon Sep 17 00:00:00 2001
From: Gabriel Almeida <gabriel@logspace.ai>
Date: Sun, 26 Mar 2023 10:03:08 -0300
Subject: [PATCH] feat: implement node templates

---
 src/backend/langflow/custom/customs.py | 40 ++--------------------
 src/backend/langflow/node/__init__.py  |  0
 src/backend/langflow/node/nodes.py     | 46 ++++++++++++++++++++++++++
 src/backend/langflow/node/template.py  | 44 ++++++++++++++++++++++++
 4 files changed, 92 insertions(+), 38 deletions(-)
 create mode 100644 src/backend/langflow/node/__init__.py
 create mode 100644 src/backend/langflow/node/nodes.py
 create mode 100644 src/backend/langflow/node/template.py

diff --git a/src/backend/langflow/custom/customs.py b/src/backend/langflow/custom/customs.py
index cea81ac1a..c7ea0d20d 100644
--- a/src/backend/langflow/custom/customs.py
+++ b/src/backend/langflow/custom/customs.py
@@ -1,42 +1,6 @@
-from langchain.agents.mrkl import prompt
+from langflow.node.nodes import ZeroShotPromptNode
 
 
 def get_custom_prompts():
     """Get custom prompts."""
-
-    return {
-        "ZeroShotPrompt": {
-            "template": {
-                "_type": "zero_shot",
-                "prefix": {
-                    "type": "str",
-                    "required": False,
-                    "placeholder": "",
-                    "list": False,
-                    "show": True,
-                    "multiline": True,
-                    "value": prompt.PREFIX,
-                },
-                "suffix": {
-                    "type": "str",
-                    "required": True,
-                    "placeholder": "",
-                    "list": False,
-                    "show": True,
-                    "multiline": True,
-                    "value": prompt.SUFFIX,
-                },
-                "format_instructions": {
-                    "type": "str",
-                    "required": False,
-                    "placeholder": "",
-                    "list": False,
-                    "show": True,
-                    "multiline": True,
-                    "value": prompt.FORMAT_INSTRUCTIONS,
-                },
-            },
-            "description": "Prompt template for Zero Shot Agent.",
-            "base_classes": ["BasePromptTemplate"],
-        }
-    }
+    return ZeroShotPromptNode().to_dict()
diff --git a/src/backend/langflow/node/__init__.py b/src/backend/langflow/node/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/backend/langflow/node/nodes.py b/src/backend/langflow/node/nodes.py
new file mode 100644
index 000000000..77e09dd16
--- /dev/null
+++ b/src/backend/langflow/node/nodes.py
@@ -0,0 +1,46 @@
+from langflow.node.template import Field, FrontendNode, Template
+from langchain.agents.mrkl import prompt
+
+
+class ZeroShotPromptNode(FrontendNode):
+    _name = "ZeroShotPrompt"
+    template = Template(
+        type_name="zero_shot",
+        fields=[
+            Field(
+                field_type="str",
+                required=False,
+                placeholder="",
+                is_list=False,
+                show=True,
+                multiline=True,
+                value=prompt.PREFIX,
+                name="prefix",
+            ),
+            Field(
+                field_type="str",
+                required=True,
+                placeholder="",
+                is_list=False,
+                show=True,
+                multiline=True,
+                value=prompt.SUFFIX,
+                name="suffix",
+            ),
+            Field(
+                field_type="str",
+                required=False,
+                placeholder="",
+                is_list=False,
+                show=True,
+                multiline=True,
+                value=prompt.FORMAT_INSTRUCTIONS,
+                name="format_instructions",
+            ),
+        ],
+    )
+    description = "Prompt template for Zero Shot Agent."
+    base_classes = ["BasePromptTemplate"]
+
+    def to_dict(self):
+        return super().to_dict()
diff --git a/src/backend/langflow/node/template.py b/src/backend/langflow/node/template.py
new file mode 100644
index 000000000..90f3333fc
--- /dev/null
+++ b/src/backend/langflow/node/template.py
@@ -0,0 +1,44 @@
+from typing import Any
+from pydantic import BaseModel
+
+
+class Field(BaseModel):
+    field_type: str = "str"
+    required: bool = False
+    placeholder: str = ""
+    is_list: bool = False
+    show: bool = True
+    multiline: bool = False
+    value: Any = None
+    # _name will be used to store the name of the field
+    # in the template
+    name: str = None
+
+    def to_dict(self):
+        return self.dict()
+
+
+class Template(BaseModel):
+    type_name: str
+    fields: list[Field]
+
+    def to_dict(self):
+        result = {field.name: field.to_dict() for field in self.fields}
+        result["_type"] = self.type_name
+        return result
+
+
+class FrontendNode(BaseModel):
+    template: Template
+    description: str
+    base_classes: list
+    _name: str = None
+
+    def to_dict(self):
+        return {
+            self._name: {
+                "template": self.template.to_dict(),
+                "description": self.description,
+                "base_classes": self.base_classes,
+            }
+        }