diff --git a/src/backend/base/langflow/components/helpers/StoreMessage.py b/src/backend/base/langflow/components/helpers/StoreMessage.py index 4864bf4a7..90e2c2769 100644 --- a/src/backend/base/langflow/components/helpers/StoreMessage.py +++ b/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 diff --git a/src/backend/base/langflow/components/prototypes/Pass.py b/src/backend/base/langflow/components/prototypes/Pass.py index 4e2f234c9..a84164380 100644 --- a/src/backend/base/langflow/components/prototypes/Pass.py +++ b/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