From 2db74fc30ed64d9d4e4baebbfba8d866243c1a1a Mon Sep 17 00:00:00 2001
From: Gabriel Almeida <gabriel@logspace.ai>
Date: Thu, 30 Mar 2023 19:39:44 -0300
Subject: [PATCH] feat: implement file type

---
 src/backend/langflow/interface/agents/base.py |  4 +-
 .../{custom_types.py => agents/custom.py}     | 38 +------------------
 src/backend/langflow/interface/base.py        | 15 ++++++--
 src/backend/langflow/interface/chains/base.py |  4 +-
 .../langflow/interface/custom/__init__.py     |  0
 .../langflow/interface/custom/types.py        | 37 ++++++++++++++++++
 src/backend/langflow/interface/listing.py     |  2 +
 src/backend/langflow/interface/llms/base.py   |  4 +-
 .../langflow/interface/memories/base.py       |  4 +-
 .../langflow/interface/prompts/base.py        |  4 +-
 .../langflow/interface/toolkits/__init__.py   |  0
 .../langflow/interface/toolkits/base.py       | 34 +++++++++++++++++
 src/backend/langflow/interface/tools/base.py  | 38 +++++++++++++------
 .../langflow/interface/tools/constants.py     |  2 +-
 src/backend/langflow/interface/types.py       |  2 +
 src/backend/langflow/template/template.py     | 13 ++++++-
 src/backend/langflow/utils/util.py            |  7 +++-
 tests/test_custom_types.py                    |  2 +-
 18 files changed, 149 insertions(+), 61 deletions(-)
 rename src/backend/langflow/interface/{custom_types.py => agents/custom.py} (63%)
 create mode 100644 src/backend/langflow/interface/custom/__init__.py
 create mode 100644 src/backend/langflow/interface/custom/types.py
 create mode 100644 src/backend/langflow/interface/toolkits/__init__.py
 create mode 100644 src/backend/langflow/interface/toolkits/base.py

diff --git a/src/backend/langflow/interface/agents/base.py b/src/backend/langflow/interface/agents/base.py
index c1d08b551..847336a24 100644
--- a/src/backend/langflow/interface/agents/base.py
+++ b/src/backend/langflow/interface/agents/base.py
@@ -10,7 +10,9 @@ class AgentCreator(LangChainTypeCreator):
 
     @property
     def type_to_loader_dict(self) -> Dict:
-        return loading.AGENT_TO_CLASS
+        if self.type_dict is None:
+            self.type_dict = loading.AGENT_TO_CLASS
+        return self.type_dict
 
     def get_signature(self, name: str) -> Dict | None:
         try:
diff --git a/src/backend/langflow/interface/custom_types.py b/src/backend/langflow/interface/agents/custom.py
similarity index 63%
rename from src/backend/langflow/interface/custom_types.py
rename to src/backend/langflow/interface/agents/custom.py
index 52d05b218..96c2f5a09 100644
--- a/src/backend/langflow/interface/custom_types.py
+++ b/src/backend/langflow/interface/agents/custom.py
@@ -1,44 +1,10 @@
-from typing import Callable, Optional
 from langchain import LLMChain
 from langchain.agents import AgentExecutor, ZeroShotAgent
-from langflow.utils import validate
-from pydantic import BaseModel, validator
 from langchain.agents.agent_toolkits.json.prompt import JSON_PREFIX, JSON_SUFFIX
-from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS
 from langchain.agents.agent_toolkits.json.toolkit import JsonToolkit
+from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS
 from langchain.schema import BaseLanguageModel
-
-
-class Function(BaseModel):
-    code: str
-    function: Optional[Callable] = None
-    imports: Optional[str] = None
-
-    # Eval code and store the function
-    def __init__(self, **data):
-        super().__init__(**data)
-
-    # Validate the function
-    @validator("code")
-    def validate_func(cls, v):
-        try:
-            validate.eval_function(v)
-        except Exception as e:
-            raise e
-
-        return v
-
-    def get_function(self):
-        """Get the function"""
-        function_name = validate.extract_function_name(self.code)
-
-        return validate.create_function(self.code, function_name)
-
-
-class PythonFunction(Function):
-    """Python function"""
-
-    code: str
+from pydantic import BaseModel
 
 
 class JsonAgent(BaseModel):
diff --git a/src/backend/langflow/interface/base.py b/src/backend/langflow/interface/base.py
index 4e1bd1fce..3e67830a4 100644
--- a/src/backend/langflow/interface/base.py
+++ b/src/backend/langflow/interface/base.py
@@ -9,11 +9,14 @@ from langflow.template.template import Template, Field, FrontendNode
 
 class LangChainTypeCreator(BaseModel, ABC):
     type_name: str
+    type_dict: Optional[Dict] = None
 
     @property
     @abstractmethod
     def type_to_loader_dict(self) -> Dict:
-        pass
+        if self.type_dict is None:
+            raise NotImplementedError
+        return self.type_dict
 
     @abstractmethod
     def get_signature(self, name: str) -> Optional[Dict[Any, Any]]:
@@ -27,7 +30,10 @@ class LangChainTypeCreator(BaseModel, ABC):
         result: Dict = {self.type_name: {}}
 
         for name in self.to_list():
-            result[self.type_name][name] = self.frontend_node(name).to_dict()
+            # frontend_node.to_dict() returns a dict with the following structure:
+            # {name: {template: {fields}, description: str}}
+            # so we should update the result dict
+            result[self.type_name].update(self.frontend_node(name).to_dict())
 
         return result
 
@@ -45,6 +51,9 @@ class LangChainTypeCreator(BaseModel, ABC):
                 show=value.get("show", True),
                 multiline=value.get("multiline", False),
                 value=value.get("value", None),
+                suffixes=value.get("suffixes", []),
+                file_types=value.get("fileTypes", []),
+                content=value.get("content", None),
             )
             for key, value in signature["template"].items()
             if key != "_type"
@@ -52,7 +61,7 @@ class LangChainTypeCreator(BaseModel, ABC):
         template = Template(type_name=name, fields=fields)
         return FrontendNode(
             template=template,
-            description=signature["description"],
+            description=signature.get("description", ""),
             base_classes=signature["base_classes"],
             name=name,
         )
diff --git a/src/backend/langflow/interface/chains/base.py b/src/backend/langflow/interface/chains/base.py
index 1e0d4fd57..f787ceebb 100644
--- a/src/backend/langflow/interface/chains/base.py
+++ b/src/backend/langflow/interface/chains/base.py
@@ -12,7 +12,9 @@ class ChainCreator(LangChainTypeCreator):
 
     @property
     def type_to_loader_dict(self) -> Dict:
-        return chains_loading.type_to_loader_dict
+        if self.type_dict is None:
+            self.type_dict = chains_loading.type_to_loader_dict
+        return self.type_dict
 
     def get_signature(self, name: str) -> Dict | None:
         try:
diff --git a/src/backend/langflow/interface/custom/__init__.py b/src/backend/langflow/interface/custom/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/backend/langflow/interface/custom/types.py b/src/backend/langflow/interface/custom/types.py
new file mode 100644
index 000000000..40029112f
--- /dev/null
+++ b/src/backend/langflow/interface/custom/types.py
@@ -0,0 +1,37 @@
+from langflow.utils import validate
+from pydantic import BaseModel, validator
+
+
+from typing import Callable, Optional
+
+
+class Function(BaseModel):
+    code: str
+    function: Optional[Callable] = None
+    imports: Optional[str] = None
+
+    # Eval code and store the function
+    def __init__(self, **data):
+        super().__init__(**data)
+
+    # Validate the function
+    @validator("code")
+    def validate_func(cls, v):
+        try:
+            validate.eval_function(v)
+        except Exception as e:
+            raise e
+
+        return v
+
+    def get_function(self):
+        """Get the function"""
+        function_name = validate.extract_function_name(self.code)
+
+        return validate.create_function(self.code, function_name)
+
+
+class PythonFunction(Function):
+    """Python function"""
+
+    code: str
diff --git a/src/backend/langflow/interface/listing.py b/src/backend/langflow/interface/listing.py
index c46c272a5..c2462d15e 100644
--- a/src/backend/langflow/interface/listing.py
+++ b/src/backend/langflow/interface/listing.py
@@ -3,6 +3,7 @@ from langflow.interface.chains.base import chain_creator
 from langflow.interface.llms.base import llm_creator
 from langflow.interface.memories.base import memory_creator
 from langflow.interface.prompts.base import prompt_creator
+from langflow.interface.toolkits.base import toolkits_creator
 from langflow.interface.tools.base import tool_creator
 
 
@@ -14,6 +15,7 @@ def get_type_dict():
         "tools": tool_creator.to_list(),
         "chains": chain_creator.to_list(),
         "memory": memory_creator.to_list(),
+        "toolkits": toolkits_creator.to_list(),
     }
 
 
diff --git a/src/backend/langflow/interface/llms/base.py b/src/backend/langflow/interface/llms/base.py
index 02ca8457c..b00fe3a84 100644
--- a/src/backend/langflow/interface/llms/base.py
+++ b/src/backend/langflow/interface/llms/base.py
@@ -10,7 +10,9 @@ class LLMCreator(LangChainTypeCreator):
 
     @property
     def type_to_loader_dict(self) -> Dict:
-        return llm_type_to_cls_dict
+        if self.type_dict is None:
+            self.type_dict = llm_type_to_cls_dict
+        return self.type_dict
 
     def get_signature(self, name: str) -> Dict | None:
         """Get the signature of an llm."""
diff --git a/src/backend/langflow/interface/memories/base.py b/src/backend/langflow/interface/memories/base.py
index f3dc7279d..d1ae4f3ff 100644
--- a/src/backend/langflow/interface/memories/base.py
+++ b/src/backend/langflow/interface/memories/base.py
@@ -10,7 +10,9 @@ class MemoryCreator(LangChainTypeCreator):
 
     @property
     def type_to_loader_dict(self) -> Dict:
-        return memory_type_to_cls_dict
+        if self.type_dict is None:
+            self.type_dict = memory_type_to_cls_dict
+        return self.type_dict
 
     def get_signature(self, name: str) -> Dict | None:
         """Get the signature of a memory."""
diff --git a/src/backend/langflow/interface/prompts/base.py b/src/backend/langflow/interface/prompts/base.py
index b3f96f69e..de77101f7 100644
--- a/src/backend/langflow/interface/prompts/base.py
+++ b/src/backend/langflow/interface/prompts/base.py
@@ -11,7 +11,9 @@ class PromptCreator(LangChainTypeCreator):
 
     @property
     def type_to_loader_dict(self) -> Dict:
-        return loading.type_to_loader_dict
+        if self.type_dict is None:
+            self.type_dict = loading.type_to_loader_dict
+        return self.type_dict
 
     def get_signature(self, name: str) -> Dict | None:
         try:
diff --git a/src/backend/langflow/interface/toolkits/__init__.py b/src/backend/langflow/interface/toolkits/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/backend/langflow/interface/toolkits/base.py b/src/backend/langflow/interface/toolkits/base.py
new file mode 100644
index 000000000..8eef3d8f8
--- /dev/null
+++ b/src/backend/langflow/interface/toolkits/base.py
@@ -0,0 +1,34 @@
+from langflow.interface.base import LangChainTypeCreator
+from langflow.utils.util import build_template_from_class
+from typing import Dict, List
+from langchain.agents import agent_toolkits
+from langflow.interface.importing.utils import import_class
+
+
+class ToolkitCreator(LangChainTypeCreator):
+    type_name: str = "toolkits"
+
+    @property
+    def type_to_loader_dict(self) -> Dict:
+        if self.type_dict is None:
+            self.type_dict = {
+                toolkit_name: import_class(
+                    f"langchain.agents.agent_toolkits.{toolkit_name}"
+                )
+                # if toolkit_name is not lower case it is a class
+                for toolkit_name in agent_toolkits.__all__
+                if not toolkit_name.islower()
+            }
+        return self.type_dict
+
+    def get_signature(self, name: str) -> Dict | None:
+        try:
+            return build_template_from_class(name, self.type_to_loader_dict)
+        except ValueError as exc:
+            raise ValueError("Prompt not found") from exc
+
+    def to_list(self) -> List[str]:
+        return list(self.type_to_loader_dict.keys())
+
+
+toolkits_creator = ToolkitCreator()
diff --git a/src/backend/langflow/interface/tools/base.py b/src/backend/langflow/interface/tools/base.py
index 496fef878..2e00ea0d2 100644
--- a/src/backend/langflow/interface/tools/base.py
+++ b/src/backend/langflow/interface/tools/base.py
@@ -1,6 +1,9 @@
 from langflow.custom import customs
-from langflow.interface.tools.constants import ALL_TOOLS_NAMES, CUSTOM_TOOLS
-import langflow.interface.tools.util
+from langflow.interface.tools.constants import (
+    ALL_TOOLS_NAMES,
+    CUSTOM_TOOLS,
+    OTHER_TOOLS,
+)
 from langflow.template.template import Field, Template
 from langflow.utils import util
 from langflow.settings import settings
@@ -12,7 +15,11 @@ from langchain.agents.load_tools import (
     _EXTRA_OPTIONAL_TOOLS,
     _LLM_TOOLS,
 )
-from langflow.interface.tools.util import get_tools_dict
+from langflow.interface.tools.util import (
+    get_tool_by_name,
+    get_tools_dict,
+    get_tool_params,
+)
 
 
 class ToolCreator(LangChainTypeCreator):
@@ -32,9 +39,7 @@ class ToolCreator(LangChainTypeCreator):
         base_classes = ["Tool"]
         all_tools = {}
         for tool in self.type_to_loader_dict.keys():
-            if tool_params := langflow.interface.tools.util.get_tool_params(
-                langflow.interface.tools.util.get_tool_by_name(tool)
-            ):
+            if tool_params := get_tool_params(get_tool_by_name(tool)):
                 tool_name = tool_params.get("name") or str(tool)
                 all_tools[tool_name] = {"type": tool, "params": tool_params}
 
@@ -67,6 +72,13 @@ class ToolCreator(LangChainTypeCreator):
                 value="",
                 multiline=True,
             ),
+            "dict_": Field(
+                field_type="file",
+                required=True,
+                is_list=False,
+                show=True,
+                value="",
+            ),
         }
 
         tool_type: str = all_tools[name]["type"]  # type: ignore
@@ -89,6 +101,8 @@ class ToolCreator(LangChainTypeCreator):
             base_classes = ["function"]
             if node := customs.get_custom_nodes("tools").get(tool_type):
                 return node
+        elif tool_type in OTHER_TOOLS:
+            params = all_tools[name]["params"]  # type: ignore
 
         else:
             params = []
@@ -108,9 +122,7 @@ class ToolCreator(LangChainTypeCreator):
 
         template = Template(fields=fields, type_name=tool_type)
 
-        tool_params = langflow.interface.tools.util.get_tool_params(
-            langflow.interface.tools.util.get_tool_by_name(tool_type)
-        )
+        tool_params = get_tool_params(get_tool_by_name(tool_type))
         if tool_params is None:
             tool_params = {}
         return {
@@ -125,9 +137,11 @@ class ToolCreator(LangChainTypeCreator):
         tools = []
 
         for tool in ALL_TOOLS_NAMES:
-            tool_params = langflow.interface.tools.util.get_tool_params(
-                langflow.interface.tools.util.get_tool_by_name(tool)
-            )
+            tool_params = get_tool_params(get_tool_by_name(tool))
+
+            if tool_params and not tool_params.get("name"):
+                tool_params["name"] = tool
+
             if tool_params and (
                 tool_params.get("name") in settings.tools
                 or (tool_params.get("name") and settings.dev)
diff --git a/src/backend/langflow/interface/tools/constants.py b/src/backend/langflow/interface/tools/constants.py
index b1e412e7d..b5db816a5 100644
--- a/src/backend/langflow/interface/tools/constants.py
+++ b/src/backend/langflow/interface/tools/constants.py
@@ -1,6 +1,6 @@
 from langchain.agents.load_tools import get_all_tool_names
 from langchain.agents import Tool
-from langflow.interface.custom_types import PythonFunction
+from langflow.interface.custom.types import PythonFunction
 from langchain.tools.json.tool import JsonSpec
 
 
diff --git a/src/backend/langflow/interface/types.py b/src/backend/langflow/interface/types.py
index f90ebe5a6..0d0d5d595 100644
--- a/src/backend/langflow/interface/types.py
+++ b/src/backend/langflow/interface/types.py
@@ -3,6 +3,7 @@ from langflow.interface.llms.base import llm_creator
 from langflow.interface.memories.base import memory_creator
 from langflow.interface.prompts.base import prompt_creator
 from langflow.interface.chains.base import chain_creator
+from langflow.interface.toolkits.base import toolkits_creator
 from langflow.interface.tools.base import tool_creator
 
 
@@ -30,6 +31,7 @@ def build_langchain_types_dict():
         llm_creator,
         memory_creator,
         tool_creator,
+        toolkits_creator,
     ]
 
     all_types = {}
diff --git a/src/backend/langflow/template/template.py b/src/backend/langflow/template/template.py
index 21b23486a..ee498ef21 100644
--- a/src/backend/langflow/template/template.py
+++ b/src/backend/langflow/template/template.py
@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Union
 from pydantic import BaseModel
 
 
@@ -10,6 +10,9 @@ class Field(BaseModel):
     show: bool = True
     multiline: bool = False
     value: Any = None
+    suffixes: list[str] = []
+    file_types: list[str] = []
+    content: Union[str, None] = None
     # _name will be used to store the name of the field
     # in the template
     name: str = ""
@@ -18,10 +21,16 @@ class Field(BaseModel):
         result = self.dict()
         # Remove key if it is None
         for key in list(result.keys()):
-            if result[key] is None:
+            if result[key] is None or result[key] == []:
                 del result[key]
         result["type"] = result.pop("field_type")
         result["list"] = result.pop("is_list")
+
+        if result.get("file_types"):
+            result["fileTypes"] = result.pop("file_types")
+
+        if self.field_type == "file":
+            result["content"] = self.content
         return result
 
 
diff --git a/src/backend/langflow/utils/util.py b/src/backend/langflow/utils/util.py
index 23bff03ad..3fd4a169e 100644
--- a/src/backend/langflow/utils/util.py
+++ b/src/backend/langflow/utils/util.py
@@ -272,6 +272,8 @@ def format_dict(d, name: Optional[str] = None):
         # Change type from str to Tool
         value["type"] = "Tool" if key in ["allowed_tools"] else _type
 
+        value["type"] = "int" if key in ["max_value_length"] else value["type"]
+
         # Show or not field
         value["show"] = bool(
             (value["required"] and key not in ["input_variables"])
@@ -307,7 +309,10 @@ def format_dict(d, name: Optional[str] = None):
         if "dict" in value["type"].lower():
             value["type"] = "code"
 
-        value["file"] = key in ["dict_"]
+        if key == "dict_":
+            value["type"] = "file"
+            value["suffixes"] = [".json", ".yaml", ".yml"]
+            value["fileTypes"] = ["json", "yaml", "yml"]
 
         # Replace default value with actual value
         if "default" in value:
diff --git a/tests/test_custom_types.py b/tests/test_custom_types.py
index 42da01696..157779ae4 100644
--- a/tests/test_custom_types.py
+++ b/tests/test_custom_types.py
@@ -1,5 +1,5 @@
 # Test this:
-from langflow.interface.custom_types import PythonFunction
+from langflow.interface.custom.types import PythonFunction
 from langflow.utils import constants
 import pytest