From 1909acca5232ca017ada2dcd405bd2c82c1d09b1 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Wed, 19 Jun 2024 18:54:40 -0300
Subject: [PATCH] Refactor input classes to include traceability and metadata
 support

---
 .../base/langflow/inputs/input_mixin.py       |  8 +++++
 src/backend/base/langflow/inputs/inputs.py    | 30 ++++++++++---------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/backend/base/langflow/inputs/input_mixin.py b/src/backend/base/langflow/inputs/input_mixin.py
index d9107d272..411c98190 100644
--- a/src/backend/base/langflow/inputs/input_mixin.py
+++ b/src/backend/base/langflow/inputs/input_mixin.py
@@ -87,6 +87,14 @@ class BaseInputMixin(BaseModel, validate_assignment=True):
         return dump
 
 
+class InputTraceMixin(BaseModel):
+    trace_as_input: bool = True
+
+
+class MetadataTraceMixin(BaseModel):
+    trace_as_metadata: bool = True
+
+
 # Mixin for input fields that can be listable
 class ListableInputMixin(BaseModel):
     is_list: bool = Field(default=False, serialization_alias="list")
diff --git a/src/backend/base/langflow/inputs/inputs.py b/src/backend/base/langflow/inputs/inputs.py
index e9c415c9a..0449dcda1 100644
--- a/src/backend/base/langflow/inputs/inputs.py
+++ b/src/backend/base/langflow/inputs/inputs.py
@@ -13,14 +13,16 @@ from .input_mixin import (
     DropDownMixin,
     FieldTypes,
     FileMixin,
+    InputTraceMixin,
     ListableInputMixin,
+    MetadataTraceMixin,
     MultilineMixin,
     RangeMixin,
     SerializableFieldTypes,
 )
 
 
-class HandleInput(BaseInputMixin, ListableInputMixin):
+class HandleInput(BaseInputMixin, ListableInputMixin, MetadataTraceMixin):
     """
     Represents an Input that has a Handle to a specific type (e.g. BaseLanguageModel, BaseRetriever, etc.)
 
@@ -35,7 +37,7 @@ class HandleInput(BaseInputMixin, ListableInputMixin):
     field_type: Optional[SerializableFieldTypes] = FieldTypes.OTHER
 
 
-class DataInput(HandleInput):
+class DataInput(HandleInput, InputTraceMixin):
     """
     Represents an Input that has a Handle that receives a Data object.
 
@@ -46,12 +48,12 @@ class DataInput(HandleInput):
     input_types: list[str] = ["Data"]
 
 
-class PromptInput(BaseInputMixin, ListableInputMixin):
+class PromptInput(BaseInputMixin, ListableInputMixin, InputTraceMixin):
     field_type: Optional[SerializableFieldTypes] = FieldTypes.PROMPT
 
 
 # Applying mixins to a specific input type
-class StrInput(BaseInputMixin, ListableInputMixin, DatabaseLoadMixin):
+class StrInput(BaseInputMixin, ListableInputMixin, DatabaseLoadMixin, MetadataTraceMixin):
     field_type: Optional[SerializableFieldTypes] = FieldTypes.TEXT
     load_from_db: CoalesceBool = False
     """Defines if the field will allow the user to open a text editor. Default is False."""
@@ -101,7 +103,7 @@ class StrInput(BaseInputMixin, ListableInputMixin, DatabaseLoadMixin):
         return value
 
 
-class MessageInput(StrInput):
+class MessageInput(StrInput, InputTraceMixin):
     input_types: list[str] = ["Message"]
 
     @staticmethod
@@ -114,7 +116,7 @@ class MessageInput(StrInput):
         raise ValueError(f"Invalid value type {type(v)}")
 
 
-class TextInput(StrInput):
+class TextInput(StrInput, MetadataTraceMixin, InputTraceMixin):
     """
     Represents a text input component for the Langflow system.
 
@@ -163,7 +165,7 @@ class TextInput(StrInput):
         return value
 
 
-class MultilineInput(TextInput, MultilineMixin):
+class MultilineInput(TextInput, MultilineMixin, InputTraceMixin):
     """
     Represents a multiline input field.
 
@@ -194,7 +196,7 @@ class SecretStrInput(BaseInputMixin, DatabaseLoadMixin):
     load_from_db: CoalesceBool = True
 
 
-class IntInput(BaseInputMixin, ListableInputMixin, RangeMixin):
+class IntInput(BaseInputMixin, ListableInputMixin, RangeMixin, MetadataTraceMixin):
     """
     Represents an integer field.
 
@@ -208,7 +210,7 @@ class IntInput(BaseInputMixin, ListableInputMixin, RangeMixin):
     field_type: Optional[SerializableFieldTypes] = FieldTypes.INTEGER
 
 
-class FloatInput(BaseInputMixin, ListableInputMixin, RangeMixin):
+class FloatInput(BaseInputMixin, ListableInputMixin, RangeMixin, MetadataTraceMixin):
     """
     Represents a float field.
 
@@ -222,7 +224,7 @@ class FloatInput(BaseInputMixin, ListableInputMixin, RangeMixin):
     field_type: Optional[SerializableFieldTypes] = FieldTypes.FLOAT
 
 
-class BoolInput(BaseInputMixin, ListableInputMixin):
+class BoolInput(BaseInputMixin, ListableInputMixin, MetadataTraceMixin):
     """
     Represents a boolean field.
 
@@ -238,7 +240,7 @@ class BoolInput(BaseInputMixin, ListableInputMixin):
     value: CoalesceBool = False
 
 
-class NestedDictInput(BaseInputMixin, ListableInputMixin):
+class NestedDictInput(BaseInputMixin, ListableInputMixin, MetadataTraceMixin, InputTraceMixin):
     """
     Represents a nested dictionary field.
 
@@ -254,7 +256,7 @@ class NestedDictInput(BaseInputMixin, ListableInputMixin):
     value: Optional[dict] = {}
 
 
-class DictInput(BaseInputMixin, ListableInputMixin):
+class DictInput(BaseInputMixin, ListableInputMixin, InputTraceMixin):
     """
     Represents a dictionary field.
 
@@ -270,7 +272,7 @@ class DictInput(BaseInputMixin, ListableInputMixin):
     value: Optional[dict] = {}
 
 
-class DropdownInput(BaseInputMixin, DropDownMixin):
+class DropdownInput(BaseInputMixin, DropDownMixin, MetadataTraceMixin):
     """
     Represents a dropdown input field.
 
@@ -287,7 +289,7 @@ class DropdownInput(BaseInputMixin, DropDownMixin):
     options: list[str] = Field(default_factory=list)
 
 
-class FileInput(BaseInputMixin, ListableInputMixin, FileMixin):
+class FileInput(BaseInputMixin, ListableInputMixin, FileMixin, MetadataTraceMixin):
     """
     Represents a file field.