Refactor store_message function and add StoreMessageComponent (#1817)

* Refactor store_message function in chat.py and memory.py

* Refactor store_message function in chat.py and memory.py

* Add StoreMessageComponent to langflow components

* Refactor store_message function in chat.py and memory.py to require session_id, sender, and sender_name

* Refactor StoreMessageComponent to use Optional[str] for sender_name and session_id parameters

---------

Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

src/backend/base/langflow/base/io/chat.py

@@ -1,10 +1,9 @@
-import warnings
 from typing import Optional, Union
 
 from langflow.field_typing import Text
 from langflow.helpers.record import records_to_text
 from langflow.interface.custom.custom_component import CustomComponent
-from langflow.memory import add_messages
+from langflow.memory import store_message
 from langflow.schema import Record
 
 
@@ -50,34 +49,16 @@ class ChatComponent(CustomComponent):
         sender: Optional[str] = None,
         sender_name: Optional[str] = None,
     ) -> list[Record]:
-        if not message:
-            warnings.warn("No message provided.")
-            return []
 
-        if not session_id or not sender or not sender_name:
+        records = store_message(
-            raise ValueError("All of session_id, sender, and sender_name must be provided.")
+            message,
-        if isinstance(message, Record):
+            session_id=session_id,
-            record = message
+            sender=sender,
-            record.data.update(
+            sender_name=sender_name,
-                {
-                    "session_id": session_id,
-                    "sender": sender,
-                    "sender_name": sender_name,
-                }
-            )
-        else:
-            record = Record(
-                data={
-                    "text": message,
-                    "session_id": session_id,
-                    "sender": sender,
-                    "sender_name": sender_name,
-                },
         )
 
-        self.status = record
+        self.status = records
-        records = add_messages([record])
+        return records
-        return records[0]
 
     def build_with_record(
         self,

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

@@ -0,0 +1,44 @@
+from typing import List, Optional
+
+from langflow.interface.custom.custom_component import CustomComponent
+from langflow.memory import get_messages, store_message
+from langflow.schema import Record
+
+
+class StoreMessageComponent(CustomComponent):
+    display_name = "Store Message"
+    description = "Stores a chat message given a Session ID."
+    beta: bool = True
+
+    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[Record]:
+
+        store_message(
+            sender=sender,
+            sender_name=sender_name,
+            session_id=session_id,
+            message=message,
+        )
+
+        self.status = get_messages(session_id=session_id)
+        return get_messages(session_id=session_id)

src/backend/base/langflow/memory.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Optional, Union
 
 from loguru import logger
@@ -51,6 +52,7 @@ def get_messages(
                 "sender": row.sender,
                 "sender_name": row.sender_name,
                 "session_id": row.session_id,
+                "timestamp": row.timestamp,
             },
         )
         records.append(record)
@@ -98,3 +100,39 @@ def delete_messages(session_id: str):
     """
     monitor_service = get_monitor_service()
     monitor_service.delete_messages(session_id)
+
+
+def store_message(
+    message: Union[str, Record],
+    session_id: Optional[str] = None,
+    sender: Optional[str] = None,
+    sender_name: Optional[str] = None,
+) -> list[Record]:
+
+    if not message:
+        warnings.warn("No message provided.")
+        return []
+
+    if not session_id or not sender or not sender_name:
+        raise ValueError("All of session_id, sender, and sender_name must be provided.")
+
+    if isinstance(message, Record):
+        record = message
+        record.data.update(
+            {
+                "session_id": session_id,
+                "sender": sender,
+                "sender_name": sender_name,
+            }
+        )
+    elif isinstance(message, str):
+        record = Record(
+            data={
+                "text": message,
+                "session_id": session_id,
+                "sender": sender,
+                "sender_name": sender_name,
+            },
+        )
+
+    return add_messages([record])