Update StoreMessageComponent for enhanced message handling

src/backend/base/langflow/components/helpers/StoreMessage.py

@@ -1,5 +1,5 @@
 from langflow.custom import Component
-from langflow.inputs import MessageTextInput, StrInput
+from langflow.inputs import MessageInput, StrInput
 from langflow.schema.message import Message
 from langflow.template import Output
 from langflow.memory import get_messages, store_message
@@ -11,22 +11,25 @@ class StoreMessageComponent(Component):
     icon = "save"
 
     inputs = [
-        MessageTextInput(
+        MessageInput(
             name="message",
             display_name="Message",
             info="The chat message to be stored.",
-            input_types=["Message", "str"],
-            required=True,
+            required=True
         ),
         StrInput(
             name="sender",
             display_name="Sender",
             info="The sender of the message.",
-            value="",
+            value="AI",
             advanced=True,
         ),
         StrInput(
-            name="sender_name", display_name="Sender Name", info="The name of the sender.", value="", advanced=True
+            name="sender_name",
+            display_name="Sender Name",
+            info="The name of the sender.",
+            value="AI",
+            advanced=True
         ),
         StrInput(
             name="session_id",
@@ -41,26 +44,13 @@ class StoreMessageComponent(Component):
     ]
 
     def store_message(self) -> Message:
-        if isinstance(self.message, str):
-            if not self.session_id:
-                raise ValueError("If passing a text, Session ID cannot be empty.")
-            message = Message(
-                text=self.message, sender=self.sender, sender_name=self.sender_name, session_id=self.session_id
-            )
+        message = self.message
 
-        elif isinstance(self.message, Message):
-            message = self.message
-            if self.session_id:
-                message.session_id = self.session_id
-            if self.sender:
-                message.sender = self.sender
-            if self.sender_name:
-                message.sender_name = self.sender_name
-        else:
-            raise ValueError("Message should be either string or Message.")
+        message.session_id = self.session_id or message.session_id
+        message.sender = self.sender or message.sender
+        message.sender_name = self.sender_name or message.sender_name
 
         store_message(message, flow_id=self.graph.flow_id)
-
         stored = get_messages(session_id=message.session_id, sender_name=message.sender_name, sender=message.sender)
         self.status = stored
         return stored

src/backend/base/langflow/components/prototypes/Pass.py

@@ -1,31 +1,31 @@
-from typing import Union
+from langflow.custom import Component
+from langflow.io import MessageInput
+from langflow.schema.message import Message
+from langflow.template import Output
 
-from langflow.custom import CustomComponent
-from langflow.field_typing import Text
-from langflow.schema import Data
-
-
-class PassComponent(CustomComponent):
+class PassMessageComponent(Component):
     display_name = "Pass"
-    description = "A pass-through component that forwards the second input while ignoring the first, used for controlling workflow direction."
-    field_order = ["ignored_input", "forwarded_input"]
-    beta = True
+    description = "Forwards the input message, unchanged."
+    icon = "arrow-right"
 
-    def build_config(self) -> dict:
-        return {
-            "ignored_input": {
-                "display_name": "Ignored Input",
-                "info": "This input is ignored. It's used to control the flow in the graph.",
-                "input_types": ["Text", "Data"],
-            },
-            "forwarded_input": {
-                "display_name": "Input",
-                "info": "This input is forwarded by the component.",
-                "input_types": ["Text", "Data"],
-            },
-        }
+    inputs = [
+        MessageInput(
+            name="input_message",
+            display_name="Input Message",
+            info="The message to be passed forward.",
+        ),
+        MessageInput(
+            name="ignored_message",
+            display_name="Ignored Message",
+            info="A second message to be ignored. Used as a workaround for continuity.",
+            advanced=True
+        ),
+    ]
 
-    def build(self, ignored_input: Text, forwarded_input: Text) -> Union[Text, Data]:
-        # The ignored_input is not used in the logic, it's just there for graph flow control
-        self.status = forwarded_input
-        return forwarded_input
+    outputs = [
+        Output(display_name="Output Message", name="output_message", method="pass_message"),
+    ]
+
+    def pass_message(self) -> Message:
+        self.status = self.input_message
+        return self.input_message