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

vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py

@@ -42,7 +42,7 @@ def test_try_node_display__serialize_with_error_output() -> None:
             {
                 "data": {
                     "node_id": str(MyNode.__wrapped_node__.__id__),
-                    "output_id": "c8fbe459-c9ee-4639-a82d-961180cf9411",
+                    "output_id": "75bb6e13-0cf0-4651-9141-658be1d5acaa",
                 },
                 "type": "NODE_OUTPUT",
             }
@@ -54,7 +54,7 @@ def test_try_node_display__serialize_with_error_output() -> None:
             {
                 "data": {
                     "node_id": str(MyNode.__wrapped_node__.__id__),
-                    "output_id": "efe6e307-3ea4-4862-a26f-4c4416bb4537",
+                    "output_id": "16b4d0ee-d309-4622-b812-3bd8f8c8a73c",
                 },
                 "type": "NODE_OUTPUT",
             }

vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -8,6 +8,7 @@ from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.references import LazyReference
+from vellum_ee.workflows.display.base import WorkflowInputsDisplay, WorkflowMetaDisplay
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
@@ -21,7 +22,6 @@ from vellum_ee.workflows.display.utils.vellum import (
     NodeOutputData,
     NodeOutputPointer,
 )
-from vellum_ee.workflows.display.vellum import WorkflowInputsVellumDisplayOverrides, WorkflowMetaVellumDisplay
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
@@ -57,7 +57,7 @@ class MyNodeB(BaseNode):
             [
                 NodeOutputPointer(
                     data=NodeOutputData(
-                        node_id="b48fa5e0-d7d3-4fe3-ae48-615415011cc5",
+                        node_id="f2f8f47a-25de-474c-95d7-df4f9b4cb5ac",
                         output_id="4b16a629-11a1-4b3f-a965-a57b872d13b8",
                     )
                 )
@@ -68,7 +68,7 @@ class MyNodeB(BaseNode):
             [
                 NodeOutputPointer(
                     data=NodeOutputData(
-                        node_id="b48fa5e0-d7d3-4fe3-ae48-615415011cc5",
+                        node_id="f2f8f47a-25de-474c-95d7-df4f9b4cb5ac",
                         output_id="4b16a629-11a1-4b3f-a965-a57b872d13b8",
                     ),
                 ),
@@ -92,7 +92,7 @@ class MyNodeB(BaseNode):
             [
                 NodeOutputPointer(
                     data=NodeOutputData(
-                        node_id="b48fa5e0-d7d3-4fe3-ae48-615415011cc5",
+                        node_id="f2f8f47a-25de-474c-95d7-df4f9b4cb5ac",
                         output_id="4b16a629-11a1-4b3f-a965-a57b872d13b8",
                     ),
                 ),
@@ -110,20 +110,19 @@ def test_create_node_input_value_pointer_rules(
         descriptor,
         WorkflowDisplayContext(
             workflow_display_class=BaseWorkflowDisplay,
-            workflow_display=WorkflowMetaVellumDisplay(
+            workflow_display=WorkflowMetaDisplay(
                 entrypoint_node_id=uuid4(),
                 entrypoint_node_source_handle_id=uuid4(),
                 entrypoint_node_display=NodeDisplayData(),
             ),
             global_workflow_input_displays={
-                Inputs.example_workflow_input: WorkflowInputsVellumDisplayOverrides(
+                Inputs.example_workflow_input: WorkflowInputsDisplay(
                     id=UUID("a154c29d-fac0-4cd0-ba88-bc52034f5470"),
                 ),
             },
             global_node_output_displays={
-                MyNodeA.Outputs.output: (
-                    MyNodeA,
-                    NodeOutputDisplay(id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"),
+                MyNodeA.Outputs.output: NodeOutputDisplay(
+                    id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"
                 ),
             },
             global_node_displays={

vellum_ee/workflows/display/nodes/vellum/try_node.py

@@ -1,8 +1,7 @@
 import inspect
 from uuid import UUID
-from typing import Any, ClassVar, Generic, Optional, Tuple, Type, TypeVar
+from typing import Any, ClassVar, Generic, Optional, TypeVar
 
-from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
 from vellum.workflows.references.output import OutputReference
@@ -37,13 +36,13 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
                 {
                     "id": id,
                     "name": attribute.name,
-                    "value": serialize_value(display_context, attribute.instance),
+                    "value": serialize_value(node_id, display_context, attribute.instance),
                 }
             )
 
         adornment: JsonObject = {
             "id": str(node_id),
-            "label": node.__qualname__,
+            "label": self.label,
             "base": self.get_base().dict(),
             "attributes": attributes,
         }
@@ -76,7 +75,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
 
         return serialized_node
 
-    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
         inner_node = self._node.__wrapped_node__
         if not inner_node:
             return super().get_node_output_display(output)
@@ -84,7 +83,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
         node_display_class = get_node_display_class(inner_node)
         node_display = node_display_class()
         if output.name == "error":
-            return inner_node, NodeOutputDisplay(
+            return NodeOutputDisplay(
                 id=self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id"),
                 name="error",
             )

vellum_ee/workflows/display/tests/test_base_workflow_display.py

@@ -1,13 +1,19 @@
 from uuid import UUID
-from typing import Dict
+from typing import Any, Dict, List, cast
 
+from vellum.workflows.constants import VellumIntegrationProviderType
+from vellum.workflows.graph import Graph
 from vellum.workflows.inputs import BaseInputs
-from vellum.workflows.nodes import BaseNode
+from vellum.workflows.nodes import BaseNode, InlineSubworkflowNode
+from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state import BaseState
+from vellum.workflows.triggers.integration import IntegrationTrigger
 from vellum.workflows.workflows.base import BaseWorkflow
-from vellum_ee.workflows.display.vellum import WorkflowInputsVellumDisplayOverrides
+from vellum_ee.workflows.display.base import EdgeDisplay, WorkflowInputsDisplay
+from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
+from vellum_ee.workflows.display.nodes import BaseNodeDisplay
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -39,7 +45,7 @@ def test_base_workflow_display__serialize_empty_workflow():
                     "data": {"label": "Entrypoint Node", "source_handle_id": "0af025a4-3b25-457d-a7ae-e3a7ba15c86c"},
                     "base": None,
                     "definition": None,
-                    "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                    "display_data": {"position": {"x": 0.0, "y": -50.0}},
                     "id": "3c41cdd9-999a-48b8-9088-f6dfa1369bfd",
                     "inputs": [],
                     "type": "ENTRYPOINT",
@@ -64,9 +70,7 @@ def test_vellum_workflow_display__serialize_input_variables_with_capitalized_var
 
     class ExampleWorkflowDisplay(BaseWorkflowDisplay[ExampleWorkflow]):
         inputs_display = {
-            Inputs.foo: WorkflowInputsVellumDisplayOverrides(
-                id=UUID("97b63d71-5413-417f-9cf5-49e1b4fd56e4"), name="Foo", required=True
-            )
+            Inputs.foo: WorkflowInputsDisplay(id=UUID("97b63d71-5413-417f-9cf5-49e1b4fd56e4"), name="Foo")
         }
 
     display = get_workflow_display(
@@ -229,6 +233,7 @@ def test_vellum_workflow_display__serialize_with_parse_json_expression():
     from vellum.workflows.references.constant import ConstantValueReference
 
     class JsonNode(BaseNode):
+
         class Outputs(BaseNode.Outputs):
             json_result = ConstantValueReference('{"key": "value"}').parse_json()
 
@@ -272,7 +277,7 @@ def test_vellum_workflow_display__serialize_with_parse_json_expression():
             break
 
     assert json_result == {
-        "id": "44c7d94c-a76a-4151-9b95-85a31764f18f",
+        "id": "e73fd6b1-1109-4a97-8510-c9ba8e6f5dbe",
         "name": "json_result",
         "type": "JSON",
         "value": {
@@ -292,6 +297,7 @@ def test_vellum_workflow_display__serialize_with_parse_json_expression():
 def test_serialize__port_with_lazy_reference():
     # GIVEN a node with a lazy reference in a Port
     class MyNode(BaseNode):
+
         class Ports(BaseNode.Ports):
             foo = Port.on_if(LazyReference(lambda: MyNode.Outputs.bar))
 
@@ -319,7 +325,7 @@ def test_serialize__port_with_lazy_reference():
     assert isinstance(ports, list)
     assert ports == [
         {
-            "id": "6c26bc2b-6469-47c1-b858-d63f0d311ea6",
+            "id": "02e90580-e6b6-441c-a5d4-7cd083c11fc7",
             "name": "foo",
             "type": "IF",
             "expression": {
@@ -329,3 +335,371 @@ def test_serialize__port_with_lazy_reference():
             },
         }
     ]
+
+
+def test_global_propagation_deep_nested_subworkflows():
+    # GIVEN the root workflow, a middle workflow, and an inner workflow
+
+    class RootInputs(BaseInputs):
+        root_param: str
+
+    class MiddleInputs(BaseInputs):
+        middle_param: str
+
+    class InnerInputs(BaseInputs):
+        inner_param: str
+
+    class InnerNode(BaseNode):
+        class Outputs(BaseOutputs):
+            done: bool
+
+        def run(self) -> Outputs:
+            return self.Outputs(done=True)
+
+    class InnerWorkflow(BaseWorkflow[InnerInputs, BaseState]):
+        graph = InnerNode
+
+    class MiddleInlineSubworkflowNode(InlineSubworkflowNode):
+        subworkflow_inputs = {"inner_param": "x"}
+        subworkflow = InnerWorkflow
+
+    class MiddleWorkflow(BaseWorkflow[MiddleInputs, BaseState]):
+        graph = MiddleInlineSubworkflowNode
+
+    class OuterInlineSubworkflowNode(InlineSubworkflowNode):
+        subworkflow_inputs = {"middle_param": "y"}
+        subworkflow = MiddleWorkflow
+
+    class RootWorkflow(BaseWorkflow[RootInputs, BaseState]):
+        graph = OuterInlineSubworkflowNode
+
+    # WHEN we build the displays
+    root_display = get_workflow_display(workflow_class=RootWorkflow)
+    middle_display = get_workflow_display(
+        workflow_class=MiddleWorkflow, parent_display_context=root_display.display_context
+    )
+    inner_display = get_workflow_display(
+        workflow_class=InnerWorkflow, parent_display_context=middle_display.display_context
+    )
+
+    # THEN the deepest display must include root + middle + inner inputs in its GLOBAL view
+    inner_global_names = {ref.name for ref in inner_display.display_context.global_workflow_input_displays.keys()}
+
+    assert inner_global_names == {"middle_param", "inner_param", "root_param"}
+
+
+def test_serialize_workflow_with_edge_display_data():
+    """
+    Tests that edges with z_index values serialize display_data correctly.
+    """
+
+    # GIVEN a workflow with connected nodes
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class EndNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            final: str
+
+    class TestWorkflow(BaseWorkflow):
+        graph = StartNode >> EndNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_result = EndNode.Outputs.final
+
+    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
+        edge_displays = {
+            (StartNode.Ports.default, EndNode): EdgeDisplay(id=UUID("12345678-1234-5678-1234-567812345678"), z_index=5)
+        }
+
+    # WHEN we serialize the workflow with the custom display
+    display = get_workflow_display(
+        base_display_class=TestWorkflowDisplay,
+        workflow_class=TestWorkflow,
+    )
+    serialized_workflow = display.serialize()
+
+    # THEN the edge should include display_data with z_index
+    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
+    edges = cast(List[Dict[str, Any]], workflow_raw_data["edges"])
+
+    edge_with_display_data = None
+    for edge in edges:
+        if edge["id"] == "12345678-1234-5678-1234-567812345678":
+            edge_with_display_data = edge
+            break
+
+    assert edge_with_display_data is not None, "Edge with custom UUID not found"
+    assert edge_with_display_data["display_data"] == {"z_index": 5}
+
+    assert edge_with_display_data["type"] == "DEFAULT"
+    assert "source_node_id" in edge_with_display_data
+    assert "target_node_id" in edge_with_display_data
+
+
+def test_serialize_workflow_with_node_display_data():
+    """
+    Tests that nodes with z_index values serialize display_data correctly.
+    """
+
+    # GIVEN a workflow with a node that has custom display data
+    class TestNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class TestWorkflow(BaseWorkflow):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_result = TestNode.Outputs.result
+
+    class TestNodeDisplay(BaseNodeDisplay[TestNode]):
+        display_data = NodeDisplayData(position=NodeDisplayPosition(x=100, y=200), z_index=10, width=300, height=150)
+
+    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
+        pass
+
+    # WHEN we serialize the workflow with the custom node display
+    display = get_workflow_display(
+        base_display_class=TestWorkflowDisplay,
+        workflow_class=TestWorkflow,
+    )
+    serialized_workflow = display.serialize()
+
+    # THEN the node should include display_data with z_index
+    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
+    nodes = cast(List[Dict[str, Any]], workflow_raw_data["nodes"])
+
+    test_node = None
+    for node in nodes:
+        if node.get("type") == "GENERIC":
+            definition = node.get("definition")
+            if isinstance(definition, dict) and definition.get("name") == "TestNode":
+                test_node = node
+                break
+
+    assert test_node is not None, "TestNode not found in serialized nodes"
+    assert test_node["display_data"] == {"position": {"x": 100, "y": 200}, "z_index": 10, "width": 300, "height": 150}
+
+
+def test_serialize_workflow_with_node_icon_and_color():
+    """
+    Tests that nodes with icon and color serialize correctly in workflow context.
+    """
+
+    # GIVEN a workflow with a node that has icon and color
+    class TestNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class TestWorkflow(BaseWorkflow):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_result = TestNode.Outputs.result
+
+    class TestNodeDisplay(BaseNodeDisplay[TestNode]):
+        display_data = NodeDisplayData(position=NodeDisplayPosition(x=100, y=200), icon="vellum:icon:cog", color="navy")
+
+    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
+        pass
+
+    # WHEN we serialize the workflow
+    display = get_workflow_display(
+        base_display_class=TestWorkflowDisplay,
+        workflow_class=TestWorkflow,
+    )
+    serialized_workflow = display.serialize()
+
+    # THEN the node should include icon and color in display_data
+    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
+    nodes = cast(List[Dict[str, Any]], workflow_raw_data["nodes"])
+
+    test_node = None
+    for node in nodes:
+        if node.get("type") == "GENERIC":
+            definition = node.get("definition")
+            if isinstance(definition, dict) and definition.get("name") == "TestNode":
+                test_node = node
+                break
+
+    assert test_node is not None, "TestNode not found in serialized nodes"
+    assert test_node["display_data"]["icon"] == "vellum:icon:cog"
+    assert test_node["display_data"]["color"] == "navy"
+
+
+def test_base_workflow_display__graph_with_trigger_and_regular_node():
+    """
+    Tests that a workflow with both a trigger edge and a regular node edge serializes correctly.
+    """
+
+    # GIVEN a workflow with a trigger and regular nodes
+    class SlackMessageTrigger(IntegrationTrigger):
+        message: str
+
+        class Config(IntegrationTrigger.Config):
+            provider = VellumIntegrationProviderType.COMPOSIO
+            integration_name = "SLACK"
+            slug = "slack_message"
+
+    class TopNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class BottomNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            final: str
+
+    class TestWorkflow(BaseWorkflow):
+        graph = {
+            Graph.from_node(TopNode),
+            SlackMessageTrigger >> BottomNode,
+        }
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
+    exec_config = workflow_display.serialize()
+
+    # THEN the serialized workflow should have the expected structure
+    workflow_raw_data = exec_config["workflow_raw_data"]
+    assert isinstance(workflow_raw_data, dict)
+
+    nodes = workflow_raw_data["nodes"]
+    edges = workflow_raw_data["edges"]
+    triggers = exec_config["triggers"]
+
+    assert isinstance(nodes, list)
+    assert isinstance(edges, list)
+    assert isinstance(triggers, list)
+
+    assert len(triggers) == 1
+    trigger = triggers[0]
+    assert isinstance(trigger, dict)
+    trigger_id = trigger["id"]
+
+    entrypoint_nodes = [n for n in nodes if isinstance(n, dict) and n.get("type") == "ENTRYPOINT"]
+    assert len(entrypoint_nodes) == 1
+    entrypoint_node = entrypoint_nodes[0]
+    assert isinstance(entrypoint_node, dict)
+    entrypoint_node_id = entrypoint_node["id"]
+
+    top_node = next(
+        (
+            node
+            for node in nodes
+            if isinstance(node, dict)
+            and isinstance(node.get("definition"), dict)
+            and cast(Dict[str, Any], node.get("definition")).get("name") == "TopNode"
+        ),
+        None,
+    )
+    bottom_node = next(
+        (
+            node
+            for node in nodes
+            if isinstance(node, dict)
+            and isinstance(node.get("definition"), dict)
+            and cast(Dict[str, Any], node.get("definition")).get("name") == "BottomNode"
+        ),
+        None,
+    )
+
+    assert top_node is not None, "TopNode not found in serialized nodes"
+    assert bottom_node is not None, "BottomNode not found in serialized nodes"
+    top_node_id = top_node["id"]
+    bottom_node_id = bottom_node["id"]
+
+    entrypoint_edges = [e for e in edges if isinstance(e, dict) and e.get("source_node_id") == entrypoint_node_id]
+    assert len(entrypoint_edges) == 1, "Should have exactly one edge from the entrypoint node"
+
+    entrypoint_edge = entrypoint_edges[0]
+    assert isinstance(entrypoint_edge, dict)
+    assert entrypoint_edge["target_node_id"] == top_node_id, "Entrypoint edge should connect to TopNode"
+
+    trigger_edges = [e for e in edges if isinstance(e, dict) and e.get("source_node_id") == trigger_id]
+    assert len(trigger_edges) == 1, "Should have exactly one edge from the trigger"
+
+    trigger_edge = trigger_edges[0]
+    assert isinstance(trigger_edge, dict)
+    assert trigger_edge["target_node_id"] == bottom_node_id, "Trigger edge should connect to BottomNode"
+
+
+def test_serialize_subworkflow_output_reference_without_display():
+    """
+    Tests that a workflow with a subworkflow node followed by a generic node
+    that references the subworkflow output can be serialized correctly without a display directory.
+    """
+
+    # GIVEN a workflow module with a subworkflow node and a follow-on node
+    module_path = "tests.workflows.subworkflow_output_reference_without_display"
+
+    # WHEN we serialize the module
+    result = BaseWorkflowDisplay.serialize_module(module_path)
+
+    # THEN it should serialize successfully
+    assert hasattr(result, "exec_config")
+    assert isinstance(result.exec_config, dict)
+    assert "workflow_raw_data" in result.exec_config
+
+    # AND the workflow should have the expected nodes
+    workflow_raw_data = result.exec_config["workflow_raw_data"]
+    assert isinstance(workflow_raw_data, dict)
+    nodes = workflow_raw_data.get("nodes")
+    assert isinstance(nodes, list)
+
+    subworkflow_node = next(
+        (
+            node
+            for node in nodes
+            if isinstance(node, dict)
+            and isinstance(node.get("definition"), dict)
+            and node.get("definition", {}).get("name") == "SubworkflowNodeExample"
+        ),
+        None,
+    )
+    followon_node = next(
+        (
+            node
+            for node in nodes
+            if isinstance(node, dict)
+            and isinstance(node.get("definition"), dict)
+            and node.get("definition", {}).get("name") == "FollowOnNode"
+        ),
+        None,
+    )
+
+    assert subworkflow_node is not None, "SubworkflowNodeExample not found in serialized nodes"
+    assert followon_node is not None, "FollowOnNode not found in serialized nodes"
+
+    # AND the subworkflow node should have output variables
+    subworkflow_data = subworkflow_node.get("data", {})
+    subworkflow_outputs = subworkflow_data.get("output_variables", [])
+    assert len(subworkflow_outputs) > 0, "SubworkflowNodeExample should have output variables"
+
+    result_output = next((out for out in subworkflow_outputs if out.get("key") == "result"), None)
+    assert result_output is not None, "SubworkflowNodeExample should have 'result' output"
+    expected_output_id = result_output["id"]
+
+    # AND the follow-on node should have an attribute that references the subworkflow output
+    followon_attributes = followon_node.get("attributes", [])
+    assert len(followon_attributes) > 0, "FollowOnNode should have attributes"
+
+    input_value_attr = next(
+        (attr for attr in followon_attributes if isinstance(attr, dict) and attr.get("name") == "input_value"), None
+    )
+
+    assert input_value_attr is not None, "input_value attribute not found in FollowOnNode"
+
+    # AND the attribute value should reference the subworkflow node's output
+    value = input_value_attr.get("value")
+    assert isinstance(value, dict)
+    assert value.get("type") == "NODE_OUTPUT"
+    assert value.get("node_id") == subworkflow_node.get("id")
+
+    # AND the node_output_id should match the actual output ID from the subworkflow
+    actual_output_id = value.get("node_output_id")
+    assert actual_output_id == expected_output_id, (
+        f"node_output_id mismatch: expected {expected_output_id}, got {actual_output_id}. "
+        "The reference to the subworkflow output is broken."
+    )

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py

@@ -4,6 +4,7 @@ from typing import Type
 
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
+from vellum_ee.workflows.display.base import WorkflowMetaDisplay
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.types import (
@@ -13,7 +14,6 @@ from vellum_ee.workflows.display.types import (
     WorkflowDisplayContext,
     WorkflowInputsDisplays,
 )
-from vellum_ee.workflows.display.vellum import WorkflowMetaVellumDisplay
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
@@ -31,7 +31,7 @@ def serialize_node():
 
         context: WorkflowDisplayContext = WorkflowDisplayContext(
             workflow_display_class=BaseWorkflowDisplay,
-            workflow_display=WorkflowMetaVellumDisplay(
+            workflow_display=WorkflowMetaDisplay(
                 entrypoint_node_id=uuid4(),
                 entrypoint_node_source_handle_id=uuid4(),
                 entrypoint_node_display=NodeDisplayData(),