diff --git a/src/backend/base/langflow/components/helpers/StoreMessage.py b/src/backend/base/langflow/components/helpers/StoreMessage.py index 5d0abfbb9..6e37573a8 100644 --- a/src/backend/base/langflow/components/helpers/StoreMessage.py +++ b/src/backend/base/langflow/components/helpers/StoreMessage.py @@ -1,22 +1,68 @@ -from langflow.custom import CustomComponent -from langflow.memory import get_messages, store_message +from langflow.custom import Component +from langflow.inputs import MessageTextInput, StrInput from langflow.schema.message import Message +from langflow.template import Output +from langflow.memory import get_messages, store_message - -class StoreMessageComponent(CustomComponent): +class StoreMessageComponent(Component): display_name = "Store Message" - description = "Stores a chat message." + description = "Stores a chat message or text." + icon = "save" - def build_config(self): - return { - "message": {"display_name": "Message"}, - } + inputs = [ + MessageTextInput( + name="message", + display_name="Message", + info="The chat message to be stored.", + input_types=["Message", "str"], + required=True + ), + StrInput( + name="sender", + display_name="Sender", + info="The sender of the message.", + value="", + advanced=True, + ), + StrInput( + name="sender_name", + display_name="Sender Name", + info="The name of the sender.", + value="", + advanced=True + ), + StrInput( + name="session_id", + display_name="Session ID", + info="The session ID of the chat.", + value="", + + ), + ] + + outputs = [ + Output(display_name="Stored Messages", name="stored_messages", method="store_message"), + ] + + 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) + + 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.") - def build( - self, - message: Message, - ) -> Message: store_message(message, flow_id=self.graph.flow_id) - self.status = get_messages() - return message + stored = get_messages(session_id=message.session_id, sender_name=message.sender_name, sender=message.sender) + self.status = stored + return stored