vellum-ai 0.14.39__py3-none-any.whl → 0.14.40__py3-none-any.whl

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,33 +1,27 @@
 import types
 from uuid import UUID
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, Generic, Type, TypeVar
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.utils.registry import get_from_node_display_registry
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 
 
-def get_node_display_class(
-    base_class: Type["BaseNodeDisplay"], node_class: Type[NodeType], root_node_class: Optional[Type[NodeType]] = None
-) -> Type["BaseNodeDisplay"]:
-    node_display_class = base_class.get_from_node_display_registry(node_class)
+def get_node_display_class(node_class: Type[NodeType]) -> Type["BaseNodeDisplay"]:
+    node_display_class = get_from_node_display_registry(node_class)
     if node_display_class:
-        if not issubclass(node_display_class, base_class):
-            raise TypeError(
-                f"Expected to find a subclass of '{base_class.__name__}' for node class '{node_class.__name__}'"
-            )
-
         return node_display_class
 
-    base_node_display_class = get_node_display_class(
-        base_class, node_class.__bases__[0], node_class if root_node_class is None else root_node_class
-    )
+    base_node_display_class = get_node_display_class(node_class.__bases__[0])
 
+    # mypy gets upset at dynamic TypeVar's, but it's technically allowed by python
+    _NodeClassType = TypeVar(f"_{node_class.__name__}Type", bound=node_class)  # type: ignore[misc]
     # `base_node_display_class` is always a Generic class, so it's safe to index into it
-    NodeDisplayBaseClass = base_node_display_class[node_class]  # type: ignore[index]
+    NodeDisplayBaseClass = base_node_display_class[_NodeClassType]  # type: ignore[index]
 
     def _get_node_input_ids_by_ref(path: str, inst: Any):
         if isinstance(inst, dict):
@@ -51,7 +45,7 @@ def get_node_display_class(
 
     NodeDisplayClass = types.new_class(
         f"{node_class.__name__}Display",
-        bases=(NodeDisplayBaseClass,),
+        bases=(NodeDisplayBaseClass, Generic[_NodeClassType]),
         exec_body=exec_body,
     )
 

vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -2,7 +2,11 @@ import pytest
 from uuid import UUID
 
 from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.ports.port import Port
+from vellum.workflows.references.constant import ConstantValueReference
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 
 @pytest.fixture
@@ -46,3 +50,35 @@ def test_get_id(node_info):
 
     assert node_display().node_id == expected_id
     assert node_display.infer_node_class().__id__ == expected_id
+
+
+def test_serialize_condition__accessor_expression():
+    # GIVEN a node with an accessor expression in a Port
+    class MyNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            foo = Port.on_if(ConstantValueReference({"hello": "world"})["hello"])
+
+    # WHEN we serialize the node
+    node_display_class = get_node_display_class(MyNode)
+    data = node_display_class().serialize(WorkflowDisplayContext())
+
+    # THEN the condition should be serialized correctly
+    assert data["ports"] == [
+        {
+            "id": "7de6ea94-7f6c-475e-8f38-ec8ac317fd19",
+            "name": "foo",
+            "type": "IF",
+            "expression": {
+                "type": "BINARY_EXPRESSION",
+                "lhs": {
+                    "type": "CONSTANT_VALUE",
+                    "value": {
+                        "type": "JSON",
+                        "value": {"hello": "world"},
+                    },
+                },
+                "operator": "accessField",
+                "rhs": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "hello"}},
+            },
+        }
+    ]

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

@@ -14,6 +14,7 @@ from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeV
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.registry import register_node_display_class
 
 _BaseAdornmentNodeType = TypeVar("_BaseAdornmentNodeType", bound=BaseAdornmentNode)
 
@@ -31,7 +32,7 @@ def _recursively_replace_wrapped_node(node_class: Type[BaseNode], wrapped_node_d
     # 1. The node display class' parameterized type
     original_base_node_display = get_original_base(wrapped_node_display_class)
     original_base_node_display.__args__ = (wrapped_node_class,)
-    wrapped_node_display_class._node_display_registry[wrapped_node_class] = wrapped_node_display_class
+    register_node_display_class(node_class=wrapped_node_class, node_display_class=wrapped_node_display_class)
     wrapped_node_display_class.__annotate_node__()
 
     # 2. The node display class' output displays
@@ -89,7 +90,7 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
                 "Unable to serialize standalone adornment nodes. Please use adornment nodes as a decorator."
             )
 
-        wrapped_node_display_class = get_node_display_class(BaseNodeDisplay, wrapped_node)
+        wrapped_node_display_class = get_node_display_class(wrapped_node)
         wrapped_node_display = wrapped_node_display_class()
         additional_kwargs = get_additional_kwargs(wrapped_node_display.node_id) if get_additional_kwargs else {}
         serialized_wrapped_node = wrapped_node_display.serialize(display_context, **kwargs, **additional_kwargs)

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

@@ -9,7 +9,6 @@ from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.workflows.base import BaseWorkflow
-from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
@@ -72,7 +71,7 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
         if not inner_node:
             return super().get_node_output_display(output)
 
-        node_display_class = get_node_display_class(BaseNodeDisplay, inner_node)
+        node_display_class = get_node_display_class(inner_node)
         node_display = node_display_class()
 
         inner_output = getattr(inner_node.Outputs, output.name)

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

@@ -6,7 +6,6 @@ from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExec
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.vellum.code_execution_node import BaseCodeExecutionNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 def _no_display_class(Node: Type[CodeExecutionNode]):
@@ -53,7 +52,7 @@ def test_serialize_node__code_node_inputs(GetDisplayClass, expected_input_id):
     GetDisplayClass(MyCodeExecutionNode)
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the inputs

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

@@ -4,7 +4,6 @@ from vellum.client.types.vellum_error import VellumError
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes.core.error_node.node import ErrorNode
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 def test_error_node_display__serialize_with_vellum_error() -> None:
@@ -20,7 +19,7 @@ def test_error_node_display__serialize_with_vellum_error() -> None:
         graph = MyNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
     serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
 
     # THEN the correct inputs should be serialized on the node

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

@@ -2,7 +2,6 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes.displayable.note_node.node import NoteNode
 from vellum_ee.workflows.display.nodes.vellum.note_node import BaseNoteNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 def test_serialize_node__note_node():
@@ -22,7 +21,7 @@ def test_serialize_node__note_node():
         graph = MyNoteNode
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the inputs

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

@@ -6,9 +6,9 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
 from vellum.workflows.references.lazy import LazyReference
+from vellum.workflows.state.base import BaseState
 from vellum_ee.workflows.display.nodes.vellum.inline_prompt_node import BaseInlinePromptNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 def test_serialize_node__lazy_reference_in_prompt_inputs():
@@ -27,7 +27,7 @@ def test_serialize_node__lazy_reference_in_prompt_inputs():
         graph = LazyReferencePromptNode >> OtherNode
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the attribute reference
@@ -103,7 +103,7 @@ def test_serialize_node__prompt_inputs(GetDisplayClass, expected_input_id):
     GetDisplayClass(MyPromptNode)
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the inputs
@@ -129,3 +129,55 @@ def test_serialize_node__prompt_inputs(GetDisplayClass, expected_input_id):
             },
         }
     ]
+
+
+def test_serialize_node__prompt_inputs__state_reference():
+    # GIVEN a state definition
+    class MyState(BaseState):
+        foo: str
+
+    # AND a prompt node with inputs
+    class MyPromptNode(InlinePromptNode):
+        prompt_inputs = {"foo": MyState.foo, "bar": "baz"}
+        blocks = []
+        ml_model = "gpt-4o"
+
+    # AND a workflow with the prompt node
+    class Workflow(BaseWorkflow):
+        graph = MyPromptNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should skip the state reference input rule
+    my_prompt_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
+    )
+
+    assert my_prompt_node["inputs"] == [
+        {
+            "id": "e47e0a80-afbb-4888-b06b-8dc78edd8572",
+            "key": "foo",
+            "value": {
+                "rules": [],
+                "combinator": "OR",
+            },
+        },
+        {
+            "id": "3750feb9-5d5c-4150-b62d-a9924f466888",
+            "key": "bar",
+            "value": {
+                "rules": [
+                    {
+                        "type": "CONSTANT_VALUE",
+                        "data": {
+                            "type": "STRING",
+                            "value": "baz",
+                        },
+                    }
+                ],
+                "combinator": "OR",
+            },
+        },
+    ]

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

@@ -5,7 +5,6 @@ from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 def test_retry_node_parameters():
@@ -21,7 +20,7 @@ def test_retry_node_parameters():
         graph = MyRetryNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
     serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
 
     # THEN the correct inputs should be serialized on the node

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

@@ -6,7 +6,6 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum_ee.workflows.display.nodes.vellum.templating_node import BaseTemplatingNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 def _no_display_class(Node: Type[TemplatingNode]):
@@ -53,7 +52,7 @@ def test_serialize_node__templating_node_inputs(GetDisplayClass, expected_input_
     GetDisplayClass(MyTemplatingNode)
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the inputs

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

@@ -5,7 +5,6 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 def test_try_node_display__serialize_with_error_output() -> None:
@@ -28,7 +27,7 @@ def test_try_node_display__serialize_with_error_output() -> None:
         graph = MyNode >> OtherNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
     serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
 
     # THEN the correct inputs should be serialized on the node

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

@@ -22,7 +22,7 @@ from vellum_ee.workflows.display.utils.vellum import (
     NodeOutputPointer,
 )
 from vellum_ee.workflows.display.vellum import WorkflowInputsVellumDisplayOverrides, WorkflowMetaVellumDisplay
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
 class Inputs(BaseInputs):
@@ -109,7 +109,7 @@ def test_create_node_input_value_pointer_rules(
     rules = create_node_input_value_pointer_rules(
         descriptor,
         WorkflowDisplayContext(
-            workflow_display_class=VellumWorkflowDisplay,
+            workflow_display_class=BaseWorkflowDisplay,
             workflow_display=WorkflowMetaVellumDisplay(
                 entrypoint_node_id=uuid4(),
                 entrypoint_node_source_handle_id=uuid4(),

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

@@ -10,7 +10,6 @@ from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.workflows.base import BaseWorkflow
-from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
@@ -82,7 +81,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
         if not inner_node:
             return super().get_node_output_display(output)
 
-        node_display_class = get_node_display_class(BaseNodeDisplay, inner_node)
+        node_display_class = get_node_display_class(inner_node)
         node_display = node_display_class()
         if output.name == "error":
             return inner_node, NodeOutputDisplay(

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

@@ -8,6 +8,7 @@ from vellum.workflows.references import NodeReference
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
 from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 from vellum_ee.workflows.display.utils.vellum import (
     ConstantValuePointer,
@@ -76,7 +77,12 @@ def create_node_input_value_pointer_rules(
                 node_input_value_pointer_rules.extend(rhs_rules)
         else:
             # Non-CoalesceExpression case
-            node_input_value_pointer_rules.append(create_node_input_value_pointer_rule(value, display_context))
+            try:
+                rule = create_node_input_value_pointer_rule(value, display_context)
+            except UnsupportedSerializationException:
+                return node_input_value_pointer_rules
+
+            node_input_value_pointer_rules.append(rule)
     else:
         pointer = create_pointer(value, pointer_type)
         node_input_value_pointer_rules.append(pointer)

vellum_ee/workflows/display/tests/{test_vellum_workflow_display.py → test_base_workflow_display.py}

@@ -6,19 +6,16 @@ from vellum.workflows.nodes import BaseNode
 from vellum.workflows.state import BaseState
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.vellum import WorkflowInputsVellumDisplayOverrides
+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
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
-def test_vellum_workflow_display__serialize_empty_workflow():
+def test_base_workflow_display__serialize_empty_workflow():
     # GIVEN an empty workflow
     class ExampleWorkflow(BaseWorkflow):
         pass
 
-    display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=ExampleWorkflow,
-    )
+    display = get_workflow_display(workflow_class=ExampleWorkflow)
 
     # WHEN serializing the workflow
     exec_config = display.serialize()
@@ -30,18 +27,18 @@ def test_vellum_workflow_display__serialize_empty_workflow():
         "output_variables": [],
         "workflow_raw_data": {
             "definition": {
-                "module": ["vellum_ee", "workflows", "display", "tests", "test_vellum_workflow_display"],
+                "module": ["vellum_ee", "workflows", "display", "tests", "test_base_workflow_display"],
                 "name": "ExampleWorkflow",
             },
             "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
             "edges": [],
             "nodes": [
                 {
-                    "data": {"label": "Entrypoint Node", "source_handle_id": "508b8b82-3517-4672-a155-18c9c7b9c545"},
+                    "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}},
-                    "id": "9eef0c18-f322-4d56-aa89-f088d3e53f6a",
+                    "id": "3c41cdd9-999a-48b8-9088-f6dfa1369bfd",
                     "inputs": [],
                     "type": "ENTRYPOINT",
                 }
@@ -62,7 +59,7 @@ def test_vellum_workflow_display__serialize_input_variables_with_capitalized_var
     class ExampleWorkflow(BaseWorkflow[Inputs, BaseState]):
         graph = StartNode
 
-    class ExampleWorkflowDisplay(VellumWorkflowDisplay[ExampleWorkflow]):
+    class ExampleWorkflowDisplay(BaseWorkflowDisplay[ExampleWorkflow]):
         inputs_display = {
             Inputs.foo: WorkflowInputsVellumDisplayOverrides(
                 id=UUID("97b63d71-5413-417f-9cf5-49e1b4fd56e4"), name="Foo", required=True
@@ -108,10 +105,7 @@ def test_vellum_workflow_display_serialize_valid_handle_ids_for_base_nodes():
             final_value = EndNode.Outputs.hello
 
     # AND a display class for this workflow
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=Workflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=Workflow)
 
     # WHEN we serialize the workflow
     exec_config = workflow_display.serialize()
@@ -169,10 +163,7 @@ def test_vellum_workflow_display__serialize_with_unused_nodes_and_edges():
             final = NodeA.Outputs.result
 
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=Workflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=Workflow)
 
     # WHEN we serialize the workflow
     exec_config = workflow_display.serialize()
@@ -245,10 +236,7 @@ def test_vellum_workflow_display__serialize_with_parse_json_expression():
             final = JsonNode.Outputs.json_result
 
     # AND a display class for this workflow
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=Workflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=Workflow)
 
     # WHEN we serialize the workflow
     exec_config = workflow_display.serialize()

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

@@ -1,11 +1,10 @@
 import pytest
 from uuid import uuid4
-from typing import Any, Type
+from typing import Type
 
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
 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.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.types import (
     NodeDisplays,
@@ -15,24 +14,23 @@ from vellum_ee.workflows.display.types import (
     WorkflowInputsDisplays,
 )
 from vellum_ee.workflows.display.vellum import WorkflowMetaVellumDisplay
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
 @pytest.fixture()
 def serialize_node():
     def _serialize_node(
         node_class: Type[NodeType],
-        base_class: type[BaseNodeDisplay[Any]] = BaseNodeDisplay,
         global_workflow_input_displays: WorkflowInputsDisplays = {},
         global_state_value_displays: StateValueDisplays = {},
         global_node_displays: NodeDisplays = {},
         global_node_output_displays: NodeOutputDisplays = {},
     ) -> JsonObject:
-        node_display_class = get_node_display_class(base_class, node_class)
+        node_display_class = get_node_display_class(node_class)
         node_display = node_display_class()
 
         context: WorkflowDisplayContext = WorkflowDisplayContext(
-            workflow_display_class=VellumWorkflowDisplay,
+            workflow_display_class=BaseWorkflowDisplay,
             workflow_display=WorkflowMetaVellumDisplay(
                 entrypoint_node_id=uuid4(),
                 entrypoint_node_source_handle_id=uuid4(),

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

@@ -10,11 +10,9 @@ from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.base import WorkflowInputsDisplay
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.vellum.retry_node import BaseRetryNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.try_node import BaseTryNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 class Inputs(BaseInputs):
@@ -124,10 +122,7 @@ def test_serialize_node__retry__no_display():
         graph = StartNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=MyWorkflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
     exec_config = workflow_display.serialize()
 
     # THEN the workflow display is created successfully
@@ -148,7 +143,6 @@ def test_serialize_node__try(serialize_node):
 
     input_id = uuid4()
     serialized_node = serialize_node(
-        base_class=BaseNodeVellumDisplay,
         node_class=InnerTryGenericNode,
         global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)},
         global_node_displays={
@@ -220,10 +214,7 @@ def test_serialize_node__try__no_display():
         graph = StartNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=MyWorkflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
 
     exec_config = workflow_display.serialize()
 
@@ -242,10 +233,7 @@ def test_serialize_node__stacked():
         graph = InnerStackedGenericNode
 
     # WHEN we serialize the workflow
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=StackedWorkflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=StackedWorkflow)
     exec_config = workflow_display.serialize()
 
     # THEN the workflow display is created successfully
@@ -283,7 +271,10 @@ def test_serialize_node__stacked():
                                 "test_adornments_serialization",
                             ],
                         },
-                        "trigger": {"id": "6e4af17f-bbee-4777-b10d-af042cd6e16a", "merge_behavior": "AWAIT_ATTRIBUTES"},
+                        "trigger": {
+                            "id": "6e4af17f-bbee-4777-b10d-af042cd6e16a",
+                            "merge_behavior": "AWAIT_ATTRIBUTES",
+                        },
                         "ports": [{"id": "408cd5fb-3a3e-4eb2-9889-61111bd6a129", "name": "default", "type": "DEFAULT"}],
                         "adornments": [
                             {

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

@@ -16,7 +16,6 @@ from vellum_ee.workflows.display.base import WorkflowInputsDisplay
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 class Inputs(BaseInputs):
@@ -169,7 +168,7 @@ def test_serialize_node__lazy_reference_with_string():
         graph = LazyReferenceGenericNode >> OtherNode
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN the node should properly serialize the attribute reference
@@ -257,10 +256,7 @@ def test_serialize_node__workflow_input_as_nested_chat_history():
         graph = GenericNode
 
     # WHEN the workflow is serialized
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay,
-        workflow_class=Workflow,
-    )
+    workflow_display = get_workflow_display(workflow_class=Workflow)
     with pytest.raises(Exception) as exc_info:
         workflow_display.serialize()
 

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

@@ -4,7 +4,6 @@ from uuid import uuid4
 from deepdiff import DeepDiff
 
 from vellum import WorkspaceSecretRead
-from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 from tests.workflows.basic_api_node.workflow import SimpleAPIWorkflow
@@ -24,7 +23,7 @@ def test_serialize_workflow(vellum_client):
     vellum_client.workspace_secrets.retrieve.return_value = workspace_secret
 
     # WHEN we serialize it
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleAPIWorkflow)
+    workflow_display = get_workflow_display(workflow_class=SimpleAPIWorkflow)
 
     serialized_workflow: dict = workflow_display.serialize()