feat: Add Message support to Smart Transform component (#11160)

* feat: Add Message support to Smart Transform component

Extends the Smart Transform component to handle Message inputs/outputs in addition to Data and DataFrame. Users can now apply LLM-generated transformations to text messages.

Changes:
- Add Message to accepted input types
- Add new process_as_message output method
- Implement text-specific prompt for Message transformations
- Add error handling with informative messages for all output methods
- Update examples to include text transformation use cases

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes

* improve code wuality

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes

---------

Co-authored-by: Rodrigo Nader <rodrigonader@MacBook-Pro-de-Rodrigo.local>
Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Edwin Jose <edwin.jose@datastax.com>
Co-authored-by: cristhianzl <cristhian.lousa@gmail.com>

src/backend/tests/unit/components/llm_operations/test_lambda_filter.py

@@ -3,6 +3,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import pytest
 from lfx.components.llm_operations.lambda_filter import LambdaFilterComponent
 from lfx.schema import Data
+from lfx.schema.dataframe import DataFrame
+from lfx.schema.message import Message
 
 from tests.base import ComponentTestBaseWithoutClient
 
@@ -50,83 +52,668 @@ class TestLambdaFilterComponent(ComponentTestBaseWithoutClient):
     def file_names_mapping(self):
         return []
 
-    @patch("lfx.base.models.unified_models.get_model_classes")
-    async def test_invalid_lambda_response(self, mock_get_model_classes, component_class, default_kwargs, mock_llm):
-        # Mock get_model_classes to return MockLanguageModel factory
-        mock_model_class = MagicMock(return_value=mock_llm)
-        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
 
-        component = await self.component_setup(component_class, default_kwargs)
-        mock_llm.ainvoke.return_value.content = "invalid lambda syntax"
+class TestValidateLambda(TestLambdaFilterComponent):
+    """Tests for _validate_lambda method."""
 
-        # Test exception handling
-        with pytest.raises(ValueError, match="Could not find lambda in response"):
-            await component.process_as_data()
+    def test_should_return_true_when_lambda_is_valid(self, component_class):
+        # Arrange
+        component = component_class()
+        valid_lambda = "lambda x: x + 1"
+
+        # Act
+        result = component._validate_lambda(valid_lambda)
+
+        # Assert
+        assert result is True
+
+    def test_should_return_false_when_lambda_keyword_missing(self, component_class):
+        # Arrange
+        component = component_class()
+        invalid_lambda = "x: x + 1"
+
+        # Act
+        result = component._validate_lambda(invalid_lambda)
+
+        # Assert
+        assert result is False
+
+    def test_should_return_false_when_colon_missing(self, component_class):
+        # Arrange
+        component = component_class()
+        invalid_lambda = "lambda x x + 1"
+
+        # Act
+        result = component._validate_lambda(invalid_lambda)
+
+        # Assert
+        assert result is False
+
+    def test_should_return_true_when_lambda_has_whitespace(self, component_class):
+        # Arrange
+        component = component_class()
+        valid_lambda = "  lambda x: x + 1  "
+
+        # Act
+        result = component._validate_lambda(valid_lambda)
+
+        # Assert
+        assert result is True
+
+
+class TestGetDataStructure(TestLambdaFilterComponent):
+    """Tests for get_data_structure method."""
+
+    def test_should_return_type_name_when_input_is_primitive(self, component_class):
+        # Arrange
+        component = component_class()
+        bool_value = True
+
+        # Act & Assert
+        assert component.get_data_structure("test") == "str"
+        assert component.get_data_structure(42) == "int"
+        assert component.get_data_structure(3.14) == "float"
+        assert component.get_data_structure(bool_value) == "bool"
+
+    def test_should_return_dict_structure_when_input_is_dict(self, component_class):
+        # Arrange
+        component = component_class()
+        test_data = {"key": "value", "number": 42}
+
+        # Act
+        result = component.get_data_structure(test_data)
+
+        # Assert
+        assert result == {"key": "str", "number": "int"}
+
+    def test_should_return_list_structure_when_input_is_list(self, component_class):
+        # Arrange
+        component = component_class()
+        test_data = [1, 2, 3]
+
+        # Act
+        result = component.get_data_structure(test_data)
+
+        # Assert
+        assert result == ["int"]
+
+    def test_should_return_empty_list_when_input_is_empty_list(self, component_class):
+        # Arrange
+        component = component_class()
+
+        # Act
+        result = component.get_data_structure([])
+
+        # Assert
+        assert result == []
+
+    def test_should_return_nested_structure_when_input_is_nested(self, component_class):
+        # Arrange
+        component = component_class()
+        test_data = {"nested": {"a": [{"b": 1}]}}
+
+        # Act
+        result = component.get_data_structure(test_data)
+
+        # Assert
+        assert result == {"nested": {"a": [{"b": "int"}]}}
+
+
+class TestGetInputTypeName(TestLambdaFilterComponent):
+    """Tests for _get_input_type_name method."""
+
+    def test_should_return_message_when_input_is_single_message(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = Message(text="test")
+
+        # Act
+        result = component._get_input_type_name()
+
+        # Assert
+        assert result == "Message"
+
+    def test_should_return_message_when_input_is_list_of_messages(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = [Message(text="test1"), Message(text="test2")]
+
+        # Act
+        result = component._get_input_type_name()
+
+        # Assert
+        assert result == "Message"
+
+    def test_should_return_dataframe_when_input_is_dataframe(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = DataFrame([{"a": 1}])
+
+        # Act
+        result = component._get_input_type_name()
+
+        # Assert
+        assert result == "DataFrame"
+
+    def test_should_return_data_when_input_is_data(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = Data(data={"key": "value"})
+
+        # Act
+        result = component._get_input_type_name()
+
+        # Assert
+        assert result == "Data"
+
+    def test_should_return_unknown_when_input_is_empty_list(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = []
+
+        # Act
+        result = component._get_input_type_name()
+
+        # Assert
+        assert result == "unknown"
+
+
+class TestIsMessageInput(TestLambdaFilterComponent):
+    """Tests for _is_message_input method."""
+
+    def test_should_return_true_when_input_is_single_message(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = Message(text="test")
+
+        # Act
+        result = component._is_message_input()
+
+        # Assert
+        assert result is True
+
+    def test_should_return_true_when_input_is_list_of_messages(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = [Message(text="test1"), Message(text="test2")]
+
+        # Act
+        result = component._is_message_input()
+
+        # Assert
+        assert result is True
+
+    def test_should_return_false_when_input_is_data(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = Data(data={"key": "value"})
+
+        # Act
+        result = component._is_message_input()
+
+        # Assert
+        assert result is False
+
+    def test_should_return_false_when_input_is_empty_list(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = []
+
+        # Act
+        result = component._is_message_input()
+
+        # Assert
+        assert result is False
+
+
+class TestExtractMessageText(TestLambdaFilterComponent):
+    """Tests for _extract_message_text method."""
+
+    def test_should_return_text_when_input_is_single_message(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = Message(text="Hello World")
+
+        # Act
+        result = component._extract_message_text()
+
+        # Assert
+        assert result == "Hello World"
+
+    def test_should_return_empty_string_when_message_text_is_none(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = Message(text=None)
+
+        # Act
+        result = component._extract_message_text()
+
+        # Assert
+        assert result == ""
+
+    def test_should_join_texts_when_input_is_list_of_messages(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = [Message(text="Hello"), Message(text="World")]
+
+        # Act
+        result = component._extract_message_text()
+
+        # Assert
+        assert result == "Hello\n\nWorld"
+
+    def test_should_return_single_text_when_list_has_one_message(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = [Message(text="Only one")]
+
+        # Act
+        result = component._extract_message_text()
+
+        # Assert
+        assert result == "Only one"
+
+
+class TestExtractStructuredData(TestLambdaFilterComponent):
+    """Tests for _extract_structured_data method."""
+
+    def test_should_return_dict_when_input_is_single_data(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = Data(data={"key": "value"})
+
+        # Act
+        result = component._extract_structured_data()
+
+        # Assert
+        assert result == {"key": "value"}
+
+    def test_should_return_records_when_input_is_dataframe(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = DataFrame([{"a": 1}, {"a": 2}])
+
+        # Act
+        result = component._extract_structured_data()
+
+        # Assert
+        assert result == [{"a": 1}, {"a": 2}]
+
+    def test_should_combine_data_when_input_is_list_of_data(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = [Data(data={"a": 1}), Data(data={"b": 2})]
+
+        # Act
+        result = component._extract_structured_data()
+
+        # Assert
+        assert result == [{"a": 1}, {"b": 2}]
+
+    def test_should_unwrap_single_dict_when_list_has_one_item(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = [Data(data={"only": "one"})]
+
+        # Act
+        result = component._extract_structured_data()
+
+        # Assert
+        assert result == {"only": "one"}
+
+    def test_should_return_empty_dict_when_no_data_extracted(self, component_class):
+        # Arrange
+        component = component_class()
+        component.data = []
+
+        # Act
+        result = component._extract_structured_data()
+
+        # Assert
+        assert result == {}
+
+
+class TestBuildTextPrompt(TestLambdaFilterComponent):
+    """Tests for _build_text_prompt method."""
+
+    def test_should_include_full_text_when_text_is_small(self, component_class):
+        # Arrange
+        component = component_class()
+        component.max_size = 1000
+        component.sample_size = 100
+        component.filter_instruction = "Transform to uppercase"
+        text = "Short text"
+
+        # Act
+        result = component._build_text_prompt(text)
+
+        # Assert
+        assert "Short text" in result
+        assert "Transform to uppercase" in result
+
+    def test_should_truncate_text_when_text_is_large(self, component_class):
+        # Arrange
+        component = component_class()
+        component.max_size = 50
+        component.sample_size = 10
+        component.filter_instruction = "Summarize"
+        text = "A" * 100
+
+        # Act
+        result = component._build_text_prompt(text)
+
+        # Assert
+        assert "Text length: 100 characters" in result
+        assert "First 10 characters" in result
+        assert "Last 10 characters" in result
+
+
+class TestBuildDataPrompt(TestLambdaFilterComponent):
+    """Tests for _build_data_prompt method."""
+
+    def test_should_include_full_data_when_data_is_small(self, component_class):
+        # Arrange
+        component = component_class()
+        component.max_size = 1000
+        component.sample_size = 100
+        component.filter_instruction = "Filter by value"
+        data = {"key": "value"}
+
+        # Act
+        result = component._build_data_prompt(data)
+
+        # Assert
+        assert '"key": "value"' in result
+        assert "Filter by value" in result
+
+    def test_should_truncate_data_when_data_is_large(self, component_class):
+        # Arrange
+        component = component_class()
+        component.max_size = 50
+        component.sample_size = 10
+        component.filter_instruction = "Filter"
+        data = {"key": "A" * 100}
+
+        # Act
+        result = component._build_data_prompt(data)
+
+        # Assert
+        assert "Data is too long to display" in result
+        assert "First lines (head)" in result
+        assert "Last lines (tail)" in result
+
+
+class TestConvertResultToData(TestLambdaFilterComponent):
+    """Tests for _convert_result_to_data method."""
+
+    def test_should_wrap_dict_when_result_is_dict(self, component_class):
+        # Arrange
+        component = component_class()
+        result = {"key": "value"}
+
+        # Act
+        data = component._convert_result_to_data(result)
+
+        # Assert
+        assert isinstance(data, Data)
+        assert data.data == {"key": "value"}
+
+    def test_should_wrap_list_in_results_key_when_result_is_list(self, component_class):
+        # Arrange
+        component = component_class()
+        result = [1, 2, 3]
+
+        # Act
+        data = component._convert_result_to_data(result)
+
+        # Assert
+        assert isinstance(data, Data)
+        assert data.data == {"_results": [1, 2, 3]}
+
+    def test_should_convert_to_string_when_result_is_other_type(self, component_class):
+        # Arrange
+        component = component_class()
+        result = 42
+
+        # Act
+        data = component._convert_result_to_data(result)
+
+        # Assert
+        assert isinstance(data, Data)
+        assert data.data == {"text": "42"}
+
+
+class TestConvertResultToDataframe(TestLambdaFilterComponent):
+    """Tests for _convert_result_to_dataframe method."""
+
+    def test_should_create_dataframe_when_result_is_list_of_dicts(self, component_class):
+        # Arrange
+        component = component_class()
+        result = [{"a": 1}, {"a": 2}]
+
+        # Act
+        df = component._convert_result_to_dataframe(result)
+
+        # Assert
+        assert isinstance(df, DataFrame)
+
+    def test_should_wrap_values_when_result_is_list_of_non_dicts(self, component_class):
+        # Arrange
+        component = component_class()
+        result = [1, 2, 3]
+
+        # Act
+        df = component._convert_result_to_dataframe(result)
+
+        # Assert
+        assert isinstance(df, DataFrame)
+
+    def test_should_create_single_row_when_result_is_dict(self, component_class):
+        # Arrange
+        component = component_class()
+        result = {"a": 1}
+
+        # Act
+        df = component._convert_result_to_dataframe(result)
+
+        # Assert
+        assert isinstance(df, DataFrame)
+
+
+class TestConvertResultToMessage(TestLambdaFilterComponent):
+    """Tests for _convert_result_to_message method."""
+
+    def test_should_return_message_when_result_is_string(self, component_class):
+        # Arrange
+        component = component_class()
+        result = "Hello World"
+
+        # Act
+        msg = component._convert_result_to_message(result)
+
+        # Assert
+        assert isinstance(msg, Message)
+        assert msg.text == "Hello World"
+
+    def test_should_join_items_when_result_is_list(self, component_class):
+        # Arrange
+        component = component_class()
+        result = ["Line 1", "Line 2"]
+
+        # Act
+        msg = component._convert_result_to_message(result)
+
+        # Assert
+        assert isinstance(msg, Message)
+        assert msg.text == "Line 1\nLine 2"
+
+    def test_should_format_json_when_result_is_dict(self, component_class):
+        # Arrange
+        component = component_class()
+        result = {"key": "value"}
+
+        # Act
+        msg = component._convert_result_to_message(result)
+
+        # Assert
+        assert isinstance(msg, Message)
+        assert '"key": "value"' in msg.text
+
+
+class TestProcessAsDataIntegration(TestLambdaFilterComponent):
+    """Integration tests for process_as_data method."""
 
     @patch("lfx.base.models.unified_models.get_model_classes")
-    async def test_successful_lambda_generation(
+    async def test_should_return_filtered_data_when_lambda_is_valid(
         self, mock_get_model_classes, component_class, default_kwargs, mock_llm
     ):
-        """Test that a lambda function is successfully generated and applied."""
-        # Mock get_model_classes to return MockLanguageModel factory
+        # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
         mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
-
         component = await self.component_setup(component_class, default_kwargs)
         mock_llm.ainvoke.return_value.content = "lambda x: [item for item in x['items'] if item['value'] > 15]"
 
-        # Execute the lambda filter
+        # Act
         result = await component.process_as_data()
 
-        # Assertions - process_as_data() returns a Data object
-        assert isinstance(result, Data), f"Expected Data object, got {type(result)}"
-        assert "_results" in result.data, "Expected '_results' key in Data object"
-
-        # Check the filtered results
+        # Assert
+        assert isinstance(result, Data)
+        assert "_results" in result.data
         filtered_items = result.data["_results"]
-        assert isinstance(filtered_items, list), "Expected list of filtered items"
-        assert len(filtered_items) == 1, f"Expected 1 item, got {len(filtered_items)}"
-        assert filtered_items[0]["name"] == "test2", f"Expected 'test2', got {filtered_items[0]['name']}"
-        assert filtered_items[0]["value"] == 20, f"Expected value 20, got {filtered_items[0]['value']}"
+        assert len(filtered_items) == 1
+        assert filtered_items[0]["name"] == "test2"
+        assert filtered_items[0]["value"] == 20
 
     @patch("lfx.base.models.unified_models.get_model_classes")
-    async def test_lambda_with_large_dataset(self, mock_get_model_classes, component_class, default_kwargs, mock_llm):
-        """Test lambda execution with a large dataset."""
-        # Mock get_model_classes to return MockLanguageModel factory
+    async def test_should_raise_error_when_lambda_not_found_in_response(
+        self, mock_get_model_classes, component_class, default_kwargs, mock_llm
+    ):
+        # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
         mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        component = await self.component_setup(component_class, default_kwargs)
+        mock_llm.ainvoke.return_value.content = "invalid response without lambda"
 
+        # Act & Assert
+        with pytest.raises(ValueError, match="Could not find lambda in response"):
+            await component.process_as_data()
+
+
+class TestProcessAsMessageIntegration(TestLambdaFilterComponent):
+    """Integration tests for process_as_message with Message input."""
+
+    @pytest.fixture
+    def message_kwargs(self, model_metadata):
+        """Return kwargs with Message input."""
+        return {
+            "data": [Message(text="Hello World")],
+            "model": model_metadata,
+            "api_key": "test-api-key",
+            "filter_instruction": "Convert to uppercase",
+            "sample_size": 1000,
+            "max_size": 30000,
+        }
+
+    @patch("lfx.base.models.unified_models.get_model_classes")
+    async def test_should_transform_message_when_input_is_message(
+        self, mock_get_model_classes, component_class, message_kwargs, mock_llm
+    ):
+        # Arrange
+        mock_model_class = MagicMock(return_value=mock_llm)
+        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        component = await self.component_setup(component_class, message_kwargs)
+        mock_llm.ainvoke.return_value.content = "lambda text: text.upper()"
+
+        # Act
+        result = await component.process_as_message()
+
+        # Assert
+        assert isinstance(result, Message)
+        assert result.text == "HELLO WORLD"
+
+    @patch("lfx.base.models.unified_models.get_model_classes")
+    async def test_should_join_multiple_messages_when_input_is_list_of_messages(
+        self, mock_get_model_classes, component_class, model_metadata, mock_llm
+    ):
+        # Arrange
+        mock_model_class = MagicMock(return_value=mock_llm)
+        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        kwargs = {
+            "data": [Message(text="Hello"), Message(text="World")],
+            "model": model_metadata,
+            "api_key": "test-api-key",
+            "filter_instruction": "Convert to uppercase",
+            "sample_size": 1000,
+            "max_size": 30000,
+        }
+        component = await self.component_setup(component_class, kwargs)
+        mock_llm.ainvoke.return_value.content = "lambda text: text.upper()"
+
+        # Act
+        result = await component.process_as_message()
+
+        # Assert
+        assert isinstance(result, Message)
+        assert result.text == "HELLO\n\nWORLD"
+
+
+class TestProcessAsDataframeIntegration(TestLambdaFilterComponent):
+    """Integration tests for process_as_dataframe method."""
+
+    @patch("lfx.base.models.unified_models.get_model_classes")
+    async def test_should_return_dataframe_when_lambda_returns_list_of_dicts(
+        self, mock_get_model_classes, component_class, default_kwargs, mock_llm
+    ):
+        # Arrange
+        mock_model_class = MagicMock(return_value=mock_llm)
+        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        component = await self.component_setup(component_class, default_kwargs)
+        mock_llm.ainvoke.return_value.content = "lambda x: x['items']"
+
+        # Act
+        result = await component.process_as_dataframe()
+
+        # Assert
+        assert isinstance(result, DataFrame)
+
+
+class TestLargeDataset(TestLambdaFilterComponent):
+    """Tests for handling large datasets."""
+
+    @patch("lfx.base.models.unified_models.get_model_classes")
+    async def test_should_filter_large_dataset_when_data_exceeds_max_size(
+        self, mock_get_model_classes, component_class, default_kwargs, mock_llm
+    ):
+        # Arrange
+        mock_model_class = MagicMock(return_value=mock_llm)
+        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
         large_data = {"items": [{"name": f"test{i}", "value": i} for i in range(2000)]}
         default_kwargs["data"] = [Data(data=large_data)]
         default_kwargs["filter_instruction"] = "Filter items with value greater than 1500"
         component = await self.component_setup(component_class, default_kwargs)
         mock_llm.ainvoke.return_value.content = "lambda x: [item for item in x['items'] if item['value'] > 1500]"
 
-        # Execute filter on the data
+        # Act
         result = await component.process_as_data()
 
-        # Assertions - process_as_data() returns a Data object
-        assert isinstance(result, Data), f"Expected Data object, got {type(result)}"
-        assert "_results" in result.data, "Expected '_results' key in Data object"
-
-        # Check the filtered results from the lambda
+        # Assert
+        assert isinstance(result, Data)
         filtered_items = result.data["_results"]
-        assert isinstance(filtered_items, list), "Expected list of filtered items"
-        assert len(filtered_items) == 499, f"Expected 499 items (1501-1999), got {len(filtered_items)}"
+        assert len(filtered_items) == 499
+        assert filtered_items[0]["value"] == 1501
+        assert filtered_items[-1]["value"] == 1999
 
-        # Verify first and last items
-        assert filtered_items[0]["value"] == 1501, f"Expected first value 1501, got {filtered_items[0]['value']}"
-        assert filtered_items[-1]["value"] == 1999, f"Expected last value 1999, got {filtered_items[-1]['value']}"
+
+class TestComplexDataStructure(TestLambdaFilterComponent):
+    """Tests for handling complex nested data structures."""
 
     @patch("lfx.base.models.unified_models.get_model_classes")
-    async def test_lambda_with_complex_data_structure(
+    async def test_should_handle_nested_data_when_structure_is_complex(
         self, mock_get_model_classes, component_class, default_kwargs, mock_llm
     ):
-        """Test lambda execution with complex nested data structures."""
-        # Mock get_model_classes to return MockLanguageModel factory
+        # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
         mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
-
         complex_data = {
             "categories": {
                 "A": [{"id": 1, "score": 90}, {"id": 2, "score": 85}],
@@ -140,61 +727,12 @@ class TestLambdaFilterComponent(ComponentTestBaseWithoutClient):
             "lambda x: [item for cat in x['categories'].values() for item in cat if item['score'] > 90]"
         )
 
-        # Execute filter
+        # Act
         result = await component.process_as_data()
 
-        # Assertions - process_as_data() returns a Data object
-        assert isinstance(result, Data), f"Expected Data object, got {type(result)}"
-        assert "_results" in result.data, "Expected '_results' key in Data object"
-
-        # Check the filtered results
+        # Assert
+        assert isinstance(result, Data)
         filtered_items = result.data["_results"]
-        assert isinstance(filtered_items, list), "Expected list of filtered items"
-        assert len(filtered_items) == 1, f"Expected 1 item with score > 90, got {len(filtered_items)}"
-        assert filtered_items[0]["id"] == 3, f"Expected id 3, got {filtered_items[0]['id']}"
-        assert filtered_items[0]["score"] == 95, f"Expected score 95, got {filtered_items[0]['score']}"
-
-    def test_validate_lambda(self, component_class):
-        component = component_class()
-
-        # Valid lambda
-        valid_lambda = "lambda x: x + 1"
-        assert component._validate_lambda(valid_lambda) is True
-
-        # Invalid lambda: missing 'lambda'
-        invalid_lambda_1 = "x: x + 1"
-        assert component._validate_lambda(invalid_lambda_1) is False
-
-        # Invalid lambda: missing ':'
-        invalid_lambda_2 = "lambda x x + 1"
-        assert component._validate_lambda(invalid_lambda_2) is False
-
-    def test_get_data_structure(self, component_class):
-        """Test that get_data_structure returns a mirror of the data with types."""
-        component = component_class()
-        test_data = {
-            "string": "test",
-            "number": 42,
-            "list": [1, 2, 3],
-            "dict": {"key": "value"},
-            "nested": {"a": [{"b": 1}]},
-        }
-
-        structure = component.get_data_structure(test_data)
-
-        # Verify the structure returns type names for primitive types
-        assert structure["string"] == "str", f"Expected 'str', got {structure['string']}"
-        assert structure["number"] == "int", f"Expected 'int', got {structure['number']}"
-
-        # Verify list structure
-        assert isinstance(structure["list"], list), "List should return a list structure"
-        assert structure["list"] == ["int"], f"Expected ['int'], got {structure['list']}"
-
-        # Verify dict structure
-        assert isinstance(structure["dict"], dict), "Dict should return a dict structure"
-        assert structure["dict"] == {"key": "str"}, f"Expected {{'key': 'str'}}, got {structure['dict']}"
-
-        # Verify nested structure
-        assert structure["nested"] == {"a": [{"b": "int"}]}, (
-            f"Expected nested structure {{'a': [{{'b': 'int'}}]}}, got {structure['nested']}"
-        )
+        assert len(filtered_items) == 1
+        assert filtered_items[0]["id"] == 3
+        assert filtered_items[0]["score"] == 95

src/lfx/src/lfx/components/llm_operations/lambda_filter.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 import json
 import re
-from typing import TYPE_CHECKING, Any
+from collections.abc import Callable  # noqa: TC003 - required at runtime for dynamic exec()
+from typing import Any
 
 from lfx.base.models.unified_models import (
     get_language_model_options,
@@ -13,18 +14,34 @@ from lfx.custom.custom_component.component import Component
 from lfx.io import DataInput, IntInput, ModelInput, MultilineInput, Output, SecretStrInput
 from lfx.schema.data import Data
 from lfx.schema.dataframe import DataFrame
+from lfx.schema.message import Message
+from lfx.utils.constants import MESSAGE_SENDER_AI
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
+TEXT_TRANSFORM_PROMPT = (
+    "Given this text, create a Python lambda function that transforms it "
+    "according to the instruction.\n"
+    "The lambda should take a string parameter and return the transformed string.\n\n"
+    "Text Preview:\n{text_preview}\n\n"
+    "Instruction: {instruction}\n\n"
+    "Return ONLY the lambda function and nothing else. No need for ```python or whatever.\n"
+    "Just a string starting with lambda.\n"
+    "Example: lambda text: text.upper()"
+)
 
-# # Compute model options once at module level
-# _MODEL_OPTIONS = get_language_model_options()
-# _PROVIDERS = [provider["provider"] for provider in _MODEL_OPTIONS]
+DATA_TRANSFORM_PROMPT = (
+    "Given this data structure and examples, create a Python lambda function "
+    "that implements the following instruction:\n\n"
+    "Data Structure:\n{dump_structure}\n\n"
+    "Example Items:\n{data_sample}\n\n"
+    "Instruction: {instruction}\n\n"
+    "Return ONLY the lambda function and nothing else. No need for ```python or whatever.\n"
+    "Just a string starting with lambda."
+)
 
 
 class LambdaFilterComponent(Component):
     display_name = "Smart Transform"
-    description = "Uses an LLM to generate a function for filtering or transforming structured data."
+    description = "Uses an LLM to generate a function for filtering or transforming structured data and messages."
     documentation: str = "https://docs.langflow.org/smart-transform"
     icon = "square-function"
     name = "Smart Transform"
@@ -33,8 +50,8 @@ class LambdaFilterComponent(Component):
         DataInput(
             name="data",
             display_name="Data",
-            info="The structured data to filter or transform using a lambda function.",
-            input_types=["Data", "DataFrame"],
+            info="The structured data or text messages to filter or transform using a lambda function.",
+            input_types=["Data", "DataFrame", "Message"],
             is_list=True,
             required=True,
         ),
@@ -57,9 +74,10 @@ class LambdaFilterComponent(Component):
             display_name="Instructions",
             info=(
                 "Natural language instructions for how to filter or transform the data using a lambda function. "
-                "Example: Filter the data to only include items where the 'status' is 'active'."
+                "Examples: 'Filter the data to only include items where status is active', "
+                "'Convert the text to uppercase', 'Keep only first 100 characters'"
             ),
-            value="Filter the data to...",
+            value="Transform the data to...",
             required=True,
         ),
         IntInput(
@@ -89,6 +107,11 @@ class LambdaFilterComponent(Component):
             name="dataframe_output",
             method="process_as_dataframe",
         ),
+        Output(
+            display_name="Output",
+            name="message_output",
+            method="process_as_message",
+        ),
     ]
 
     def update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None):
@@ -119,127 +142,189 @@ class LambdaFilterComponent(Component):
         # Return False if the lambda function does not start with 'lambda' or does not contain a colon
         return lambda_text.strip().startswith("lambda") and ":" in lambda_text
 
-    async def _execute_lambda(self) -> Any:
-        self.log(str(self.data))
+    def _get_input_type_name(self) -> str:
+        """Detect and return the input type name for error messages."""
+        if isinstance(self.data, Message):
+            return "Message"
+        if isinstance(self.data, DataFrame):
+            return "DataFrame"
+        if isinstance(self.data, Data):
+            return "Data"
+        if isinstance(self.data, list) and len(self.data) > 0:
+            first = self.data[0]
+            if isinstance(first, Message):
+                return "Message"
+            if isinstance(first, DataFrame):
+                return "DataFrame"
+            if isinstance(first, Data):
+                return "Data"
+        return "unknown"
 
-        # Convert input to a unified format
-        if isinstance(self.data, list):
-            # Handle list of Data or DataFrame objects
-            combined_data = []
-            for item in self.data:
-                if isinstance(item, DataFrame):
-                    # DataFrame to list of dicts
-                    combined_data.extend(item.to_dict(orient="records"))
-                elif hasattr(item, "data"):
-                    # Data object
-                    if isinstance(item.data, dict):
-                        combined_data.append(item.data)
-                    elif isinstance(item.data, list):
-                        combined_data.extend(item.data)
+    def _extract_message_text(self) -> str:
+        """Extract text content from Message input(s)."""
+        if isinstance(self.data, Message):
+            return self.data.text or ""
 
-            # If we have a single dict, unwrap it so lambdas can access it directly
-            if len(combined_data) == 1 and isinstance(combined_data[0], dict):
-                data = combined_data[0]
-            elif len(combined_data) == 0:
-                data = {}
-            else:
-                data = combined_data  # type: ignore[assignment]
-        elif isinstance(self.data, DataFrame):
-            # Single DataFrame to list of dicts
-            data = self.data.to_dict(orient="records")
-        elif hasattr(self.data, "data"):
-            # Single Data object
-            data = self.data.data
+        texts = [msg.text or "" for msg in self.data if isinstance(msg, Message)]
+        return "\n\n".join(texts) if len(texts) > 1 else (texts[0] if texts else "")
+
+    def _extract_structured_data(self) -> dict | list:
+        """Extract structured data from Data or DataFrame input(s)."""
+        if isinstance(self.data, DataFrame):
+            return self.data.to_dict(orient="records")
+
+        if hasattr(self.data, "data"):
+            return self.data.data
+
+        if not isinstance(self.data, list):
+            return self.data
+
+        combined_data: list[dict] = []
+        for item in self.data:
+            if isinstance(item, DataFrame):
+                combined_data.extend(item.to_dict(orient="records"))
+            elif hasattr(item, "data"):
+                if isinstance(item.data, dict):
+                    combined_data.append(item.data)
+                elif isinstance(item.data, list):
+                    combined_data.extend(item.data)
+
+        if len(combined_data) == 1 and isinstance(combined_data[0], dict):
+            return combined_data[0]
+        if len(combined_data) == 0:
+            return {}
+        return combined_data
+
+    def _is_message_input(self) -> bool:
+        """Check if input is Message type."""
+        if isinstance(self.data, Message):
+            return True
+        return isinstance(self.data, list) and len(self.data) > 0 and isinstance(self.data[0], Message)
+
+    def _build_text_prompt(self, text: str) -> str:
+        """Build prompt for text/Message transformation."""
+        text_length = len(text)
+        if text_length > self.max_size:
+            text_preview = (
+                f"Text length: {text_length} characters\n\n"
+                f"First {self.sample_size} characters:\n{text[: self.sample_size]}\n\n"
+                f"Last {self.sample_size} characters:\n{text[-self.sample_size :]}"
+            )
         else:
-            data = self.data
+            text_preview = text
 
+        return TEXT_TRANSFORM_PROMPT.format(text_preview=text_preview, instruction=self.filter_instruction)
+
+    def _build_data_prompt(self, data: dict | list) -> str:
+        """Build prompt for structured data transformation."""
         dump = json.dumps(data)
-        self.log(str(data))
+        dump_structure = json.dumps(self.get_data_structure(data))
 
-        llm = get_llm(model=self.model, user_id=self.user_id, api_key=self.api_key)
-        instruction = self.filter_instruction
-        sample_size = self.sample_size
-
-        # Get data structure and samples
-        data_structure = self.get_data_structure(data)
-        dump_structure = json.dumps(data_structure)
-        self.log(dump_structure)
-
-        # For large datasets, sample from head and tail
         if len(dump) > self.max_size:
             data_sample = (
-                f"Data is too long to display... \n\n First lines (head): {dump[:sample_size]} \n\n"
-                f" Last lines (tail): {dump[-sample_size:]})"
+                f"Data is too long to display...\n\nFirst lines (head): {dump[: self.sample_size]}\n\n"
+                f"Last lines (tail): {dump[-self.sample_size :]}"
             )
         else:
             data_sample = dump
 
-        self.log(data_sample)
+        return DATA_TRANSFORM_PROMPT.format(
+            dump_structure=dump_structure, data_sample=data_sample, instruction=self.filter_instruction
+        )
 
-        prompt = f"""Given this data structure and examples, create a Python lambda function that
-                    implements the following instruction:
-
-                    Data Structure:
-                    {dump_structure}
-
-                    Example Items:
-                    {data_sample}
-
-                    Instruction: {instruction}
-
-                    Return ONLY the lambda function and nothing else. No need for ```python or whatever.
-                    Just a string starting with lambda.
-                    """
-
-        response = await llm.ainvoke(prompt)
-        response_text = response.content if hasattr(response, "content") else str(response)
-        self.log(response_text)
-
-        # Extract lambda using regex
+    def _parse_lambda_from_response(self, response_text: str) -> Callable[[Any], Any]:
+        """Extract and validate lambda function from LLM response."""
         lambda_match = re.search(r"lambda\s+\w+\s*:.*?(?=\n|$)", response_text)
         if not lambda_match:
             msg = f"Could not find lambda in response: {response_text}"
             raise ValueError(msg)
 
         lambda_text = lambda_match.group().strip()
-        self.log(lambda_text)
+        self.log(f"Generated lambda: {lambda_text}")
 
-        # Validation is commented out as requested
         if not self._validate_lambda(lambda_text):
             msg = f"Invalid lambda format: {lambda_text}"
             raise ValueError(msg)
 
-        # Create and apply the function
-        fn: Callable[[Any], Any] = eval(lambda_text)  # noqa: S307
+        return eval(lambda_text)  # noqa: S307
 
-        # Apply the lambda function to the data
+    async def _execute_lambda(self) -> Any:
+        """Generate and execute a lambda function based on input type."""
+        if self._is_message_input():
+            data: Any = self._extract_message_text()
+            prompt = self._build_text_prompt(data)
+        else:
+            data = self._extract_structured_data()
+            prompt = self._build_data_prompt(data)
+
+        llm = get_llm(model=self.model, user_id=self.user_id, api_key=self.api_key)
+        response = await llm.ainvoke(prompt)
+        response_text = response.content if hasattr(response, "content") else str(response)
+
+        fn = self._parse_lambda_from_response(response_text)
         return fn(data)
 
-    async def process_as_data(self) -> Data:
-        """Process the data and return as a Data object."""
-        result = await self._execute_lambda()
+    def _handle_process_error(self, error: Exception, output_type: str) -> None:
+        """Handle errors from process methods with context-aware messages."""
+        input_type = self._get_input_type_name()
+        error_msg = (
+            f"Failed to convert result to {output_type} output. "
+            f"Error: {error}. "
+            f"Input type was {input_type}. "
+            f"Try using the same output type as the input."
+        )
+        raise ValueError(error_msg) from error
 
-        # Convert result to Data based on type
+    def _convert_result_to_data(self, result: Any) -> Data:
+        """Convert lambda result to Data object."""
         if isinstance(result, dict):
             return Data(data=result)
         if isinstance(result, list):
             return Data(data={"_results": result})
-        # For other types, convert to string
         return Data(data={"text": str(result)})
 
+    def _convert_result_to_dataframe(self, result: Any) -> DataFrame:
+        """Convert lambda result to DataFrame object."""
+        if isinstance(result, list):
+            if all(isinstance(item, dict) for item in result):
+                return DataFrame(result)
+            return DataFrame([{"value": item} for item in result])
+        if isinstance(result, dict):
+            return DataFrame([result])
+        return DataFrame([{"value": str(result)}])
+
+    def _convert_result_to_message(self, result: Any) -> Message:
+        """Convert lambda result to Message object."""
+        if isinstance(result, str):
+            return Message(text=result, sender=MESSAGE_SENDER_AI)
+        if isinstance(result, list):
+            text = "\n".join(str(item) for item in result)
+            return Message(text=text, sender=MESSAGE_SENDER_AI)
+        if isinstance(result, dict):
+            text = json.dumps(result, indent=2)
+            return Message(text=text, sender=MESSAGE_SENDER_AI)
+        return Message(text=str(result), sender=MESSAGE_SENDER_AI)
+
+    async def process_as_data(self) -> Data:
+        """Process the data and return as a Data object."""
+        try:
+            result = await self._execute_lambda()
+            return self._convert_result_to_data(result)
+        except Exception as e:  # noqa: BLE001 - dynamic lambda can raise any exception
+            self._handle_process_error(e, "Data")
+
     async def process_as_dataframe(self) -> DataFrame:
         """Process the data and return as a DataFrame."""
-        result = await self._execute_lambda()
+        try:
+            result = await self._execute_lambda()
+            return self._convert_result_to_dataframe(result)
+        except Exception as e:  # noqa: BLE001 - dynamic lambda can raise any exception
+            self._handle_process_error(e, "DataFrame")
 
-        # Convert result to DataFrame based on type
-        if isinstance(result, list):
-            # Check if it's a list of dicts
-            if all(isinstance(item, dict) for item in result):
-                return DataFrame(result)
-            # List of non-dicts: wrap each value
-            return DataFrame([{"value": item} for item in result])
-        if isinstance(result, dict):
-            # Single dict becomes single-row DataFrame
-            return DataFrame([result])
-        # Other types: convert to string and wrap
-        return DataFrame([{"value": str(result)}])
+    async def process_as_message(self) -> Message:
+        """Process the data and return as a Message."""
+        try:
+            result = await self._execute_lambda()
+            return self._convert_result_to_message(result)
+        except Exception as e:  # noqa: BLE001 - dynamic lambda can raise any exception
+            self._handle_process_error(e, "Message")