vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

vellum_ee/workflows/display/tests/workflow_serialization/test_list_vellum_document_serialization.py

@@ -0,0 +1,65 @@
+from typing import List
+
+from vellum import VellumDocument
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes.displayable import FinalOutputNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.state import BaseState
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+class Inputs(BaseInputs):
+    documents: List[VellumDocument]
+
+
+class MyFinalOutputNode(FinalOutputNode):
+    class Outputs(BaseOutputs):
+        value = Inputs.documents
+
+
+class WorkflowWithListVellumDocument(BaseWorkflow[Inputs, BaseState]):
+    """A workflow that takes a list of VellumDocument as input."""
+
+    graph = MyFinalOutputNode
+
+    class Outputs(BaseOutputs):
+        result = MyFinalOutputNode.Outputs.value
+
+
+def test_serialize_workflow_with_list_vellum_document():
+    """Tests that we can serialize a workflow with List[VellumDocument] input."""
+
+    workflow_display = get_workflow_display(workflow_class=WorkflowWithListVellumDocument)
+
+    # WHEN we serialize it
+    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",
+    }
+
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+
+    input_var = input_variables[0]
+    assert input_var["key"] == "documents"
+    assert input_var["type"] == "JSON"
+    # NOTE: Once custom type serialization is supported, we will want to represent this using an openapi spec
+    assert input_var["required"] is True
+    assert input_var["default"] is None
+    assert input_var["extensions"] == {"color": None}
+
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+
+    output_var = output_variables[0]
+    assert output_var["key"] == "result"
+
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    assert len(workflow_raw_data["nodes"]) == 2
+    assert len(workflow_raw_data["edges"]) == 1

vellum_ee/workflows/display/tests/workflow_serialization/test_manual_trigger_serialization.py

@@ -0,0 +1,110 @@
+import pytest
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.triggers.base import BaseTrigger
+from vellum.workflows.triggers.manual import ManualTrigger
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_manual_trigger_serialization():
+    """Workflow with ManualTrigger serializes with triggers field."""
+
+    class SimpleNode(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = ManualTrigger >> SimpleNode
+
+    result = get_workflow_display(workflow_class=TestWorkflow).serialize()
+    assert "triggers" in result
+    triggers = result["triggers"]
+    assert isinstance(triggers, list)
+
+    assert len(triggers) == 1
+    assert triggers[0] == {"id": "b3c8ab56-001f-4157-bbc2-4a7fe5ebf8c6", "type": "MANUAL", "attributes": []}
+
+
+def test_manual_trigger_multiple_entrypoints():
+    """ManualTrigger with multiple entrypoints."""
+
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    class MultiWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = ManualTrigger >> {NodeA, NodeB}
+
+    result = get_workflow_display(workflow_class=MultiWorkflow).serialize()
+    workflow_data = result["workflow_raw_data"]
+    assert isinstance(workflow_data, dict)
+    assert "nodes" in workflow_data
+    nodes = workflow_data["nodes"]
+    assert isinstance(nodes, list)
+
+    # entrypoint + 2 nodes
+    assert len(nodes) == 3
+
+    assert "triggers" in result
+    triggers = result["triggers"]
+    assert isinstance(triggers, list)
+
+    assert len(triggers) == 1
+    assert triggers[0] == {"id": "b3c8ab56-001f-4157-bbc2-4a7fe5ebf8c6", "type": "MANUAL", "attributes": []}
+
+
+def test_unknown_trigger_type():
+    """Unknown trigger types raise ValueError."""
+
+    class UnknownTrigger(BaseTrigger):
+        pass
+
+    class SimpleNode(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = UnknownTrigger >> SimpleNode
+
+    with pytest.raises(ValueError, match="Unknown trigger type: UnknownTrigger"):
+        get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+
+def test_manual_trigger_entrypoint_id_consistency():
+    """ManualTrigger ID matches entrypoint node ID for backward compatibility."""
+
+    class SimpleNode(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = ManualTrigger >> SimpleNode
+
+    result = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    # Get trigger ID
+    triggers = result["triggers"]
+    assert isinstance(triggers, list)
+    assert len(triggers) == 1
+    trigger = triggers[0]
+    assert isinstance(trigger, dict)
+    trigger_id = trigger["id"]
+
+    # Get entrypoint node ID - ManualTrigger workflows still use ENTRYPOINT nodes
+    workflow_raw_data = result["workflow_raw_data"]
+    assert isinstance(workflow_raw_data, dict)
+    nodes = workflow_raw_data["nodes"]
+    assert isinstance(nodes, list)
+    entrypoint_nodes = [n for n in nodes if isinstance(n, dict) and n["type"] == "ENTRYPOINT"]
+    assert len(entrypoint_nodes) == 1, "ManualTrigger workflows should have ENTRYPOINT node for backward compatibility"
+    entrypoint_id = entrypoint_nodes[0]["id"]
+
+    # Verify IDs match
+    assert trigger_id == entrypoint_id, (
+        f"ManualTrigger ID ({trigger_id}) must match entrypoint node ID ({entrypoint_id}) "
+        "to maintain trigger-entrypoint linkage"
+    )
+    # Also verify the expected UUID
+    assert trigger_id == "b3c8ab56-001f-4157-bbc2-4a7fe5ebf8c6"

vellum_ee/workflows/display/tests/workflow_serialization/test_map_node_with_terminal_nodes_serialization.py

@@ -0,0 +1,62 @@
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.map_node_with_terminal_nodes.workflow import MapNodeWithTerminalNodes
+
+
+def test_serialize_workflow__terminal_nodes_have_different_output_ids():
+    """
+    Tests that terminal nodes with the same base class serialize to different output IDs
+    when one is at the top level and one is inside a map node.
+    """
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MapNodeWithTerminalNodes)
+    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",
+    }
+
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+
+    map_node = None
+    top_level_terminal_node = None
+
+    for node in workflow_raw_data["nodes"]:
+        if node["type"] == "MAP":
+            map_node = node
+        elif node["type"] == "TERMINAL":
+            top_level_terminal_node = node
+
+    assert map_node is not None, "Map node should be present"
+    assert top_level_terminal_node is not None, "Top-level terminal node should be present"
+
+    subworkflow_raw_data = map_node["data"]["workflow_raw_data"]
+    map_terminal_node = None
+
+    for node in subworkflow_raw_data["nodes"]:
+        if node["type"] == "TERMINAL":
+            map_terminal_node = node
+            break
+
+    assert map_terminal_node is not None, "Terminal node inside map should be present"
+
+    assert top_level_terminal_node["base"]["name"] == "FinalOutputNode", (
+        f"Top-level terminal should have base name 'FinalOutputNode', " f"got {top_level_terminal_node['base']['name']}"
+    )
+    assert map_terminal_node["base"]["name"] == "FinalOutputNode", (
+        f"Map terminal should have base name 'FinalOutputNode', " f"got {map_terminal_node['base']['name']}"
+    )
+
+    # THEN the top-level terminal node and the map's terminal node should have different output IDs
+    top_level_output_id = top_level_terminal_node["data"]["output_id"]
+    map_terminal_output_id = map_terminal_node["data"]["output_id"]
+
+    assert top_level_output_id != map_terminal_output_id, (
+        f"Terminal nodes should have different output IDs. "
+        f"Top-level: {top_level_output_id}, Map terminal: {map_terminal_output_id}"
+    )

vellum_ee/workflows/display/tests/workflow_serialization/test_multi_trigger_same_node_serialization.py

@@ -0,0 +1,99 @@
+"""Tests for multiple triggers pointing to the same node."""
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.constants import VellumIntegrationProviderType
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.triggers.integration import IntegrationTrigger
+from vellum.workflows.triggers.manual import ManualTrigger
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_manual_and_slack_trigger_same_node():
+    """
+    Tests that when Manual and Slack triggers point to the same node,
+    we have edges for both triggers.
+    """
+
+    class MyManualTrigger(ManualTrigger):
+        pass
+
+    class SlackMessageTrigger(IntegrationTrigger):
+        message: str
+        channel: str
+
+        class Config(IntegrationTrigger.Config):
+            provider = VellumIntegrationProviderType.COMPOSIO
+            integration_name = "SLACK"
+            slug = "slack_new_message"
+
+    class ProcessNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+        def run(self) -> Outputs:
+            return self.Outputs(result="processed")
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = {
+            MyManualTrigger >> ProcessNode,
+            SlackMessageTrigger >> ProcessNode,
+        }
+
+    result = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    assert "triggers" in result
+    triggers = result["triggers"]
+    assert isinstance(triggers, list)
+    assert len(triggers) == 2
+
+    manual_trigger = next((t for t in triggers if isinstance(t, dict) and t.get("type") == "MANUAL"), None)
+    assert manual_trigger is not None, "Should have a Manual trigger"
+    assert isinstance(manual_trigger, dict)
+    manual_trigger_id = manual_trigger["id"]
+
+    slack_trigger = next((t for t in triggers if isinstance(t, dict) and t.get("type") == "INTEGRATION"), None)
+    assert slack_trigger is not None, "Should have a Slack trigger"
+    assert isinstance(slack_trigger, dict)
+    slack_trigger_id = slack_trigger["id"]
+
+    workflow_raw_data = result["workflow_raw_data"]
+    assert isinstance(workflow_raw_data, dict)
+    edges = workflow_raw_data["edges"]
+    assert isinstance(edges, list)
+
+    nodes = workflow_raw_data["nodes"]
+    assert isinstance(nodes, list)
+    process_nodes = [n for n in nodes if isinstance(n, dict) and n.get("type") not in ("TERMINAL", "ENTRYPOINT")]
+    assert len(process_nodes) > 0, "Should have at least one process node"
+    process_node = process_nodes[0]
+    process_node_id = process_node["id"]
+
+    manual_edge = next(
+        (
+            e
+            for e in edges
+            if isinstance(e, dict)
+            and e.get("source_node_id") == manual_trigger_id
+            and e.get("target_node_id") == process_node_id
+        ),
+        None,
+    )
+    assert manual_edge is not None, (
+        f"Should have edge from Manual trigger ({manual_trigger_id}) " f"to ProcessNode ({process_node_id})"
+    )
+
+    slack_edge = next(
+        (
+            e
+            for e in edges
+            if isinstance(e, dict)
+            and e.get("source_node_id") == slack_trigger_id
+            and e.get("target_node_id") == process_node_id
+        ),
+        None,
+    )
+    assert slack_edge is not None, (
+        f"Should have edge from Slack trigger ({slack_trigger_id}) " f"to ProcessNode ({process_node_id})"
+    )

vellum_ee/workflows/display/tests/workflow_serialization/test_scheduled_trigger_serialization.py

@@ -0,0 +1,267 @@
+from contextlib import contextmanager
+import json
+import os
+import tempfile
+from unittest.mock import patch
+from uuid import UUID
+from typing import Iterator, Tuple
+
+from deepdiff import DeepDiff
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.triggers.base import _get_trigger_path_to_id_mapping
+from vellum.workflows.triggers.schedule import ScheduleTrigger
+from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.editor import NodeDisplayComment
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+@contextmanager
+def mock_metadata_json(trigger_class_name: str, trigger_id: UUID) -> Iterator[Tuple[str, UUID]]:
+    """
+    Context manager that mocks metadata.json with trigger ID mapping.
+
+    Args:
+        trigger_class_name: Name of the trigger class (e.g., "DailyScheduleTrigger")
+        trigger_id: The UUID to map the trigger to
+
+    Yields:
+        Tuple of (trigger_module_path, trigger_id)
+    """
+    # Create a temporary directory to simulate a workflow module
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Create metadata.json with trigger mapping
+        metadata_path = os.path.join(tmpdir, "metadata.json")
+
+        test_module_parts = __name__.split(".")
+        workflow_root_module = ".".join(test_module_parts[:-1])  # Remove last part
+
+        relative_module_path = "." + test_module_parts[-1]
+        # Include the full trigger path with class name to match TypeScript codegen
+        relative_trigger_path = f"{relative_module_path}.{trigger_class_name}"
+
+        with open(metadata_path, "w") as f:
+            json.dump(
+                {
+                    "trigger_path_to_id_mapping": {
+                        # Use full trigger path including class name
+                        relative_trigger_path: str(trigger_id),
+                    }
+                },
+                f,
+            )
+
+        # Create a custom mock for virtual_open that returns our tmpdir file
+        def mock_virtual_open(path):
+            # Regardless of the path requested, return our test metadata.json
+            return open(metadata_path)
+
+        # Mock the workflow root finder to return the parent module path
+        with patch(
+            "vellum.workflows.triggers.base._find_workflow_root_with_metadata", return_value=workflow_root_module
+        ):
+            with patch("vellum.workflows.triggers.base.virtual_open", side_effect=mock_virtual_open):
+                # Clear the LRU cache to ensure fresh read
+                _get_trigger_path_to_id_mapping.cache_clear()
+
+                yield __name__, trigger_id
+
+
+def test_scheduled_trigger_serialization_with_metadata_json():
+    """ScheduleTrigger uses mapped ID from metadata.json when available."""
+
+    # GIVEN we have an expected trigger ID
+    expected_trigger_id = UUID("12345678-1234-1234-1234-123456789abc")
+
+    # AND we have metadata.json set up with the trigger mapping
+    with mock_metadata_json(
+        "test_scheduled_trigger_serialization_with_metadata_json.<locals>.DailyScheduleTrigger",
+        expected_trigger_id,
+    ):
+
+        class DailyScheduleTrigger(ScheduleTrigger):
+            class Config(ScheduleTrigger.Config):
+                cron = "0 9 * * *"  # Every day at 9am
+                timezone = "America/New_York"
+
+        class ProcessNode(BaseNode):
+            class Outputs(BaseNode.Outputs):
+                timestamp = DailyScheduleTrigger.current_run_at
+
+            def run(self) -> Outputs:
+                return self.Outputs()
+
+        # AND a workflow that has the scheduled trigger as the entrypoint
+        class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+            graph = DailyScheduleTrigger >> ProcessNode
+
+        # WHEN we serialize the workflow
+        # The trigger ID should be automatically loaded from metadata.json by the metaclass
+        result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+        # THEN we get the expected trigger
+        assert "triggers" in result
+        triggers = result["triggers"]
+        assert isinstance(triggers, list)
+        assert len(triggers) == 1
+
+        trigger = triggers[0]
+        assert isinstance(trigger, dict)
+        assert trigger["type"] == "SCHEDULED"
+
+        # AND the id is the one from metadata.json
+        assert trigger["id"] == str(expected_trigger_id)
+
+
+def test_scheduled_trigger_serialization_without_metadata_json():
+    """ScheduleTrigger uses default hash-based __id__ when no metadata.json is available."""
+
+    # GIVEN a scheduled trigger defined without metadata.json
+    class WeeklyScheduleTrigger(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "0 9 * * MON"  # Every Monday at 9am
+            timezone = "UTC"
+
+    class ProcessNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            timestamp = WeeklyScheduleTrigger.current_run_at
+
+        def run(self) -> Outputs:
+            return self.Outputs()
+
+    # AND a workflow that has the scheduled trigger as the entrypoint
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = WeeklyScheduleTrigger >> ProcessNode
+
+    # WHEN we serialize the workflow without metadata.json
+    result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    # THEN we get the expected trigger
+    assert "triggers" in result
+    triggers = result["triggers"]
+    assert isinstance(triggers, list)
+    assert len(triggers) == 1
+
+    trigger = triggers[0]
+    assert isinstance(trigger, dict)
+    assert trigger["type"] == "SCHEDULED"
+
+    # AND the id is the trigger class's default __id__ (hash-based)
+    expected_module_path = (
+        f"{__name__}.test_scheduled_trigger_serialization_without_metadata_json.<locals>.WeeklyScheduleTrigger"
+    )
+    expected_default_id = str(uuid4_from_hash(expected_module_path))
+    assert trigger["id"] == expected_default_id
+
+    # AND attributes are serialized
+    assert "attributes" in trigger
+    attributes = trigger["attributes"]
+    assert isinstance(attributes, list)
+
+
+def test_scheduled_trigger_serialization_display_data():
+    """ScheduleTrigger with Display class serializes all display attributes correctly."""
+
+    # GIVEN a scheduled trigger with comprehensive Display attributes
+    class DailyTriggerWithDisplay(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "0 9 * * *"
+            timezone = "UTC"
+
+        class Display(ScheduleTrigger.Display):
+            label = "Daily Schedule"
+            x = 100.5
+            y = 200.75
+            z_index = 3
+            icon = "vellum:icon:calendar"
+            color = "#4A90E2"
+
+    class ProcessNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            timestamp = DailyTriggerWithDisplay.current_run_at
+
+        def run(self) -> Outputs:
+            return self.Outputs()
+
+    # AND a workflow that uses the trigger
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = DailyTriggerWithDisplay >> ProcessNode
+
+    # WHEN we serialize the workflow
+    result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    # THEN we get the expected trigger
+    assert "triggers" in result
+    triggers = result["triggers"]
+    assert isinstance(triggers, list)
+    assert len(triggers) == 1
+
+    trigger = triggers[0]
+    assert isinstance(trigger, dict)
+    assert trigger["type"] == "SCHEDULED"
+
+    # AND display_data is serialized with icon and color
+    assert "display_data" in trigger
+    display_data = trigger["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["icon"] == "vellum:icon:calendar"
+    assert display_data["color"] == "#4A90E2"
+
+
+def test_scheduled_trigger_serialization_full():
+    # GIVEN a scheduled trigger with comprehensive Display attributes
+    class DailyTriggerWithDisplay(ScheduleTrigger):
+        class Config(ScheduleTrigger.Config):
+            cron = "0 9 * * *"
+            timezone = "UTC"
+
+        class Display(ScheduleTrigger.Display):
+            label = "Daily Schedule"
+            x = 100.5
+            y = 200.75
+            z_index = 3
+            icon = "vellum:icon:calendar"
+            color = "#4A90E2"
+            comment = NodeDisplayComment(value="This is scheduled trigger", expanded=True)
+
+    class ProcessNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            timestamp = DailyTriggerWithDisplay.current_run_at
+
+        def run(self) -> Outputs:
+            return self.Outputs()
+
+    # AND a workflow that uses the trigger
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = DailyTriggerWithDisplay >> ProcessNode
+
+    # WHEN we serialize the workflow
+    result: dict = get_workflow_display(workflow_class=TestWorkflow).serialize()
+
+    # THEN we get the expected trigger
+    assert len(result["triggers"]) == 1
+    trigger = result["triggers"][0]
+    assert not DeepDiff(
+        trigger,
+        {
+            "id": "f3e5eddb-75da-42e6-9abf-d616f30c145c",
+            "type": "SCHEDULED",
+            "cron": "0 9 * * *",
+            "timezone": "UTC",
+            "attributes": [],
+            "display_data": {
+                "label": "Daily Schedule",
+                "position": {"x": 100.5, "y": 200.75},
+                "z_index": 3,
+                "icon": "vellum:icon:calendar",
+                "color": "#4A90E2",
+                "comment": {
+                    "value": "This is scheduled trigger",
+                    "expanded": True,
+                },
+            },
+        },
+    )

vellum_ee/workflows/display/tests/workflow_serialization/test_terminal_node_any_serialization.py

@@ -0,0 +1,49 @@
+from typing import Any
+
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+class Inputs(BaseInputs):
+    input: Any
+
+
+class AnyFinalOutputNode(FinalOutputNode[BaseState, Any]):
+    class Outputs(FinalOutputNode.Outputs):
+        value: Any = Inputs.input
+
+
+class AnyFinalOutputNodeWorkflow(BaseWorkflow[Inputs, BaseState]):
+    graph = AnyFinalOutputNode
+
+    class Outputs(BaseWorkflow.Outputs):
+        value = AnyFinalOutputNode.Outputs.value
+
+
+def test_serialize_workflow_with_any_output_type():
+    """
+    Tests that a terminal node with Any output type can be serialized correctly.
+    """
+    # GIVEN a Workflow that uses a Final Output Node with Any type
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=AnyFinalOutputNodeWorkflow)
+    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",
+    }
+
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    terminal_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "TERMINAL")
+    assert terminal_node["data"]["output_type"] == "JSON"
+
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert output_variables[0]["type"] == "JSON"