From b5b02e8a95448e3ba83e56cd6a4ce1ce3efa4128 Mon Sep 17 00:00:00 2001
From: Rodrigo <rodrigosilvanader@gmail.com>
Date: Mon, 10 Jun 2024 18:48:09 -0300
Subject: [PATCH] feat: Add Message and StoreMessage components

The code changes introduce two new components, `MessageComponent` and `StoreMessageComponent`, to the `langflow` project.

The `MessageComponent` creates a `Message` object given a session ID, allowing for the creation of chat messages with customizable sender type, sender name, and text.

The `StoreMessageComponent` stores a chat message, utilizing the `store_message` function and the `get_messages` function to retrieve the stored messages.

These changes aim to enhance the chat functionality of the project.
---
 .../components/experimental/Message.py        | 38 +++++++++++++++++++
 .../components/experimental/StoreMessage.py   | 32 +++-------------
 src/backend/base/langflow/memory.py           |  8 +++-
 3 files changed, 51 insertions(+), 27 deletions(-)
 create mode 100644 src/backend/base/langflow/components/experimental/Message.py

diff --git a/src/backend/base/langflow/components/experimental/Message.py b/src/backend/base/langflow/components/experimental/Message.py
new file mode 100644
index 000000000..9fdd8a0da
--- /dev/null
+++ b/src/backend/base/langflow/components/experimental/Message.py
@@ -0,0 +1,38 @@
+from typing import Optional
+
+from langflow.custom import CustomComponent
+from langflow.schema.message import Message
+
+
+class MessageComponent(CustomComponent):
+    display_name = "Message"
+    description = "Creates a Message object given a Session ID."
+
+    def build_config(self):
+        return {
+            "sender": {
+                "options": ["Machine", "User"],
+                "display_name": "Sender Type",
+            },
+            "sender_name": {"display_name": "Sender Name"},
+            "text": {"display_name": "Text"},
+            "session_id": {
+                "display_name": "Session ID",
+                "info": "Session ID of the chat history.",
+                "input_types": ["Text"],
+            },
+        }
+
+    def build(
+        self,
+        sender: str = "User",
+        sender_name: Optional[str] = None,
+        session_id: Optional[str] = None,
+        text: str = "",
+    ) -> Message:
+        message = Message(
+            text=text, sender=sender, sender_name=sender_name, flow_id=self.graph.flow_id, session_id=session_id
+        )
+
+        self.status = message
+        return message
diff --git a/src/backend/base/langflow/components/experimental/StoreMessage.py b/src/backend/base/langflow/components/experimental/StoreMessage.py
index 19be36068..5d0abfbb9 100644
--- a/src/backend/base/langflow/components/experimental/StoreMessage.py
+++ b/src/backend/base/langflow/components/experimental/StoreMessage.py
@@ -1,5 +1,3 @@
-from typing import List, Optional
-
 from langflow.custom import CustomComponent
 from langflow.memory import get_messages, store_message
 from langflow.schema.message import Message
@@ -7,36 +5,18 @@ from langflow.schema.message import Message
 
 class StoreMessageComponent(CustomComponent):
     display_name = "Store Message"
-    description = "Stores a chat message given a Session ID."
-    beta: bool = True
+    description = "Stores a chat message."
 
     def build_config(self):
         return {
-            "sender": {
-                "options": ["Machine", "User"],
-                "display_name": "Sender Type",
-            },
-            "sender_name": {"display_name": "Sender Name"},
             "message": {"display_name": "Message"},
-            "session_id": {
-                "display_name": "Session ID",
-                "info": "Session ID of the chat history.",
-                "input_types": ["Text"],
-            },
         }
 
     def build(
         self,
-        sender: str = "User",
-        sender_name: Optional[str] = None,
-        session_id: Optional[str] = None,
-        message: str = "",
-    ) -> List[Message]:
-        store_message(
-            message=Message(
-                text=message, sender=sender, sender_name=sender_name, flow_id=self.graph.flow_id, session_id=session_id
-            )
-        )
+        message: Message,
+    ) -> Message:
+        store_message(message, flow_id=self.graph.flow_id)
+        self.status = get_messages()
 
-        self.status = get_messages(session_id=session_id)
-        return get_messages(session_id=session_id)
+        return message
diff --git a/src/backend/base/langflow/memory.py b/src/backend/base/langflow/memory.py
index 9a07f4bc1..e812f449c 100644
--- a/src/backend/base/langflow/memory.py
+++ b/src/backend/base/langflow/memory.py
@@ -46,7 +46,13 @@ def get_messages(
     # so we need to reverse the order
     messages_df = messages_df[::-1] if order == "DESC" else messages_df
     for row in messages_df.itertuples():
-        msg = Message(text=row.text, sender=row.sender, sender_name=row.sender_name, timestamp=row.timestamp)
+        msg = Message(
+            text=row.text,
+            sender=row.sender,
+            session_id=row.session_id,
+            sender_name=row.sender_name,
+            timestamp=row.timestamp,
+        )
 
         messages.append(msg)