From c7a3624f21f5529892b66e9f568f178697bee36f Mon Sep 17 00:00:00 2001
From: ogabrielluiz <gabriel@langflow.org>
Date: Mon, 17 Jun 2024 11:32:41 -0300
Subject: [PATCH] Refactor inputs.py to add MessageInput and TextInput classes

---
 src/backend/base/langflow/inputs/__init__.py | 12 ++++++------
 src/backend/base/langflow/inputs/inputs.py   | 16 +++++++++++++++-
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/backend/base/langflow/inputs/__init__.py b/src/backend/base/langflow/inputs/__init__.py
index 6939c470c..10f297b59 100644
--- a/src/backend/base/langflow/inputs/__init__.py
+++ b/src/backend/base/langflow/inputs/__init__.py
@@ -17,12 +17,12 @@ from .inputs import (
 __all__ = [
     "SecretStrInput",
     "StrInput",
-    "IntInput",
-    "FloatInput",
-    "BoolInput",
-    "NestedDictInput",
-    "DictInput",
-    "DropdownInput",
+    MultilineInput,
+    NestedDictInput,
+    PromptInput,
+    SecretStrInput,
+    StrInput,
+    TextInput,
     "FileInput",
     "PromptInput",
     "MultilineInput",
diff --git a/src/backend/base/langflow/inputs/inputs.py b/src/backend/base/langflow/inputs/inputs.py
index 420e3468b..c84feed20 100644
--- a/src/backend/base/langflow/inputs/inputs.py
+++ b/src/backend/base/langflow/inputs/inputs.py
@@ -40,8 +40,22 @@ class StrInput(BaseInputMixin, ListableInputMixin, DatabaseLoadMixin):  # noqa:
     """Defines if the field will allow the user to open a text editor. Default is False."""
 
 
+class MessageInput(StrInput):
+    input_types: list[str] = ["Message"]
+
+    @field_validator("value")
+    @classmethod
+    def validate_value(cls, v: Any, _info):
+        # If v is a instance of Message, then its fine
+        if isinstance(v, Message):
+            return v
+        if isinstance(v, str):
+            return Message(text=v)
+        raise ValueError(f"Invalid value type {type(v)}")
+
+
 class TextInput(StrInput):
-    input_types: list[str] = ["Data", "Message", "Text"]
+    input_types: list[str] = ["Message"]
 
     @field_validator("value")
     @classmethod