Add serialization support for ChatMessageTrigger

ee/vellum_ee/workflows/display/base.py

@@ -18,6 +18,7 @@ class WorkflowTriggerType(Enum):
     MANUAL = "MANUAL"
     INTEGRATION = "INTEGRATION"
     SCHEDULED = "SCHEDULED"
+    CHAT_MESSAGE = "CHAT_MESSAGE"
 
 
 def get_trigger_type_mapping() -> Dict[Type["BaseTrigger"], WorkflowTriggerType]:

ee/vellum_ee/workflows/display/tests/workflow_serialization/test_chat_message_trigger_serialization.py

@@ -0,0 +1,62 @@
+"""Tests for ChatMessageTrigger serialization."""
+
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.chat_message_trigger_execution.workflows.simple_chat_workflow import SimpleChatWorkflow
+
+
+def test_simple_chat_workflow_serialization():
+    """SimpleChatWorkflow from tests/workflows serializes correctly with ChatMessageTrigger."""
+
+    # GIVEN a Workflow that uses a ChatMessageTrigger
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=SimpleChatWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+        "triggers",
+    }
+
+    # AND the triggers should be serialized correctly
+    assert serialized_workflow["triggers"] == [
+        {
+            "id": "9e14c49b-c6d9-4fe5-9ff2-835fd695fe5f",
+            "type": "CHAT_MESSAGE",
+            "attributes": [
+                {
+                    "id": "5edbfd78-b634-4305-b2ad-d9feecbd5e5f",
+                    "key": "message",
+                    "type": "JSON",
+                    "required": True,
+                    "default": {
+                        "type": "JSON",
+                        "value": None,
+                    },
+                    "extensions": None,
+                    "schema": None,
+                }
+            ],
+            "exec_config": {
+                "output": {
+                    "type": "NODE_OUTPUT",
+                    "node_id": "6c43f557-304c-4f08-a8fd-13d1fb02d96a",
+                    "node_output_id": "14f1265b-d5fb-4b60-b06b-9012029f6c6c",
+                },
+            },
+            "display_data": {
+                "label": "Chat Message",
+                "position": {
+                    "x": 0.0,
+                    "y": 0.0,
+                },
+                "z_index": 0,
+                "icon": "vellum:icon:message",
+                "color": "blue",
+            },
+        }
+    ]

ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -31,6 +31,7 @@
 from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, StateValueReference, WorkflowInputReference
+from vellum.workflows.triggers.chat_message import ChatMessageTrigger
 from vellum.workflows.triggers.integration import IntegrationTrigger
 from vellum.workflows.triggers.manual import ManualTrigger
 from vellum.workflows.triggers.schedule import ScheduleTrigger
@@ -634,6 +635,8 @@ def _serialize_workflow_trigger(self) -> Optional[JsonArray]:
                     trigger_type = WorkflowTriggerType.INTEGRATION
                 elif issubclass(trigger_class, ScheduleTrigger):
                     trigger_type = WorkflowTriggerType.SCHEDULED
+                elif issubclass(trigger_class, ChatMessageTrigger):
+                    trigger_type = WorkflowTriggerType.CHAT_MESSAGE
                 else:
                     raise ValueError(
                         f"Unknown trigger type: {trigger_class.__name__}. "
@@ -644,6 +647,13 @@ def _serialize_workflow_trigger(self) -> Optional[JsonArray]:
 
             # Serialize trigger attributes from attribute_references as VellumVariables
             attribute_references = trigger_class.attribute_references().values()
+
+            def get_attribute_type(reference: Any) -> str:
+                try:
+                    return primitive_type_to_vellum_variable_type(reference)
+                except ValueError:
+                    return "JSON"
+
             trigger_attributes: JsonArray = cast(
                 JsonArray,
                 [
@@ -652,10 +662,10 @@ def _serialize_workflow_trigger(self) -> Optional[JsonArray]:
                         {
                             "id": str(reference.id),
                             "key": reference.name,
-                            "type": primitive_type_to_vellum_variable_type(reference),
+                            "type": get_attribute_type(reference),
                             "required": type(None) not in reference.types,
                             "default": {
-                                "type": primitive_type_to_vellum_variable_type(reference),
+                                "type": get_attribute_type(reference),
                                 "value": None,
                             },
                             "extensions": None,
@@ -695,6 +705,11 @@ def _serialize_workflow_trigger(self) -> Optional[JsonArray]:
                     # Validate trigger attributes against the expected types from the API
                     self._validate_integration_trigger_attributes(trigger_class, trigger_attributes)
 
+                if trigger_type == WorkflowTriggerType.CHAT_MESSAGE and issubclass(trigger_class, ChatMessageTrigger):
+                    chat_exec_config = self._serialize_chat_message_trigger_exec_config(trigger_class)
+                    if chat_exec_config:
+                        trigger_data["exec_config"] = chat_exec_config
+
             # Serialize display_data from trigger's Display class
             display_class = trigger_class.Display
             display_data: JsonObject = {}
@@ -903,6 +918,28 @@ def _validate_integration_trigger_attributes(
                         trigger_class_name=trigger_class.__name__,
                     )
 
+    def _serialize_chat_message_trigger_exec_config(
+        self, trigger_class: Type[ChatMessageTrigger]
+    ) -> Optional[JsonObject]:
+        config_class = trigger_class.Config
+        output = getattr(config_class, "output", None)
+
+        if output is None:
+            return None
+
+        serialized_output = serialize_value(
+            executable_id=trigger_class.__id__,
+            display_context=self.display_context,
+            value=output,
+        )
+
+        return cast(
+            JsonObject,
+            {
+                "output": serialized_output,
+            },
+        )
+
     @staticmethod
     def _model_dump(value: Any) -> Any:
         if hasattr(value, "model_dump"):