vellum-ai 0.14.44__py3-none-any.whl → 0.14.46__py3-none-any.whl

vellum_ee/workflows/display/utils/expressions.py

@@ -1,12 +1,98 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+from vellum.client.types.logical_operator import LogicalOperator
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.expressions.accessor import AccessorExpression
+from vellum.workflows.expressions.and_ import AndExpression
+from vellum.workflows.expressions.begins_with import BeginsWithExpression
+from vellum.workflows.expressions.between import BetweenExpression
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
+from vellum.workflows.expressions.contains import ContainsExpression
+from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
+from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
+from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
+from vellum.workflows.expressions.does_not_equal import DoesNotEqualExpression
+from vellum.workflows.expressions.ends_with import EndsWithExpression
+from vellum.workflows.expressions.equals import EqualsExpression
+from vellum.workflows.expressions.greater_than import GreaterThanExpression
+from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
+from vellum.workflows.expressions.in_ import InExpression
+from vellum.workflows.expressions.is_nil import IsNilExpression
+from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
+from vellum.workflows.expressions.is_not_null import IsNotNullExpression
+from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
+from vellum.workflows.expressions.is_null import IsNullExpression
+from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
+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.workflows.expressions.or_ import OrExpression
+from vellum.workflows.expressions.parse_json import ParseJsonExpression
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
+from vellum.workflows.references.constant import ConstantValueReference
+from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.lazy import LazyReference
+from vellum.workflows.references.output import OutputReference
+from vellum.workflows.references.state_value import StateValueReference
+from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.references.workflow_input import WorkflowInputReference
+from vellum.workflows.types.core import JsonObject
+from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 
+def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperator:
+    if isinstance(descriptor, EqualsExpression):
+        return "="
+    elif isinstance(descriptor, DoesNotEqualExpression):
+        return "!="
+    elif isinstance(descriptor, LessThanExpression):
+        return "<"
+    elif isinstance(descriptor, GreaterThanExpression):
+        return ">"
+    elif isinstance(descriptor, LessThanOrEqualToExpression):
+        return "<="
+    elif isinstance(descriptor, GreaterThanOrEqualToExpression):
+        return ">="
+    elif isinstance(descriptor, ContainsExpression):
+        return "contains"
+    elif isinstance(descriptor, BeginsWithExpression):
+        return "beginsWith"
+    elif isinstance(descriptor, EndsWithExpression):
+        return "endsWith"
+    elif isinstance(descriptor, DoesNotContainExpression):
+        return "doesNotContain"
+    elif isinstance(descriptor, DoesNotBeginWithExpression):
+        return "doesNotBeginWith"
+    elif isinstance(descriptor, DoesNotEndWithExpression):
+        return "doesNotEndWith"
+    elif isinstance(descriptor, (IsNullExpression, IsNilExpression, IsUndefinedExpression)):
+        return "null"
+    elif isinstance(descriptor, (IsNotNullExpression, IsNotNilExpression, IsNotUndefinedExpression)):
+        return "notNull"
+    elif isinstance(descriptor, InExpression):
+        return "in"
+    elif isinstance(descriptor, NotInExpression):
+        return "notIn"
+    elif isinstance(descriptor, BetweenExpression):
+        return "between"
+    elif isinstance(descriptor, NotBetweenExpression):
+        return "notBetween"
+    elif isinstance(descriptor, AndExpression):
+        return "and"
+    elif isinstance(descriptor, OrExpression):
+        return "or"
+    elif isinstance(descriptor, CoalesceExpression):
+        return "coalesce"
+    elif isinstance(descriptor, ParseJsonExpression):
+        return "parseJson"
+    else:
+        raise ValueError(f"Unsupported descriptor type: {descriptor}")
+
+
 def get_child_descriptor(value: LazyReference, display_context: "WorkflowDisplayContext") -> BaseDescriptor:
     if isinstance(value._get, str):
         reference_parts = value._get.split(".")
@@ -21,10 +107,144 @@ def get_child_descriptor(value: LazyReference, display_context: "WorkflowDisplay
             )
 
         node_class_name = ".".join(reference_parts[:-2])
-        for node in display_context.node_displays.keys():
+        for node in display_context.global_node_displays.keys():
             if node.__name__ == node_class_name:
                 return getattr(node.Outputs, output_name)
 
         raise Exception(f"Failed to parse lazy reference: {value._get}")
 
     return value._get()
+
+
+def serialize_condition(display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
+    if isinstance(
+        condition,
+        (
+            IsNullExpression,
+            IsNotNullExpression,
+            IsNilExpression,
+            IsNotNilExpression,
+            IsUndefinedExpression,
+            IsNotUndefinedExpression,
+            ParseJsonExpression,
+        ),
+    ):
+        lhs = serialize_value(display_context, condition._expression)
+        return {
+            "type": "UNARY_EXPRESSION",
+            "lhs": lhs,
+            "operator": convert_descriptor_to_operator(condition),
+        }
+    elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
+        base = serialize_value(display_context, condition._value)
+        lhs = serialize_value(display_context, condition._start)
+        rhs = serialize_value(display_context, condition._end)
+
+        return {
+            "type": "TERNARY_EXPRESSION",
+            "base": base,
+            "operator": convert_descriptor_to_operator(condition),
+            "lhs": lhs,
+            "rhs": rhs,
+        }
+    elif isinstance(
+        condition,
+        (
+            AndExpression,
+            BeginsWithExpression,
+            CoalesceExpression,
+            ContainsExpression,
+            DoesNotBeginWithExpression,
+            DoesNotContainExpression,
+            DoesNotEndWithExpression,
+            DoesNotEqualExpression,
+            EndsWithExpression,
+            EqualsExpression,
+            GreaterThanExpression,
+            GreaterThanOrEqualToExpression,
+            InExpression,
+            LessThanExpression,
+            LessThanOrEqualToExpression,
+            NotInExpression,
+            OrExpression,
+        ),
+    ):
+        lhs = serialize_value(display_context, condition._lhs)
+        rhs = serialize_value(display_context, condition._rhs)
+
+        return {
+            "type": "BINARY_EXPRESSION",
+            "lhs": lhs,
+            "operator": convert_descriptor_to_operator(condition),
+            "rhs": rhs,
+        }
+    elif isinstance(condition, AccessorExpression):
+        return {
+            "type": "BINARY_EXPRESSION",
+            "lhs": serialize_value(display_context, condition._base),
+            "operator": "accessField",
+            "rhs": serialize_value(display_context, condition._field),
+        }
+
+    raise UnsupportedSerializationException(f"Unsupported condition type: {condition.__class__.__name__}")
+
+
+def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> JsonObject:
+    if isinstance(value, ConstantValueReference):
+        return serialize_value(display_context, value._value)
+
+    if isinstance(value, LazyReference):
+        child_descriptor = get_child_descriptor(value, display_context)
+        return serialize_value(display_context, child_descriptor)
+
+    if isinstance(value, WorkflowInputReference):
+        workflow_input_display = display_context.global_workflow_input_displays[value]
+        return {
+            "type": "WORKFLOW_INPUT",
+            "input_variable_id": str(workflow_input_display.id),
+        }
+
+    if isinstance(value, StateValueReference):
+        state_value_display = display_context.global_state_value_displays[value]
+        return {
+            "type": "STATE_VALUE",
+            "state_variable_id": str(state_value_display.id),
+        }
+
+    if isinstance(value, OutputReference):
+        upstream_node, output_display = display_context.global_node_output_displays[value]
+        upstream_node_display = display_context.global_node_displays[upstream_node]
+
+        return {
+            "type": "NODE_OUTPUT",
+            "node_id": str(upstream_node_display.node_id),
+            "node_output_id": str(output_display.id),
+        }
+
+    if isinstance(value, VellumSecretReference):
+        return {
+            "type": "VELLUM_SECRET",
+            "vellum_secret_name": value.name,
+        }
+
+    if isinstance(value, ExecutionCountReference):
+        node_class_display = display_context.global_node_displays[value.node_class]
+
+        return {
+            "type": "EXECUTION_COUNTER",
+            "node_id": str(node_class_display.node_id),
+        }
+
+    if isinstance(value, dict) and any(isinstance(v, BaseDescriptor) for v in value.values()):
+        raise ValueError("Nested references are not supported.")
+
+    if not isinstance(value, BaseDescriptor):
+        vellum_value = primitive_to_vellum_value(value)
+        return {
+            "type": "CONSTANT_VALUE",
+            "value": vellum_value.dict(),
+        }
+
+    # If it's not any of the references we know about,
+    # then try to serialize it as a nested value
+    return serialize_condition(display_context, value)

vellum_ee/workflows/display/utils/vellum.py

@@ -2,36 +2,9 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.array_vellum_value import ArrayVellumValue
-from vellum.client.types.logical_operator import LogicalOperator
 from vellum.client.types.vellum_value import VellumValue
 from vellum.client.types.vellum_variable_type import VellumVariableType
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.expressions.and_ import AndExpression
-from vellum.workflows.expressions.begins_with import BeginsWithExpression
-from vellum.workflows.expressions.between import BetweenExpression
-from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
-from vellum.workflows.expressions.contains import ContainsExpression
-from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
-from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
-from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
-from vellum.workflows.expressions.does_not_equal import DoesNotEqualExpression
-from vellum.workflows.expressions.ends_with import EndsWithExpression
-from vellum.workflows.expressions.equals import EqualsExpression
-from vellum.workflows.expressions.greater_than import GreaterThanExpression
-from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
-from vellum.workflows.expressions.in_ import InExpression
-from vellum.workflows.expressions.is_nil import IsNilExpression
-from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
-from vellum.workflows.expressions.is_not_null import IsNotNullExpression
-from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
-from vellum.workflows.expressions.is_null import IsNullExpression
-from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
-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.workflows.expressions.or_ import OrExpression
-from vellum.workflows.expressions.parse_json import ParseJsonExpression
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.references import OutputReference, WorkflowInputReference
@@ -40,7 +13,6 @@ from vellum.workflows.references.lazy import LazyReference
 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
 
@@ -145,8 +117,7 @@ def create_node_input_value_pointer_rule(
         workflow_input_display = display_context.global_workflow_input_displays[value]
         return InputVariablePointer(data=InputVariableData(input_variable_id=str(workflow_input_display.id)))
     if isinstance(value, VellumSecretReference):
-        vellum_client = create_vellum_client()
-        workspace_secret = vellum_client.workspace_secrets.retrieve(
+        workspace_secret = display_context.client.workspace_secrets.retrieve(
             id=value.name,
         )
         return WorkspaceSecretPointer(
@@ -166,52 +137,3 @@ def create_node_input_value_pointer_rule(
         return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_value)
 
     raise UnsupportedSerializationException(f"Unsupported descriptor type: {value.__class__.__name__}")
-
-
-def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperator:
-    if isinstance(descriptor, EqualsExpression):
-        return "="
-    elif isinstance(descriptor, DoesNotEqualExpression):
-        return "!="
-    elif isinstance(descriptor, LessThanExpression):
-        return "<"
-    elif isinstance(descriptor, GreaterThanExpression):
-        return ">"
-    elif isinstance(descriptor, LessThanOrEqualToExpression):
-        return "<="
-    elif isinstance(descriptor, GreaterThanOrEqualToExpression):
-        return ">="
-    elif isinstance(descriptor, ContainsExpression):
-        return "contains"
-    elif isinstance(descriptor, BeginsWithExpression):
-        return "beginsWith"
-    elif isinstance(descriptor, EndsWithExpression):
-        return "endsWith"
-    elif isinstance(descriptor, DoesNotContainExpression):
-        return "doesNotContain"
-    elif isinstance(descriptor, DoesNotBeginWithExpression):
-        return "doesNotBeginWith"
-    elif isinstance(descriptor, DoesNotEndWithExpression):
-        return "doesNotEndWith"
-    elif isinstance(descriptor, (IsNullExpression, IsNilExpression, IsUndefinedExpression)):
-        return "null"
-    elif isinstance(descriptor, (IsNotNullExpression, IsNotNilExpression, IsNotUndefinedExpression)):
-        return "notNull"
-    elif isinstance(descriptor, InExpression):
-        return "in"
-    elif isinstance(descriptor, NotInExpression):
-        return "notIn"
-    elif isinstance(descriptor, BetweenExpression):
-        return "between"
-    elif isinstance(descriptor, NotBetweenExpression):
-        return "notBetween"
-    elif isinstance(descriptor, AndExpression):
-        return "and"
-    elif isinstance(descriptor, OrExpression):
-        return "or"
-    elif isinstance(descriptor, CoalesceExpression):
-        return "coalesce"
-    elif isinstance(descriptor, ParseJsonExpression):
-        return "parseJson"
-    else:
-        raise ValueError(f"Unsupported descriptor type: {descriptor}")

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -6,6 +6,7 @@ import logging
 from uuid import UUID
 from typing import Any, Dict, ForwardRef, Generic, Iterator, List, Optional, Tuple, Type, TypeVar, Union, cast, get_args
 
+from vellum.client import Vellum as VellumClient
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -21,6 +22,7 @@ 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.workflows.vellum_client import create_vellum_client
 from vellum_ee.workflows.display.base import (
     EdgeDisplay,
     EntrypointDisplay,
@@ -46,6 +48,7 @@ from vellum_ee.workflows.display.types import (
     WorkflowInputsDisplays,
     WorkflowOutputDisplays,
 )
+from vellum_ee.workflows.display.utils.expressions import serialize_value
 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
@@ -83,9 +86,16 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         self,
         *,
         parent_display_context: Optional[WorkflowDisplayContext] = None,
+        client: Optional[VellumClient] = None,
         dry_run: bool = False,
     ):
         self._parent_display_context = parent_display_context
+        self._client = client or (
+            # propagate the client from the parent display context if it is not provided
+            self._parent_display_context.client
+            if self._parent_display_context
+            else create_vellum_client()
+        )
         self._errors: List[Exception] = []
         self._dry_run = dry_run
 
@@ -124,26 +134,24 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 }
             )
 
-        nodes: JsonArray = []
+        serialized_nodes: Dict[UUID, JsonObject] = {}
         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,
+        serialized_nodes[entrypoint_node_id] = {
+            "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():
@@ -155,7 +163,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 self.add_error(e)
                 continue
 
-            nodes.append(serialized_node)
+            serialized_nodes[node_display.node_id] = serialized_node
 
         # Add all unused nodes in the workflow
         for node in self._workflow.get_unused_nodes():
@@ -167,10 +175,11 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 self.add_error(e)
                 continue
 
-            nodes.append(serialized_node)
+            serialized_nodes[node_display.node_id] = serialized_node
 
         synthetic_output_edges: JsonArray = []
         output_variables: JsonArray = []
+        output_values: JsonArray = []
         final_output_nodes = [
             node for node in self.display_context.node_displays.keys() if issubclass(node, FinalOutputNode)
         ]
@@ -185,9 +194,9 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         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 isinstance(workflow_output.instance, OutputReference):
+                unreferenced_final_output_node_outputs.discard(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
@@ -220,24 +229,22 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 )
                 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,
-                    }
-                )
+                serialized_nodes[final_output_node_id] = {
+                    "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(
@@ -255,6 +262,19 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                         }
                     )
 
+            elif isinstance(workflow_output.instance, OutputReference):
+                terminal_node_id = workflow_output.instance.outputs_class._node_class.__id__
+                serialized_terminal_node = serialized_nodes.get(terminal_node_id)
+                if serialized_terminal_node and isinstance(serialized_terminal_node["data"], dict):
+                    serialized_terminal_node["data"]["name"] = workflow_output_display.name
+
+            output_values.append(
+                {
+                    "output_variable_id": str(workflow_output_display.id),
+                    "value": serialize_value(self.display_context, workflow_output.instance),
+                }
+            )
+
             output_variables.append(
                 {
                     "id": str(workflow_output_display.id),
@@ -309,13 +329,14 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
 
         return {
             "workflow_raw_data": {
-                "nodes": nodes,
+                "nodes": list(serialized_nodes.values()),
                 "edges": edges,
                 "display_data": self.display_context.workflow_display.display_data.dict(),
                 "definition": {
                     "name": self._workflow.__name__,
                     "module": cast(JsonArray, self._workflow.__module__.split(".")),
                 },
+                "output_values": output_values,
             },
             "input_variables": input_variables,
             "state_variables": state_variables,
@@ -482,6 +503,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             )
 
         return WorkflowDisplayContext(
+            client=self._client,
             workflow_display=workflow_meta_display,
             workflow_input_displays=workflow_input_displays,
             global_workflow_input_displays=global_workflow_input_displays,

vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py

@@ -1,6 +1,7 @@
 import types
 from typing import TYPE_CHECKING, Generic, Optional, Type, TypeVar
 
+from vellum.client import Vellum as VellumClient
 from vellum.workflows.types.generics import WorkflowType
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.registry import get_from_workflow_display_registry
@@ -35,6 +36,7 @@ def get_workflow_display(
     *,
     workflow_class: Type[WorkflowType],
     parent_display_context: Optional[WorkflowDisplayContext] = None,
+    client: Optional[VellumClient] = None,
     dry_run: bool = False,
     # DEPRECATED: The following arguments will be removed in 0.15.0
     root_workflow_class: Optional[Type[WorkflowType]] = None,
@@ -42,5 +44,6 @@ def get_workflow_display(
 ) -> "BaseWorkflowDisplay":
     return _get_workflow_display_class(workflow_class=workflow_class)(
         parent_display_context=parent_display_context,
+        client=client,
         dry_run=dry_run,
     )