feat: updated ChatOutput component that can accept Data, Dataframe and Message (#6643)

* update chatoutput

* [autofix.ci] apply automated fixes

* update

* tests

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* Update chat.py

* update template

* fix lint errors

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* fix: rename variable for clarity in chat output component test

* [autofix.ci] apply automated fixes

* fix: enable loading from database for API key in starter project configurations

* update templates

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Co-authored-by: cristhianzl <cristhian.lousa@gmail.com>

src/backend/base/langflow/components/outputs/chat.py

@@ -1,6 +1,11 @@
+from typing import Any
+
 from langflow.base.io.chat import ChatComponent
 from langflow.inputs import BoolInput
-from langflow.io import DropdownInput, MessageInput, MessageTextInput, Output
+from langflow.inputs.inputs import HandleInput
+from langflow.io import DropdownInput, MessageTextInput, Output
+from langflow.schema.data import Data
+from langflow.schema.dataframe import DataFrame
 from langflow.schema.message import Message
 from langflow.schema.properties import Source
 from langflow.utils.constants import (
@@ -18,10 +23,12 @@ class ChatOutput(ChatComponent):
     minimized = True
 
     inputs = [
-        MessageInput(
+        HandleInput(
             name="input_value",
             display_name="Text",
             info="Message to be passed as output.",
+            input_types=["Data", "DataFrame", "Message"],
+            required=True,
         ),
         BoolInput(
             name="should_store_message",
@@ -76,6 +83,13 @@ class ChatOutput(ChatComponent):
             info="The text color of the name",
             advanced=True,
         ),
+        BoolInput(
+            name="clean_data",
+            display_name="Basic Clean Data",
+            value=True,
+            info="Whether to clean the data",
+            advanced=True,
+        ),
     ]
     outputs = [
         Output(
@@ -92,16 +106,35 @@ class ChatOutput(ChatComponent):
         if display_name:
             source_dict["display_name"] = display_name
         if source:
-            source_dict["source"] = source
+            # Handle case where source is a ChatOpenAI object
+            if hasattr(source, "model_name"):
+                source_dict["source"] = source.model_name
+            elif hasattr(source, "model"):
+                source_dict["source"] = str(source.model)
+            else:
+                source_dict["source"] = str(source)
         return Source(**source_dict)
 
     async def message_response(self) -> Message:
+        # First convert the input to string if needed
+        text = self.convert_to_string()
+
+        # Get source properties
         source, icon, display_name, source_id = self.get_properties_from_source_component()
         background_color = self.background_color
         text_color = self.text_color
         if self.chat_icon:
             icon = self.chat_icon
-        message = self.input_value if isinstance(self.input_value, Message) else Message(text=self.input_value)
+
+        # Create or use existing Message object
+        if isinstance(self.input_value, Message):
+            message = self.input_value
+            # Update message properties
+            message.text = text
+        else:
+            message = Message(text=text)
+
+        # Set message properties
         message.sender = self.sender
         message.sender_name = self.sender_name
         message.session_id = self.session_id
@@ -110,12 +143,54 @@ class ChatOutput(ChatComponent):
         message.properties.icon = icon
         message.properties.background_color = background_color
         message.properties.text_color = text_color
-        if self.session_id and isinstance(message, Message) and self.should_store_message:
-            stored_message = await self.send_message(
-                message,
-            )
+
+        # Store message if needed
+        if self.session_id and self.should_store_message:
+            stored_message = await self.send_message(message)
             self.message.value = stored_message
             message = stored_message
 
         self.status = message
         return message
+
+    def _validate_input(self) -> None:
+        """Validate the input data and raise ValueError if invalid."""
+        if self.input_value is None:
+            msg = "Input data cannot be None"
+            raise ValueError(msg)
+        if not isinstance(self.input_value, Data | DataFrame | Message | str | list):
+            msg = f"Expected Data or DataFrame or Message or str, got {type(self.input_value).__name__}"
+            raise TypeError(msg)
+
+    def _safe_convert(self, data: Any) -> str:
+        """Safely convert input data to string."""
+        try:
+            if isinstance(data, str):
+                return data
+            if isinstance(data, Message):
+                return data.get_text()
+            if isinstance(data, Data):
+                if data.get_text() is None:
+                    msg = "Empty Data object"
+                    raise ValueError(msg)
+                return data.get_text()
+            if isinstance(data, DataFrame):
+                if self.clean_data:
+                    # Remove empty rows
+                    data = data.dropna(how="all")
+                    # Remove empty lines in each cell
+                    data = data.replace(r"^\s*$", "", regex=True)
+                    # Replace multiple newlines with a single newline
+                    data = data.replace(r"\n+", "\n", regex=True)
+                return data.to_markdown(index=False)
+            return str(data)
+        except (ValueError, TypeError, AttributeError) as e:
+            msg = f"Error converting data: {e!s}"
+            raise ValueError(msg) from e
+
+    def convert_to_string(self) -> str:
+        """Convert input data to string with proper error handling."""
+        self._validate_input()
+        if isinstance(self.input_value, list):
+            return "\n".join([self._safe_convert(item) for item in self.input_value])
+        return self._safe_convert(self.input_value)

src/backend/tests/unit/components/outputs/test_chat_output_component.py

@@ -0,0 +1,98 @@
+import pytest
+from langflow.components.outputs import ChatOutput
+from langflow.schema.data import Data
+from langflow.schema.dataframe import DataFrame
+from langflow.schema.message import Message
+from langflow.utils.constants import MESSAGE_SENDER_AI, MESSAGE_SENDER_NAME_AI
+
+from tests.base import ComponentTestBaseWithClient
+
+
+@pytest.mark.usefixtures("client")
+class TestChatOutput(ComponentTestBaseWithClient):
+    @pytest.fixture
+    def component_class(self):
+        return ChatOutput
+
+    @pytest.fixture
+    def default_kwargs(self):
+        return {
+            "input_value": "Hello, how are you?",
+            "should_store_message": True,
+            "sender": MESSAGE_SENDER_AI,
+            "sender_name": MESSAGE_SENDER_NAME_AI,
+            "session_id": "test_session_123",
+            "data_template": "{text}",
+            "background_color": "#f0f0f0",
+            "chat_icon": "🤖",
+            "text_color": "#000000",
+            "clean_data": True,
+        }
+
+    @pytest.fixture
+    def file_names_mapping(self):
+        return [
+            {"version": "1.0.19", "module": "outputs", "file_name": "ChatOutput"},
+            {"version": "1.1.0", "module": "outputs", "file_name": "chat"},
+            {"version": "1.1.1", "module": "outputs", "file_name": "chat"},
+        ]
+
+    async def test_process_string_input(self, component_class, default_kwargs):
+        """Test processing a simple string input."""
+        component = component_class(**default_kwargs)
+        input_text = "Hello, this is a test message"
+        component.input_value = input_text
+        result = await component.message_response()
+        assert result.text == input_text
+        assert result.sender == MESSAGE_SENDER_AI
+        assert result.sender_name == MESSAGE_SENDER_NAME_AI
+
+    async def test_process_data_input(self, component_class, default_kwargs):
+        """Test processing a Data object input."""
+        component = component_class(**default_kwargs)
+        data = Data(text="Test data message")
+        component.input_value = data
+        result = await component.message_response()
+        assert result.text == "Test data message"
+        assert result.sender == MESSAGE_SENDER_AI
+
+    async def test_process_dataframe_input(self, component_class, default_kwargs):
+        """Test processing a DataFrame input."""
+        component = component_class(**default_kwargs)
+        sample_df = DataFrame(data={"col1": ["A", "B"], "col2": [1, 2]})
+        component.input_value = sample_df
+        result = await component.message_response()
+        assert "col1" in result.text
+        assert "col2" in result.text
+        assert "A" in result.text
+        assert "B" in result.text
+
+    async def test_process_message_input(self, component_class, default_kwargs):
+        """Test processing a Message object input."""
+        component = component_class(**default_kwargs)
+        message = Message(text="Test message content")
+        component.input_value = message
+        result = await component.message_response()
+        assert result.text == "Test message content"
+        assert result.sender == MESSAGE_SENDER_AI
+
+    async def test_process_list_input(self, component_class, default_kwargs):
+        """Test processing a list of inputs."""
+        component = component_class(**default_kwargs)
+        input_list = ["First message", Data(text="Second message"), Message(text="Third message")]
+        component.input_value = input_list
+        result = await component.message_response()
+        assert "First message" in result.text
+        assert "Second message" in result.text
+        assert "Third message" in result.text
+
+    async def test_invalid_input(self, component_class, default_kwargs):
+        """Test handling of invalid input."""
+        component = component_class(**default_kwargs)
+        component.input_value = None
+        with pytest.raises(ValueError, match="Input data cannot be None"):
+            await component.message_response()
+
+        component.input_value = 123  # Invalid type
+        with pytest.raises(TypeError, match="Expected Data or DataFrame or Message or str"):
+            await component.message_response()

src/backend/tests/unit/components/outputs/test_output_components.py

@@ -1,37 +1,7 @@
 import pytest
-from langflow.components.outputs import ChatOutput, TextOutputComponent
-from langflow.utils.constants import MESSAGE_SENDER_AI, MESSAGE_SENDER_NAME_AI
+from langflow.components.outputs import TextOutputComponent
 
-from tests.base import ComponentTestBaseWithClient, ComponentTestBaseWithoutClient
-
-
-@pytest.mark.usefixtures("client")
-class TestChatOutput(ComponentTestBaseWithClient):
-    @pytest.fixture
-    def component_class(self):
-        return ChatOutput
-
-    @pytest.fixture
-    def default_kwargs(self):
-        return {
-            "input_value": "Hello, how are you?",
-            "should_store_message": True,
-            "sender": MESSAGE_SENDER_AI,
-            "sender_name": MESSAGE_SENDER_NAME_AI,
-            "session_id": "test_session_123",
-            "data_template": "{text}",
-            "background_color": "#f0f0f0",
-            "chat_icon": "🤖",
-            "text_color": "#000000",
-        }
-
-    @pytest.fixture
-    def file_names_mapping(self):
-        return [
-            {"version": "1.0.19", "module": "outputs", "file_name": "ChatOutput"},
-            {"version": "1.1.0", "module": "outputs", "file_name": "chat"},
-            {"version": "1.1.1", "module": "outputs", "file_name": "chat"},
-        ]
+from tests.base import ComponentTestBaseWithoutClient
 
 
 class TestTextOutputComponent(ComponentTestBaseWithoutClient):