From 40baf204244c6af9b4038d09d4eac1d06003cbcc Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Tue, 27 Feb 2024 17:51:02 -0300
Subject: [PATCH] Refactor ChatComponent's add_message method

---
 src/backend/langflow/io/schema.py | 39 ++++++++++++++++---------------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/src/backend/langflow/io/schema.py b/src/backend/langflow/io/schema.py
index d83d3472e..4660b4276 100644
--- a/src/backend/langflow/io/schema.py
+++ b/src/backend/langflow/io/schema.py
@@ -48,27 +48,28 @@ class ChatComponent(CustomComponent):
             raise ValueError(
                 "All of session_id, sender, and sender_name must be provided."
             )
-
-        if not record:
-            record = []
-            if not session_id or not sender or not sender_name:
-                raise ValueError
-            for text in text:
-                record = Record(
-                    text=text,
-                    data={
-                        "session_id": session_id,
-                        "sender": sender,
-                        "sender_name": sender_name,
-                    },
-                )
-                record.append(record)
-        elif isinstance(record, Record):
-            record = [record]
+        if isinstance(message, Record):
+            record = message
+            record.data.update(
+                {
+                    "session_id": session_id,
+                    "sender": sender,
+                    "sender_name": sender_name,
+                }
+            )
+        else:
+            record = Record(
+                text=message,
+                data={
+                    "session_id": session_id,
+                    "sender": sender,
+                    "sender_name": sender_name,
+                },
+            )
 
         self.status = record
-        record = add_messages(record)
-        return record
+        records = add_messages([record])
+        return records[0]
 
     def build(
         self,