From 84d3bd1ba6b86418feec288ea197268973c27f1d Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 26 Feb 2024 18:53:10 -0300
Subject: [PATCH] Refactor ChatInput and ChatOutput to support returning
 messages as records

---
 .../langflow/components/io/ChatInput.py       | 37 ++++++++++--------
 .../langflow/components/io/ChatOutput.py      | 39 +++++++++++--------
 2 files changed, 44 insertions(+), 32 deletions(-)

diff --git a/src/backend/langflow/components/io/ChatInput.py b/src/backend/langflow/components/io/ChatInput.py
index 704ab732b..4d8fc509c 100644
--- a/src/backend/langflow/components/io/ChatInput.py
+++ b/src/backend/langflow/components/io/ChatInput.py
@@ -1,7 +1,6 @@
-from typing import Optional, Union
+from typing import Optional
 
 from langflow import CustomComponent
-from langflow.field_typing import Text
 from langflow.schema import Record
 
 
@@ -25,6 +24,10 @@ class ChatInput(CustomComponent):
                 "display_name": "Session ID",
                 "info": "Session ID of the chat history.",
             },
+            "return_record": {
+                "display_name": "Return Record",
+                "info": "Return the message as a record containing the sender, sender_name, and session_id.",
+            },
         }
 
     def build(
@@ -33,21 +36,23 @@ class ChatInput(CustomComponent):
         sender_name: Optional[str] = "User",
         message: Optional[str] = None,
         session_id: Optional[str] = None,
+        return_record: Optional[bool] = False,
     ) -> Record:
-        if isinstance(message, Record):
-            # Update the data of the record
-            message.data["sender"] = sender
-            message.data["sender_name"] = sender_name
-            message.data["session_id"] = session_id
-        else:
-            message = Record(
-                text=message,
-                data={
-                    "sender": sender,
-                    "sender_name": sender_name,
-                    "session_id": session_id,
-                },
-            )
+        if return_record:
+            if isinstance(message, Record):
+                # Update the data of the record
+                message.data["sender"] = sender
+                message.data["sender_name"] = sender_name
+                message.data["session_id"] = session_id
+            else:
+                message = Record(
+                    text=message,
+                    data={
+                        "sender": sender,
+                        "sender_name": sender_name,
+                        "session_id": session_id,
+                    },
+                )
         if not message:
             message = ""
         self.status = message
diff --git a/src/backend/langflow/components/io/ChatOutput.py b/src/backend/langflow/components/io/ChatOutput.py
index 6812c6d34..05639cdb2 100644
--- a/src/backend/langflow/components/io/ChatOutput.py
+++ b/src/backend/langflow/components/io/ChatOutput.py
@@ -1,6 +1,7 @@
-from typing import Optional
+from typing import Optional, Union
 
 from langflow import CustomComponent
+from langflow.field_typing import Text
 from langflow.schema import Record
 
 
@@ -27,6 +28,10 @@ class ChatOutput(CustomComponent):
                 "info": "Session ID of the chat history.",
                 "input_types": ["Text"],
             },
+            "return_record": {
+                "display_name": "Return Record",
+                "info": "Return the message as a record containing the sender, sender_name, and session_id.",
+            },
         }
 
     def build(
@@ -35,21 +40,23 @@ class ChatOutput(CustomComponent):
         sender_name: Optional[str] = "AI",
         session_id: Optional[str] = None,
         message: Optional[str] = None,
-    ) -> Record:
-        if isinstance(message, Record):
-            # Update the data of the record
-            message.data["sender"] = sender
-            message.data["sender_name"] = sender_name
-            message.data["session_id"] = session_id
-        else:
-            message = Record(
-                text=message,
-                data={
-                    "sender": sender,
-                    "sender_name": sender_name,
-                    "session_id": session_id,
-                },
-            )
+        return_record: Optional[bool] = False,
+    ) -> Union[Text, Record]:
+        if return_record:
+            if isinstance(message, Record):
+                # Update the data of the record
+                message.data["sender"] = sender
+                message.data["sender_name"] = sender_name
+                message.data["session_id"] = session_id
+            else:
+                message = Record(
+                    text=message,
+                    data={
+                        "sender": sender,
+                        "sender_name": sender_name,
+                        "session_id": session_id,
+                    },
+                )
         if not message:
             message = ""
         self.status = message