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

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()
 

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

@@ -1,7 +1,6 @@
 from deepdiff import DeepDiff
 
 from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
-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_code_execution_node.try_workflow import TrySimpleCodeExecutionWorkflow
@@ -12,9 +11,7 @@ from tests.workflows.basic_code_execution_node.workflow_with_code import SimpleC
 def test_serialize_workflow_with_filepath():
     # GIVEN a Workflow with a code execution node
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=SimpleCodeExecutionWithFilepathWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=SimpleCodeExecutionWithFilepathWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
@@ -253,9 +250,7 @@ def test_serialize_workflow_with_filepath():
 def test_serialize_workflow_with_code():
     # GIVEN a Workflow with a code execution node
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=SimpleCodeExecutionWithCodeWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=SimpleCodeExecutionWithCodeWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
@@ -482,9 +477,7 @@ def test_serialize_workflow_with_code():
 def test_serialize_workflow__try_wrapped():
     # GIVEN a Workflow with a code execution node
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=TrySimpleCodeExecutionWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=TrySimpleCodeExecutionWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow

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

@@ -20,7 +20,6 @@ from vellum.workflows.expressions.less_than import LessThanExpression
 from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
 from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.expressions.not_in import NotInExpression
-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_conditional_node.workflow import CategoryWorkflow
@@ -30,7 +29,7 @@ from tests.workflows.basic_conditional_node.workflow_with_only_one_conditional_n
 def test_serialize_workflow():
     # GIVEN a Workflow that uses a ConditionalNode
     # WHEN we serialize it
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=CategoryWorkflow)
+    workflow_display = get_workflow_display(workflow_class=CategoryWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
@@ -779,7 +778,7 @@ def test_conditional_node_serialize_all_operators_with_lhs_and_rhs(descriptor, o
     # GIVEN a simple workflow with one conditional node
     workflow_cls = create_simple_workflow(descriptor)
 
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=workflow_cls)
+    workflow_display = get_workflow_display(workflow_class=workflow_cls)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
@@ -885,7 +884,7 @@ def test_conditional_node_serialize_all_operators_with_expression(descriptor, op
     # GIVEN a simple workflow with one conditional node
     workflow_cls = create_simple_workflow(descriptor)
 
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=workflow_cls)
+    workflow_display = get_workflow_display(workflow_class=workflow_cls)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
@@ -978,7 +977,7 @@ def test_conditional_node_serialize_all_operators_with_value_and_start_and_end(d
     # GIVEN a simple workflow with one conditional node
     workflow_cls = create_simple_workflow(descriptor)
 
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=workflow_cls)
+    workflow_display = get_workflow_display(workflow_class=workflow_cls)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow

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

@@ -1,6 +1,5 @@
 from deepdiff import DeepDiff
 
-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_default_state.workflow import BasicDefaultStateWorkflow
@@ -9,9 +8,7 @@ from tests.workflows.basic_default_state.workflow import BasicDefaultStateWorkfl
 def test_serialize_workflow():
     # GIVEN a Workflow that has a simple state definition
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicDefaultStateWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicDefaultStateWorkflow)
 
     serialized_workflow: dict = workflow_display.serialize()
     # THEN we should get a serialized representation of the Workflow

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

@@ -1,6 +1,5 @@
 from deepdiff import DeepDiff
 
-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_error_node.workflow import BasicErrorNodeWorkflow
@@ -9,9 +8,7 @@ from tests.workflows.basic_error_node.workflow import BasicErrorNodeWorkflow
 def test_serialize_workflow():
     # GIVEN a Workflow with an error node
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicErrorNodeWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicErrorNodeWorkflow)
 
     serialized_workflow: dict = workflow_display.serialize()
 

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

@@ -1,17 +1,14 @@
 from deepdiff import DeepDiff
 
-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_generic_node.workflow import BasicGenericNodeWorkflow
 
 
-def test_serialize_workflow(vellum_client):
+def test_serialize_workflow():
     # GIVEN a Workflow that uses a generic node
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicGenericNodeWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicGenericNodeWorkflow)
 
     serialized_workflow: dict = workflow_display.serialize()
 

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

@@ -1,6 +1,5 @@
 from deepdiff import DeepDiff
 
-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_guardrail_node.workflow import BasicGuardrailNodeWorkflow
@@ -9,14 +8,17 @@ from tests.workflows.basic_guardrail_node.workflow import BasicGuardrailNodeWork
 def test_serialize_workflow():
     # GIVEN a workflow that uses a guardrail node
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicGuardrailNodeWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicGuardrailNodeWorkflow)
 
     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"}
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+    }
 
     # AND its input variables should be what we expect
     input_variables = serialized_workflow["input_variables"]

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

@@ -1,4 +1,3 @@
-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_inline_prompt_node_with_functions.workflow import BasicInlinePromptWithFunctionsWorkflow
@@ -6,9 +5,7 @@ from tests.workflows.basic_inline_prompt_node_with_functions.workflow import Bas
 
 def test_serialize_workflow():
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicInlinePromptWithFunctionsWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicInlinePromptWithFunctionsWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
     assert (
         serialized_workflow["workflow_raw_data"]["nodes"][-2]["data"]["exec_config"]["prompt_template_block_data"][

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

@@ -1,6 +1,5 @@
 from deepdiff import DeepDiff
 
-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_inline_subworkflow.workflow import BasicInlineSubworkflowWorkflow
@@ -9,9 +8,7 @@ from tests.workflows.basic_inline_subworkflow.workflow import BasicInlineSubwork
 def test_serialize_workflow():
     # GIVEN a Workflow
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicInlineSubworkflowWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicInlineSubworkflowWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow

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

@@ -1,6 +1,5 @@
 from deepdiff import DeepDiff
 
-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_map_node.workflow import SimpleMapExample
@@ -9,7 +8,7 @@ from tests.workflows.basic_map_node.workflow import SimpleMapExample
 def test_serialize_workflow():
     # GIVEN a Workflow that uses a MapNode
     # WHEN we serialize it
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleMapExample)
+    workflow_display = get_workflow_display(workflow_class=SimpleMapExample)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow

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

@@ -1,6 +1,5 @@
 from deepdiff import DeepDiff
 
-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_merge_node.await_all_workflow import AwaitAllPassingWorkflow
@@ -9,9 +8,7 @@ from tests.workflows.basic_merge_node.await_all_workflow import AwaitAllPassingW
 def test_serialize_workflow__await_all():
     # GIVEN a Workflow that uses an await all merge node
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=AwaitAllPassingWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=AwaitAllPassingWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow

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

@@ -4,7 +4,6 @@ from uuid import uuid4
 from deepdiff import DeepDiff
 
 from vellum import DeploymentRead
-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_text_prompt_deployment.workflow import BasicTextPromptDeployment
@@ -25,9 +24,7 @@ def test_serialize_workflow(vellum_client):
     vellum_client.deployments.retrieve.return_value = deployment
 
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicTextPromptDeployment
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicTextPromptDeployment)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow

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

@@ -1,4 +1,3 @@
-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_search_node.workflow import BasicSearchWorkflow
@@ -8,14 +7,17 @@ def test_serialize_workflow():
     # GIVEN a Workflow with a search node
     # WHEN we serialize it
 
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicSearchWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicSearchWorkflow)
 
     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"}
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+    }
 
     # AND its input variables should be what we expect
     input_variables = serialized_workflow["input_variables"]

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

@@ -4,7 +4,6 @@ from uuid import uuid4
 from deepdiff import DeepDiff
 
 from vellum import WorkflowDeploymentRead
-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_subworkflow_deployment.workflow import BasicSubworkflowDeploymentWorkflow
@@ -25,9 +24,7 @@ def test_serialize_workflow(vellum_client):
     vellum_client.workflow_deployments.retrieve.return_value = deployment
 
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicSubworkflowDeploymentWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicSubworkflowDeploymentWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow

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

@@ -1,6 +1,5 @@
 from deepdiff import DeepDiff
 
-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_templating_node.workflow_with_json_input import BasicTemplatingNodeWorkflowWithJson
@@ -10,9 +9,7 @@ def test_serialize_workflow():
     # GIVEN a Workflow that uses a vellum templating node",
 
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicTemplatingNodeWorkflowWithJson
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicTemplatingNodeWorkflowWithJson)
 
     serialized_workflow: dict = workflow_display.serialize()
 

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

@@ -1,4 +1,3 @@
-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_final_output_node.workflow import BasicFinalOutputNodeWorkflow
@@ -7,9 +6,7 @@ from tests.workflows.basic_final_output_node.workflow import BasicFinalOutputNod
 def test_serialize_workflow():
     # GIVEN a Workflow that uses a Final Output Node
     # WHEN we serialize it
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=BasicFinalOutputNodeWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=BasicFinalOutputNodeWorkflow)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow

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

@@ -2,7 +2,6 @@ import pytest
 
 from deepdiff import DeepDiff
 
-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_try_node.workflow import SimpleTryExample, StandaloneTryExample
@@ -11,7 +10,7 @@ from tests.workflows.basic_try_node.workflow import SimpleTryExample, Standalone
 def test_serialize_workflow():
     # GIVEN a Workflow with a TryNode
     # WHEN we serialize it
-    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleTryExample)
+    workflow_display = get_workflow_display(workflow_class=SimpleTryExample)
     serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
@@ -77,9 +76,7 @@ def test_serialize_workflow__standalone():
     # GIVEN a Workflow with a standalone TryNode
     # WHEN we serialize it
     with pytest.raises(NotImplementedError) as exc:
-        workflow_display = get_workflow_display(
-            base_display_class=VellumWorkflowDisplay, workflow_class=StandaloneTryExample
-        )
+        workflow_display = get_workflow_display(workflow_class=StandaloneTryExample)
         workflow_display.serialize()
 
     # THEN we should get an error