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

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:

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -1,21 +1,25 @@
-from abc import abstractmethod
 from copy import copy
 from functools import cached_property
 import importlib
+import inspect
 import logging
 from uuid import UUID
-from typing import Any, Dict, Generic, Iterator, List, Optional, Tuple, Type, Union, get_args
+from typing import Any, Dict, ForwardRef, Generic, Iterator, List, Optional, Tuple, Type, TypeVar, Union, cast, get_args
 
 from vellum.workflows import BaseWorkflow
+from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.events.workflow import NodeEventDisplayContext, WorkflowEventDisplayContext
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.utils import get_unadorned_node, get_wrapped_node
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
+from vellum.workflows.nodes.displayable.final_output_node.node import FinalOutputNode
+from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, WorkflowInputReference
-from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
+from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.base import (
     EdgeDisplay,
@@ -27,10 +31,10 @@ from vellum_ee.workflows.display.base import (
 )
 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.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
+from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import (
     EdgeDisplays,
     EntrypointDisplays,
@@ -42,6 +46,8 @@ from vellum_ee.workflows.display.types import (
     WorkflowInputsDisplays,
     WorkflowOutputDisplays,
 )
+from vellum_ee.workflows.display.utils.registry import register_workflow_display_class
+from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 logger = logging.getLogger(__name__)
@@ -69,45 +75,257 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
     # Used to explicitly specify display data for a workflow's ports.
     port_displays: PortDisplays = {}
 
-    # Used to store the mapping between workflows and their display classes
-    _workflow_display_registry: Dict[Type[WorkflowType], Type["BaseWorkflowDisplay"]] = {}
-
     _errors: List[Exception]
 
     _dry_run: bool
 
     def __init__(
         self,
-        workflow: Type[WorkflowType],
         *,
         parent_display_context: Optional[WorkflowDisplayContext] = None,
         dry_run: bool = False,
     ):
-        self._workflow = workflow
         self._parent_display_context = parent_display_context
         self._errors: List[Exception] = []
         self._dry_run = dry_run
 
-    @abstractmethod
     def serialize(self) -> JsonObject:
-        pass
+        input_variables: JsonArray = []
+        for workflow_input_reference, workflow_input_display in self.display_context.workflow_input_displays.items():
+            default = (
+                primitive_to_vellum_value(workflow_input_reference.instance)
+                if workflow_input_reference.instance
+                else None
+            )
+            input_variables.append(
+                {
+                    "id": str(workflow_input_display.id),
+                    "key": workflow_input_display.name or workflow_input_reference.name,
+                    "type": infer_vellum_variable_type(workflow_input_reference),
+                    "default": default.dict() if default else None,
+                    "required": workflow_input_reference.instance is undefined,
+                    "extensions": {"color": workflow_input_display.color},
+                }
+            )
 
-    @classmethod
-    def get_from_workflow_display_registry(cls, workflow_class: Type[WorkflowType]) -> Type["BaseWorkflowDisplay"]:
-        try:
-            return cls._workflow_display_registry[workflow_class]
-        except KeyError:
-            return cls._workflow_display_registry[WorkflowType]  # type: ignore [misc]
+        state_variables: JsonArray = []
+        for state_value_reference, state_value_display in self.display_context.state_value_displays.items():
+            default = (
+                primitive_to_vellum_value(state_value_reference.instance) if state_value_reference.instance else None
+            )
+            state_variables.append(
+                {
+                    "id": str(state_value_display.id),
+                    "key": state_value_display.name or state_value_reference.name,
+                    "type": infer_vellum_variable_type(state_value_reference),
+                    "default": default.dict() if default else None,
+                    "required": state_value_reference.instance is undefined,
+                    "extensions": {"color": state_value_display.color},
+                }
+            )
+
+        nodes: JsonArray = []
+        edges: JsonArray = []
+
+        # Add a single synthetic node for the workflow entrypoint
+        entrypoint_node_id = self.display_context.workflow_display.entrypoint_node_id
+        entrypoint_node_source_handle_id = self.display_context.workflow_display.entrypoint_node_source_handle_id
+        nodes.append(
+            {
+                "id": str(entrypoint_node_id),
+                "type": "ENTRYPOINT",
+                "inputs": [],
+                "data": {
+                    "label": "Entrypoint Node",
+                    "source_handle_id": str(entrypoint_node_source_handle_id),
+                },
+                "display_data": self.display_context.workflow_display.entrypoint_node_display.dict(),
+                "base": None,
+                "definition": None,
+            },
+        )
+
+        # Add all the nodes in the workflow
+        for node in self._workflow.get_nodes():
+            node_display = self.display_context.node_displays[node]
+
+            try:
+                serialized_node = node_display.serialize(self.display_context)
+            except NotImplementedError as e:
+                self.add_error(e)
+                continue
+
+            nodes.append(serialized_node)
+
+        # Add all unused nodes in the workflow
+        for node in self._workflow.get_unused_nodes():
+            node_display = self.display_context.node_displays[node]
+
+            try:
+                serialized_node = node_display.serialize(self.display_context)
+            except NotImplementedError as e:
+                self.add_error(e)
+                continue
+
+            nodes.append(serialized_node)
+
+        synthetic_output_edges: JsonArray = []
+        output_variables: JsonArray = []
+        final_output_nodes = [
+            node for node in self.display_context.node_displays.keys() if issubclass(node, FinalOutputNode)
+        ]
+        final_output_node_outputs = {node.Outputs.value for node in final_output_nodes}
+        unreferenced_final_output_node_outputs = final_output_node_outputs.copy()
+        final_output_node_base: JsonObject = {
+            "name": FinalOutputNode.__name__,
+            "module": cast(JsonArray, FinalOutputNode.__module__.split(".")),
+        }
+
+        # Add a synthetic Terminal Node and track the Workflow's output variables for each Workflow output
+        for workflow_output, workflow_output_display in self.display_context.workflow_output_displays.items():
+            final_output_node_id = uuid4_from_hash(f"{self.workflow_id}|node_id|{workflow_output.name}")
+            inferred_type = infer_vellum_variable_type(workflow_output)
+
+            # Remove the terminal node output from the unreferenced set
+            unreferenced_final_output_node_outputs.discard(cast(OutputReference, workflow_output.instance))
+
+            if workflow_output.instance not in final_output_node_outputs:
+                # Create a synthetic terminal node only if there is no terminal node for this output
+                try:
+                    node_input = create_node_input(
+                        final_output_node_id,
+                        "node_input",
+                        # This is currently the wrapper node's output, but we want the wrapped node
+                        workflow_output.instance,
+                        self.display_context,
+                    )
+                except ValueError as e:
+                    raise ValueError(f"Failed to serialize output '{workflow_output.name}': {str(e)}") from e
+
+                source_node_display: Optional[BaseNodeDisplay]
+                first_rule = node_input.value.rules[0]
+                if first_rule.type == "NODE_OUTPUT":
+                    source_node_id = UUID(first_rule.data.node_id)
+                    try:
+                        source_node_display = [
+                            node_display
+                            for node_display in self.display_context.node_displays.values()
+                            if node_display.node_id == source_node_id
+                        ][0]
+                    except IndexError:
+                        source_node_display = None
+
+                synthetic_target_handle_id = str(
+                    uuid4_from_hash(f"{self.workflow_id}|target_handle_id|{workflow_output_display.name}")
+                )
+                synthetic_display_data = NodeDisplayData().dict()
+                synthetic_node_label = "Final Output"
+                nodes.append(
+                    {
+                        "id": str(final_output_node_id),
+                        "type": "TERMINAL",
+                        "data": {
+                            "label": synthetic_node_label,
+                            "name": workflow_output_display.name,
+                            "target_handle_id": synthetic_target_handle_id,
+                            "output_id": str(workflow_output_display.id),
+                            "output_type": inferred_type,
+                            "node_input_id": str(node_input.id),
+                        },
+                        "inputs": [node_input.dict()],
+                        "display_data": synthetic_display_data,
+                        "base": final_output_node_base,
+                        "definition": None,
+                    }
+                )
+
+                if source_node_display:
+                    source_handle_id = source_node_display.get_source_handle_id(
+                        port_displays=self.display_context.port_displays
+                    )
+
+                    synthetic_output_edges.append(
+                        {
+                            "id": str(uuid4_from_hash(f"{self.workflow_id}|edge_id|{workflow_output_display.name}")),
+                            "source_node_id": str(source_node_display.node_id),
+                            "source_handle_id": str(source_handle_id),
+                            "target_node_id": str(final_output_node_id),
+                            "target_handle_id": synthetic_target_handle_id,
+                            "type": "DEFAULT",
+                        }
+                    )
+
+            output_variables.append(
+                {
+                    "id": str(workflow_output_display.id),
+                    "key": workflow_output_display.name,
+                    "type": inferred_type,
+                }
+            )
+
+        # If there are terminal nodes with no workflow output reference,
+        # raise a serialization error
+        if len(unreferenced_final_output_node_outputs) > 0:
+            self.add_error(
+                ValueError("Unable to serialize terminal nodes that are not referenced by workflow outputs.")
+            )
+
+        # Add an edge for each edge in the workflow
+        for target_node, entrypoint_display in self.display_context.entrypoint_displays.items():
+            unadorned_target_node = get_unadorned_node(target_node)
+            target_node_display = self.display_context.node_displays[unadorned_target_node]
+            edges.append(
+                {
+                    "id": str(entrypoint_display.edge_display.id),
+                    "source_node_id": str(entrypoint_node_id),
+                    "source_handle_id": str(entrypoint_node_source_handle_id),
+                    "target_node_id": str(target_node_display.node_id),
+                    "target_handle_id": str(target_node_display.get_trigger_id()),
+                    "type": "DEFAULT",
+                }
+            )
+
+        for (source_node_port, target_node), edge_display in self.display_context.edge_displays.items():
+            unadorned_source_node_port = get_unadorned_port(source_node_port)
+            unadorned_target_node = get_unadorned_node(target_node)
+
+            source_node_port_display = self.display_context.port_displays[unadorned_source_node_port]
+            target_node_display = self.display_context.node_displays[unadorned_target_node]
+
+            edges.append(
+                {
+                    "id": str(edge_display.id),
+                    "source_node_id": str(source_node_port_display.node_id),
+                    "source_handle_id": str(source_node_port_display.id),
+                    "target_node_id": str(target_node_display.node_id),
+                    "target_handle_id": str(
+                        target_node_display.get_target_handle_id_by_source_node_id(source_node_port_display.node_id)
+                    ),
+                    "type": "DEFAULT",
+                }
+            )
+
+        edges.extend(synthetic_output_edges)
+
+        return {
+            "workflow_raw_data": {
+                "nodes": nodes,
+                "edges": edges,
+                "display_data": self.display_context.workflow_display.display_data.dict(),
+                "definition": {
+                    "name": self._workflow.__name__,
+                    "module": cast(JsonArray, self._workflow.__module__.split(".")),
+                },
+            },
+            "input_variables": input_variables,
+            "state_variables": state_variables,
+            "output_variables": output_variables,
+        }
 
     @cached_property
     def workflow_id(self) -> UUID:
         """Can be overridden as a class attribute to specify a custom workflow id."""
-        return uuid4_from_hash(self._workflow.__qualname__)
-
-    @property
-    @abstractmethod
-    def node_display_base_class(self) -> Type[BaseNodeDisplay]:
-        pass
+        return self._workflow.__id__
 
     def add_error(self, error: Exception) -> None:
         if self._dry_run:
@@ -159,13 +377,8 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             port_displays[port] = node_display.get_node_port_display(port)
 
     def _get_node_display(self, node: Type[BaseNode]) -> BaseNodeDisplay:
-        node_display_class = get_node_display_class(self.node_display_base_class, node)
-        node_display = node_display_class()
-
-        if not isinstance(node_display, self.node_display_base_class):
-            raise ValueError(f"{node.__name__} must be a subclass of {self.node_display_base_class.__name__}")
-
-        return node_display
+        node_display_class = get_node_display_class(node)
+        return node_display_class()
 
     @cached_property
     def display_context(self) -> WorkflowDisplayContext:
@@ -294,14 +507,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 display_data=overrides.display_data,
             )
 
-        entrypoint_node_id = uuid4_from_hash(f"{self.workflow_id}|entrypoint_node_id")
-        entrypoint_node_source_handle_id = uuid4_from_hash(f"{self.workflow_id}|entrypoint_node_source_handle_id")
-
-        return WorkflowMetaDisplay(
-            entrypoint_node_id=entrypoint_node_id,
-            entrypoint_node_source_handle_id=entrypoint_node_source_handle_id,
-            entrypoint_node_display=NodeDisplayData(),
-        )
+        return WorkflowMetaDisplay.get_default(self._workflow)
 
     def _generate_workflow_input_display(
         self, workflow_input: WorkflowInputReference, overrides: Optional[WorkflowInputsDisplay] = None
@@ -368,11 +574,13 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         super().__init_subclass__(**kwargs)
 
         workflow_class = get_args(cls.__orig_bases__[0])[0]  # type: ignore [attr-defined]
-        cls._workflow_display_registry[workflow_class] = cls
+        register_workflow_display_class(workflow_class=workflow_class, workflow_display_class=cls)
 
     @staticmethod
     def gather_event_display_context(
-        module_path: str, workflow_class: Type[BaseWorkflow]
+        module_path: str,
+        # DEPRECATED: This will be removed in the 0.15.0 release
+        workflow_class: Optional[Type[BaseWorkflow]] = None,
     ) -> Union[WorkflowEventDisplayContext, None]:
         workflow_display_module = f"{module_path}.display.workflow"
         try:
@@ -380,11 +588,11 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         except ModuleNotFoundError:
             return None
 
-        workflow_display = display_class.WorkflowDisplay(workflow_class)
-        if not isinstance(workflow_display, BaseWorkflowDisplay):
+        WorkflowDisplayClass = display_class.WorkflowDisplay
+        if not isinstance(WorkflowDisplayClass, type) or not issubclass(WorkflowDisplayClass, BaseWorkflowDisplay):
             return None
 
-        return workflow_display.get_event_display_context()
+        return WorkflowDisplayClass().get_event_display_context()
 
     def get_event_display_context(self):
         display_context = self.display_context
@@ -403,9 +611,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         node_event_displays = {}
         for node_id in node_displays:
             node, current_node_display = node_displays[node_id]
-            input_display = {}
-            if isinstance(current_node_display, BaseNodeVellumDisplay):
-                input_display = current_node_display.node_input_ids_by_name
+            input_display = current_node_display.node_input_ids_by_name
             output_display = {
                 output.name: current_node_display.output_display[output].id
                 for output in current_node_display.output_display
@@ -494,3 +700,27 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         return EdgeDisplay(
             id=uuid4_from_hash(f"{self.workflow_id}|id|{source_node_id}|{target_node_id}"),
         )
+
+    @classmethod
+    def infer_workflow_class(cls) -> Type[BaseWorkflow]:
+        original_base = get_original_base(cls)
+        workflow_class = get_args(original_base)[0]
+        if isinstance(workflow_class, TypeVar):
+            bounded_class = workflow_class.__bound__
+            if inspect.isclass(bounded_class) and issubclass(bounded_class, BaseWorkflow):
+                return bounded_class
+
+            if isinstance(bounded_class, ForwardRef) and bounded_class.__forward_arg__ == BaseWorkflow.__name__:
+                return BaseWorkflow
+
+        if issubclass(workflow_class, BaseWorkflow):
+            return workflow_class
+
+        raise ValueError(f"Workflow {cls.__name__} must be a subclass of {BaseWorkflow.__name__}")
+
+    @property
+    def _workflow(self) -> Type[WorkflowType]:
+        return cast(Type[WorkflowType], self.__class__.infer_workflow_class())
+
+
+register_workflow_display_class(workflow_class=BaseWorkflow, workflow_display_class=BaseWorkflowDisplay)

vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py

@@ -1,33 +1,46 @@
-from typing import Optional, Type
+import types
+from typing import TYPE_CHECKING, Generic, Optional, Type, TypeVar
 
 from vellum.workflows.types.generics import WorkflowType
-from vellum_ee.workflows.display.types import WorkflowDisplayContext, WorkflowDisplayType
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.registry import get_from_workflow_display_registry
+
+if TYPE_CHECKING:
+    from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
+
+
+def _get_workflow_display_class(*, workflow_class: Type[WorkflowType]) -> Type["BaseWorkflowDisplay"]:
+    workflow_display_class = get_from_workflow_display_registry(workflow_class)
+    if workflow_display_class:
+        return workflow_display_class
+
+    base_workflow_display_class = _get_workflow_display_class(
+        workflow_class=workflow_class.__bases__[0],
+    )
+
+    # mypy gets upset at dynamic TypeVar's, but it's technically allowed by python
+    _WorkflowClassType = TypeVar(f"_{workflow_class.__name__}Type", bound=workflow_class)  # type: ignore[misc]
+    # `base_workflow_display_class` is always a Generic class, so it's safe to index into it
+    WorkflowDisplayBaseClass = base_workflow_display_class[_WorkflowClassType]  # type: ignore[index]
+
+    WorkflowDisplayClass = types.new_class(
+        f"{workflow_class.__name__}Display",
+        bases=(WorkflowDisplayBaseClass, Generic[_WorkflowClassType]),
+    )
+
+    return WorkflowDisplayClass
 
 
 def get_workflow_display(
     *,
-    base_display_class: Type[WorkflowDisplayType],
     workflow_class: Type[WorkflowType],
-    root_workflow_class: Optional[Type[WorkflowType]] = None,
     parent_display_context: Optional[WorkflowDisplayContext] = None,
     dry_run: bool = False,
-) -> WorkflowDisplayType:
-    try:
-        workflow_display_class = base_display_class.get_from_workflow_display_registry(workflow_class)
-    except KeyError:
-        try:
-            return get_workflow_display(
-                base_display_class=base_display_class,
-                workflow_class=workflow_class.__bases__[0],
-                root_workflow_class=workflow_class if root_workflow_class is None else root_workflow_class,
-                parent_display_context=parent_display_context,
-                dry_run=dry_run,
-            )
-        except IndexError:
-            return base_display_class(workflow_class)
-
-    return workflow_display_class(  # type: ignore[return-value]
-        workflow_class,
+    # DEPRECATED: The following arguments will be removed in 0.15.0
+    root_workflow_class: Optional[Type[WorkflowType]] = None,
+    base_display_class: Optional[Type["BaseWorkflowDisplay"]] = None,
+) -> "BaseWorkflowDisplay":
+    return _get_workflow_display_class(workflow_class=workflow_class)(
         parent_display_context=parent_display_context,
         dry_run=dry_run,
     )