From db17aceb3b8e59181e0b378501c222b87a96ff68 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 3 Jul 2023 16:37:37 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix(base.py):=20add=20output=5Ft?=
 =?UTF-8?q?ypes=20field=20to=20FrontendNode=20class=20to=20support=20speci?=
 =?UTF-8?q?fying=20output=20types=20=F0=9F=94=A7=20fix(documentloaders.py)?=
 =?UTF-8?q?:=20add=5Fextra=5Fbase=5Fclasses=20method=20to=20set=20base=5Fc?=
 =?UTF-8?q?lasses=20and=20output=5Ftypes=20for=20DocumentLoaderFrontNode?=
 =?UTF-8?q?=20=F0=9F=94=A7=20fix(textsplitters.py):=20add=5Fextra=5Fbase?=
 =?UTF-8?q?=5Fclasses=20method=20to=20set=20base=5Fclasses=20and=20output?=
 =?UTF-8?q?=5Ftypes=20for=20TextSplittersFrontendNode=20=F0=9F=94=A7=20fix?=
 =?UTF-8?q?(vectorstores.py):=20change=20field=5Ftype=20to=20"Document"=20?=
 =?UTF-8?q?for=20documents=20field=20in=20VectorStoreFrontendNode=20?=
 =?UTF-8?q?=F0=9F=94=A7=20fix(index.ts):=20add=20output=5Ftypes=20field=20?=
 =?UTF-8?q?to=20APIClassType=20interface=20The=20changes=20were=20made=20t?=
 =?UTF-8?q?o=20add=20support=20for=20specifying=20output=20types=20in=20th?=
 =?UTF-8?q?e=20FrontendNode=20class=20and=20its=20subclasses.=20This=20all?=
 =?UTF-8?q?ows=20for=20more=20flexibility=20in=20defining=20the=20behavior?=
 =?UTF-8?q?=20and=20capabilities=20of=20the=20frontend=20nodes=20in=20the?=
 =?UTF-8?q?=20application.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/template/frontend_node/base.py         | 2 ++
 .../langflow/template/frontend_node/documentloaders.py      | 4 ++++
 .../langflow/template/frontend_node/textsplitters.py        | 6 +++++-
 src/backend/langflow/template/frontend_node/vectorstores.py | 2 +-
 src/frontend/src/types/api/index.ts                         | 1 +
 5 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/backend/langflow/template/frontend_node/base.py b/src/backend/langflow/template/frontend_node/base.py
index 47c8c53f6..55773db72 100644
--- a/src/backend/langflow/template/frontend_node/base.py
+++ b/src/backend/langflow/template/frontend_node/base.py
@@ -18,6 +18,7 @@ class FrontendNode(BaseModel):
     name: str = ""
     display_name: str = ""
     custom_fields: List[str] = []
+    output_types: List[str] = []
 
     def process_base_classes(self) -> None:
         """Removes unwanted base classes from the list of base classes."""
@@ -36,6 +37,7 @@ class FrontendNode(BaseModel):
                 "base_classes": self.base_classes,
                 "display_name": self.display_name or self.name,
                 "custom_fields": self.custom_fields,
+                "output_types": self.output_types,
             },
         }
 
diff --git a/src/backend/langflow/template/frontend_node/documentloaders.py b/src/backend/langflow/template/frontend_node/documentloaders.py
index 501aa361e..ff9cac29f 100644
--- a/src/backend/langflow/template/frontend_node/documentloaders.py
+++ b/src/backend/langflow/template/frontend_node/documentloaders.py
@@ -19,6 +19,10 @@ def build_file_field(
 
 
 class DocumentLoaderFrontNode(FrontendNode):
+    def add_extra_base_classes(self) -> None:
+        self.base_classes = ["Document"]
+        self.output_types = ["Document"]
+
     file_path_templates = {
         "AirbyteJSONLoader": build_file_field(suffixes=[".json"], fileTypes=["json"]),
         "CoNLLULoader": build_file_field(suffixes=[".csv"], fileTypes=["csv"]),
diff --git a/src/backend/langflow/template/frontend_node/textsplitters.py b/src/backend/langflow/template/frontend_node/textsplitters.py
index 03880379d..7e5f7a820 100644
--- a/src/backend/langflow/template/frontend_node/textsplitters.py
+++ b/src/backend/langflow/template/frontend_node/textsplitters.py
@@ -3,10 +3,14 @@ from langflow.template.frontend_node.base import FrontendNode
 
 
 class TextSplittersFrontendNode(FrontendNode):
+    def add_extra_base_classes(self) -> None:
+        self.base_classes = ["Document"]
+        self.output_types = ["Document"]
+
     def add_extra_fields(self) -> None:
         self.template.add_field(
             TemplateField(
-                field_type="BaseLoader",
+                field_type="Document",
                 required=True,
                 show=True,
                 name="documents",
diff --git a/src/backend/langflow/template/frontend_node/vectorstores.py b/src/backend/langflow/template/frontend_node/vectorstores.py
index 01b6bfe53..9903b13c2 100644
--- a/src/backend/langflow/template/frontend_node/vectorstores.py
+++ b/src/backend/langflow/template/frontend_node/vectorstores.py
@@ -252,7 +252,7 @@ class VectorStoreFrontendNode(FrontendNode):
             # when instantiating the vectorstores
             field.name = "documents"
 
-            field.field_type = "TextSplitter"
+            field.field_type = "Document"
             field.display_name = "Documents"
             field.required = False
             field.show = True
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index d54a8165c..91a3dcbc2 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -13,6 +13,7 @@ export type APIClassType = {
   template: APITemplateType;
   display_name: string;
   input_types?: Array<string>;
+  output_types?: Array<string>;
   [key: string]: Array<string> | string | APITemplateType;
 };