From f417378779576a808d592bc911188b82644fc1ed Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 19 Feb 2024 19:00:51 -0300
Subject: [PATCH] Add 'as_record' parameter to ChatOutput class

---
 src/backend/langflow/components/io/ChatOutput.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/backend/langflow/components/io/ChatOutput.py b/src/backend/langflow/components/io/ChatOutput.py
index 699138f9f..a78baf2c2 100644
--- a/src/backend/langflow/components/io/ChatOutput.py
+++ b/src/backend/langflow/components/io/ChatOutput.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Union
 
 from langflow import CustomComponent
 from langflow.field_typing import Text
@@ -28,6 +28,10 @@ class ChatOutput(CustomComponent):
                 "info": "Session ID of the chat history.",
                 "input_types": ["Text"],
             },
+            "as_record": {
+                "display_name": "As Record",
+                "info": "If true, the message will be returned as a Record.",
+            },
         }
 
     def build(
@@ -37,7 +41,7 @@ class ChatOutput(CustomComponent):
         session_id: Optional[str] = None,
         message: Optional[str] = None,
         as_record: Optional[bool] = False,
-    ) -> Text:
+    ) -> Union[Text, Record]:
         self.status = message
         if as_record:
             if isinstance(message, Record):