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

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

vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_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.complex_final_output_node.missing_final_output_node import MissingFinalOutputNodeWorkflow
@@ -11,9 +10,7 @@ from tests.workflows.complex_final_output_node.missing_workflow_output import Mi
 
 def test_serialize_workflow__missing_final_output_node():
     # GIVEN a Workflow that is missing a Terminal Node
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=MissingFinalOutputNodeWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=MissingFinalOutputNodeWorkflow)
 
     # WHEN we serialize it
     serialized_workflow: dict = workflow_display.serialize()
@@ -179,9 +176,7 @@ def test_serialize_workflow__missing_final_output_node():
 
 def test_serialize_workflow__missing_workflow_output():
     # GIVEN a Workflow that contains a terminal node that is unreferenced by the Workflow's Outputs
-    workflow_display = get_workflow_display(
-        base_display_class=VellumWorkflowDisplay, workflow_class=MissingWorkflowOutputWorkflow
-    )
+    workflow_display = get_workflow_display(workflow_class=MissingWorkflowOutputWorkflow)
 
     # WHEN we serialize it, it should throw an error
     with pytest.raises(ValueError) as exc_info:

vellum_ee/workflows/display/types.py

@@ -1,11 +1,12 @@
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Dict, Tuple, Type, TypeVar
+from typing import TYPE_CHECKING, Dict, Tuple, Type
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.events.workflow import WorkflowEventDisplayContext  # noqa: F401
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, StateValueReference, WorkflowInputReference
+from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.base import (
     EdgeDisplay,
     EntrypointDisplay,
@@ -16,11 +17,11 @@ from vellum_ee.workflows.display.base import (
 )
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay
+from vellum_ee.workflows.display.utils.registry import get_default_workflow_display_class
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
 
-WorkflowDisplayType = TypeVar("WorkflowDisplayType", bound="BaseWorkflowDisplay")
 
 WorkflowInputsDisplays = Dict[WorkflowInputReference, WorkflowInputsDisplay]
 StateValueDisplays = Dict[StateValueReference, StateValueDisplay]
@@ -34,8 +35,8 @@ PortDisplays = Dict[Port, PortDisplay]
 
 @dataclass
 class WorkflowDisplayContext:
-    workflow_display_class: Type["BaseWorkflowDisplay"]
-    workflow_display: WorkflowMetaDisplay
+    workflow_display_class: Type["BaseWorkflowDisplay"] = field(default_factory=get_default_workflow_display_class)
+    workflow_display: WorkflowMetaDisplay = field(default_factory=lambda: WorkflowMetaDisplay.get_default(BaseWorkflow))
     workflow_input_displays: WorkflowInputsDisplays = field(default_factory=dict)
     global_workflow_input_displays: WorkflowInputsDisplays = field(default_factory=dict)
     state_value_displays: StateValueDisplays = field(default_factory=dict)

vellum_ee/workflows/display/utils/exceptions.py

@@ -0,0 +1,7 @@
+class UserFacingException(Exception):
+    def to_message(self) -> str:
+        return str(self)
+
+
+class UnsupportedSerializationException(UserFacingException):
+    pass

vellum_ee/workflows/display/utils/registry.py

@@ -0,0 +1,37 @@
+from typing import TYPE_CHECKING, Dict, Optional, Type
+
+from vellum.workflows.nodes import BaseNode
+from vellum.workflows.workflows.base import BaseWorkflow
+
+if TYPE_CHECKING:
+    from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+    from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
+
+
+# Used to store the mapping between workflows and their display classes
+_workflow_display_registry: Dict[Type[BaseWorkflow], Type["BaseWorkflowDisplay"]] = {}
+
+# Used to store the mapping between node types and their display classes
+_node_display_registry: Dict[Type[BaseNode], Type["BaseNodeDisplay"]] = {}
+
+
+def get_from_workflow_display_registry(workflow_class: Type[BaseWorkflow]) -> Optional[Type["BaseWorkflowDisplay"]]:
+    return _workflow_display_registry.get(workflow_class)
+
+
+def register_workflow_display_class(
+    workflow_class: Type[BaseWorkflow], workflow_display_class: Type["BaseWorkflowDisplay"]
+) -> None:
+    _workflow_display_registry[workflow_class] = workflow_display_class
+
+
+def get_default_workflow_display_class() -> Type["BaseWorkflowDisplay"]:
+    return _workflow_display_registry[BaseWorkflow]
+
+
+def get_from_node_display_registry(node_class: Type[BaseNode]) -> Optional[Type["BaseNodeDisplay"]]:
+    return _node_display_registry.get(node_class)
+
+
+def register_node_display_class(node_class: Type[BaseNode], node_display_class: Type["BaseNodeDisplay"]) -> None:
+    _node_display_registry[node_class] = node_display_class

vellum_ee/workflows/display/utils/vellum.py

@@ -41,6 +41,7 @@ from vellum.workflows.references.node import NodeReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 from vellum.workflows.vellum_client import create_vellum_client
+from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
 from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 
 if TYPE_CHECKING:
@@ -164,7 +165,7 @@ def create_node_input_value_pointer_rule(
         vellum_value = primitive_to_vellum_value(value)
         return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_value)
 
-    raise ValueError(f"Unsupported descriptor type: {value.__class__.__name__}")
+    raise UnsupportedSerializationException(f"Unsupported descriptor type: {value.__class__.__name__}")
 
 
 def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperator: