vellum-ai 0.13.14__py3-none-any.whl → 0.13.18__py3-none-any.whl

vellum_cli/tests/test_pull.py

@@ -39,6 +39,7 @@ def test_pull(vellum_client, mock_module, base_command):
     # GIVEN a module on the user's filesystem
     temp_dir = mock_module.temp_dir
     module = mock_module.module
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
@@ -56,6 +57,27 @@ def test_pull(vellum_client, mock_module, base_command):
     with open(workflow_py) as f:
         assert f.read() == "print('hello')"
 
+    # AND the vellum.lock.json file is created
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    assert os.path.exists(vellum_lock_json)
+    with open(vellum_lock_json) as f:
+        lock_data = json.load(f)
+        assert lock_data == {
+            "version": "1.0",
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": workflow_sandbox_id,
+                    "container_image_name": None,
+                    "container_image_tag": None,
+                    "ignore": None,
+                    "deployments": [],
+                    "workspace": "default",
+                }
+            ],
+            "workspaces": [],
+        }
+
 
 def test_pull__second_module(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
@@ -95,7 +117,13 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
     workflow_sandbox_id = "87654321-0000-0000-0000-000000000000"
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {"workflow.py": "print('hello')", "metadata.json": json.dumps({"label": "Super Cool Workflow"})}
+            )
+        ]
+    )
 
     # AND we are currently in a new directory
     current_dir = os.getcwd()
@@ -112,7 +140,7 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
 
     # AND the pull api is called with the workflow sandbox id
     vellum_client.workflows.pull.assert_called_once()
-    workflow_py = os.path.join(temp_dir, "workflow_87654321", "workflow.py")
+    workflow_py = os.path.join(temp_dir, "super_cool_workflow", "workflow.py")
     assert os.path.exists(workflow_py)
     with open(workflow_py) as f:
         assert f.read() == "print('hello')"
@@ -127,7 +155,7 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
             "workspaces": [],
             "workflows": [
                 {
-                    "module": "workflow_87654321",
+                    "module": "super_cool_workflow",
                     "workflow_sandbox_id": "87654321-0000-0000-0000-000000000000",
                     "ignore": None,
                     "deployments": [],
@@ -147,7 +175,13 @@ def test_pull__sandbox_id_with_other_workflow_configured(vellum_client, mock_mod
     workflow_sandbox_id = "87654321-0000-0000-0000-000000000000"
 
     # AND the workflow pull API call returns a zip file
-    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+    vellum_client.workflows.pull.return_value = iter(
+        [
+            _zip_file_map(
+                {"workflow.py": "print('hello')", "metadata.json": json.dumps({"label": "Super Cool Workflow"})}
+            )
+        ]
+    )
 
     # WHEN the user runs the pull command with the new workflow sandbox id
     runner = CliRunner()
@@ -162,7 +196,7 @@ def test_pull__sandbox_id_with_other_workflow_configured(vellum_client, mock_mod
     assert call_args[0] == workflow_sandbox_id
 
     # AND the workflow.py file is written to the module directory
-    workflow_py = os.path.join(temp_dir, "workflow_87654321", "workflow.py")
+    workflow_py = os.path.join(temp_dir, "super_cool_workflow", "workflow.py")
     assert os.path.exists(workflow_py)
     with open(workflow_py) as f:
         assert f.read() == "print('hello')"
@@ -420,8 +454,12 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
             _zip_file_map(
                 {
                     "workflow.py": "print('hello')",
-                    "metadata.json": '{"runner_config": { "container_image_name": "test", '
-                    '"container_image_tag": "1.0" } }',
+                    "metadata.json": json.dumps(
+                        {
+                            "runner_config": {"container_image_name": "test", "container_image_tag": "1.0"},
+                            "label": "Super Cool Workflow",
+                        }
+                    ),
                 }
             )
         ]
@@ -456,7 +494,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "workspace": "default",
             },
             {
-                "module": "workflow_87654321",
+                "module": "super_cool_workflow",
                 "workflow_sandbox_id": new_workflow_sandbox_id,
                 "ignore": None,
                 "deployments": [],
@@ -539,3 +577,72 @@ def test_pull__include_sandbox(vellum_client, mock_module):
     with open(lock_json) as f:
         lock_data = json.load(f)
         assert lock_data["workflows"][0]["ignore"] == "sandbox.py"
+
+
+def test_pull__same_pull_twice__one_entry_in_lockfile(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    module = mock_module.module
+    temp_dir = mock_module.temp_dir
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+
+    # AND the workflow pull API call returns a zip file both times
+    zip_contents = _zip_file_map({"workflow.py": "print('hello')"})
+    responses = iter([zip_contents, zip_contents])
+
+    def workflows_pull_side_effect(*args, **kwargs):
+        return iter([next(responses)])
+
+    vellum_client.workflows.pull.side_effect = workflows_pull_side_effect
+
+    # AND the user runs the pull command once
+    runner = CliRunner()
+    runner.invoke(cli_main, ["pull", module])
+
+    # WHEN the user runs the pull command again but with the workflow sandbox id
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-sandbox-id", workflow_sandbox_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0, (result.output, result.exception)
+
+    # AND the lockfile should only have one entry
+    lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(lock_json) as f:
+        lock_data = json.load(f)
+        assert len(lock_data["workflows"]) == 1
+
+
+def test_pull__module_not_in_config(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    module = mock_module.module
+    temp_dir = mock_module.temp_dir
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+    set_pyproject_toml = mock_module.set_pyproject_toml
+
+    # AND the pyproject.toml does not have the module configured
+    set_pyproject_toml({"workflows": []})
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+
+    # WHEN the user runs the pull command again with the workflow sandbox id and module
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", module, "--workflow-sandbox-id", workflow_sandbox_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0, (result.output, result.exception)
+
+    # AND the lockfile should have the new entry
+    lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(lock_json) as f:
+        lock_data = json.load(f)
+        assert lock_data["workflows"] == [
+            {
+                "module": module,
+                "workflow_sandbox_id": workflow_sandbox_id,
+                "ignore": None,
+                "deployments": [],
+                "container_image_name": None,
+                "container_image_tag": None,
+                "workspace": "default",
+            }
+        ]

vellum_cli/tests/test_push.py

@@ -29,6 +29,21 @@ def _extract_tar_gz(tar_gz_bytes: bytes) -> dict[str, str]:
     return files
 
 
+def _ensure_workflow_py(temp_dir: str, module: str) -> str:
+    base_dir = os.path.join(temp_dir, *module.split("."))
+    os.makedirs(base_dir, exist_ok=True)
+    workflow_py_file_content = """\
+from vellum.workflows import BaseWorkflow
+
+class ExampleWorkflow(BaseWorkflow):
+    pass
+"""
+    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
+        f.write(workflow_py_file_content)
+
+    return workflow_py_file_content
+
+
 def test_push__no_config(mock_module):
     # GIVEN no config file set
     mock_module.set_pyproject_toml({"workflows": []})
@@ -89,16 +104,7 @@ def test_push__happy_path(mock_module, vellum_client, base_command):
     module = mock_module.module
 
     # AND a workflow exists in the module successfully
-    base_dir = os.path.join(temp_dir, *module.split("."))
-    os.makedirs(base_dir, exist_ok=True)
-    workflow_py_file_content = """\
-from vellum.workflows import BaseWorkflow
-
-class ExampleWorkflow(BaseWorkflow):
-    pass
-"""
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
-        f.write(workflow_py_file_content)
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, module)
 
     # AND the push API call returns successfully
     vellum_client.workflows.push.return_value = WorkflowPushResponse(
@@ -113,14 +119,12 @@ class ExampleWorkflow(BaseWorkflow):
     assert result.exit_code == 0
 
     # Get the last part of the module path and format it
-    expected_label = mock_module.module.split(".")[-1].replace("_", " ").title()
     expected_artifact_name = f"{mock_module.module.replace('.', '__')}.tar.gz"
 
     # AND we should have called the push API with the correct args
     vellum_client.workflows.push.assert_called_once()
     call_args = vellum_client.workflows.push.call_args.kwargs
     assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
-    assert call_args["label"] == expected_label
     assert is_valid_uuid(call_args["workflow_sandbox_id"])
     assert call_args["artifact"].name == expected_artifact_name
     assert "deplyment_config" not in call_args
@@ -137,22 +141,142 @@ class ExampleWorkflow(BaseWorkflow):
     ],
     ids=["push", "workflows_push"],
 )
-def test_push__deployment(mock_module, vellum_client, base_command):
+def test_push__workflow_sandbox_option__existing_id(mock_module, vellum_client, base_command):
     # GIVEN a single workflow configured
     temp_dir = mock_module.temp_dir
     module = mock_module.module
+    existing_workflow_sandbox_id = mock_module.workflow_sandbox_id
 
     # AND a workflow exists in the module successfully
-    base_dir = os.path.join(temp_dir, *module.split("."))
-    os.makedirs(base_dir, exist_ok=True)
-    workflow_py_file_content = """\
-from vellum.workflows import BaseWorkflow
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, module)
 
-class ExampleWorkflow(BaseWorkflow):
-    pass
-"""
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
-        f.write(workflow_py_file_content)
+    # AND the push API call would return successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=existing_workflow_sandbox_id,
+    )
+
+    # WHEN calling `vellum push` with the workflow sandbox option on an existing config
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [module, "--workflow-sandbox-id", existing_workflow_sandbox_id])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # Get the last part of the module path and format it
+    expected_artifact_name = f"{mock_module.module.replace('.', '__')}.tar.gz"
+
+    # AND we should have called the push API with the correct args
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
+    assert call_args["workflow_sandbox_id"] == existing_workflow_sandbox_id
+    assert call_args["artifact"].name == expected_artifact_name
+    assert "deplyment_config" not in call_args
+
+    extracted_files = _extract_tar_gz(call_args["artifact"].read())
+    assert extracted_files["workflow.py"] == workflow_py_file_content
+
+
+def test_push__workflow_sandbox_option__existing_no_module(mock_module, vellum_client):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    first_module = mock_module.module
+    second_module = f"{first_module}2"
+    first_workflow_sandbox_id = mock_module.workflow_sandbox_id
+    second_workflow_sandbox_id = str(uuid4())
+
+    # AND the pyproject.toml has two workflow sandboxes configured
+    mock_module.set_pyproject_toml(
+        {
+            "workflows": [
+                {"module": first_module, "workflow_sandbox_id": first_workflow_sandbox_id},
+                {"module": second_module, "workflow_sandbox_id": second_workflow_sandbox_id},
+            ]
+        }
+    )
+
+    # AND a workflow exists for both modules
+    _ensure_workflow_py(temp_dir, first_module)
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, second_module)
+
+    # AND the push API call would return successfully for the second module
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=second_workflow_sandbox_id,
+    )
+
+    # WHEN calling `vellum push` with the workflow sandbox option on the second module
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", "--workflow-sandbox-id", second_workflow_sandbox_id])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # Get the last part of the module path and format it
+    expected_artifact_name = f"{second_module.replace('.', '__')}.tar.gz"
+
+    # AND we should have called the push API with the correct args
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
+    assert call_args["workflow_sandbox_id"] == second_workflow_sandbox_id
+    assert call_args["artifact"].name == expected_artifact_name
+    assert "deplyment_config" not in call_args
+
+    extracted_files = _extract_tar_gz(call_args["artifact"].read())
+    assert extracted_files["workflow.py"] == workflow_py_file_content
+
+
+def test_push__workflow_sandbox_option__existing_id_different_module(mock_module):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    second_module = f"{module}2"
+    first_workflow_sandbox_id = mock_module.workflow_sandbox_id
+    second_workflow_sandbox_id = str(uuid4())
+    set_pyproject_toml = mock_module.set_pyproject_toml
+
+    # AND the pyproject.toml has two workflow sandboxes configured
+    set_pyproject_toml(
+        {
+            "workflows": [
+                {"module": module, "workflow_sandbox_id": first_workflow_sandbox_id},
+                {"module": second_module, "workflow_sandbox_id": second_workflow_sandbox_id},
+            ]
+        }
+    )
+
+    # AND a workflow exists in both modules successfully
+    _ensure_workflow_py(temp_dir, module)
+    _ensure_workflow_py(temp_dir, second_module)
+
+    # WHEN calling `vellum push` with the first module and the second workflow sandbox id
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", module, "--workflow-sandbox-id", second_workflow_sandbox_id])
+
+    # THEN it should fail
+    assert result.exit_code == 1
+    assert result.exception
+    assert (
+        str(result.exception)
+        == "Multiple workflows found in project to push. Pushing only a single workflow is supported."
+    )
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["push"],
+        ["workflows", "push"],
+    ],
+    ids=["push", "workflows_push"],
+)
+def test_push__deployment(mock_module, vellum_client, base_command):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, module)
 
     # AND the push API call returns successfully
     vellum_client.workflows.push.return_value = WorkflowPushResponse(
@@ -167,14 +291,12 @@ class ExampleWorkflow(BaseWorkflow):
     assert result.exit_code == 0
 
     # Get the last part of the module path and format it
-    expected_label = mock_module.module.split(".")[-1].replace("_", " ").title()
     expected_artifact_name = f"{mock_module.module.replace('.', '__')}.tar.gz"
 
     # AND we should have called the push API with the correct args
     vellum_client.workflows.push.assert_called_once()
     call_args = vellum_client.workflows.push.call_args.kwargs
     assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
-    assert call_args["label"] == expected_label
     assert is_valid_uuid(call_args["workflow_sandbox_id"])
     assert call_args["artifact"].name == expected_artifact_name
     assert call_args["deployment_config"] == "{}"
@@ -252,16 +374,7 @@ def test_push__strict_option_returns_diffs(mock_module, vellum_client):
     module = mock_module.module
 
     # AND a workflow exists in the module successfully
-    base_dir = os.path.join(temp_dir, *module.split("."))
-    os.makedirs(base_dir, exist_ok=True)
-    workflow_py_file_content = """\
-from vellum.workflows import BaseWorkflow
-
-class ExampleWorkflow(BaseWorkflow):
-    pass
-"""
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
-        f.write(workflow_py_file_content)
+    _ensure_workflow_py(temp_dir, module)
 
     # AND the push API call returns a 4xx response with diffs
     vellum_client.workflows.push.side_effect = ApiError(
@@ -356,16 +469,7 @@ MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
         )
 
     # AND a workflow exists in the module successfully
-    base_dir = os.path.join(temp_dir, *module.split("."))
-    os.makedirs(base_dir, exist_ok=True)
-    workflow_py_file_content = """\
-from vellum.workflows import BaseWorkflow
-
-class ExampleWorkflow(BaseWorkflow):
-    pass
-"""
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
-        f.write(workflow_py_file_content)
+    _ensure_workflow_py(temp_dir, module)
 
     # AND the push API call returns a new workflow sandbox id
     new_workflow_sandbox_id = str(uuid4())

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -208,8 +208,8 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         return uuid4_from_hash(f"{self.node_id}|trigger")
 
     @classmethod
-    def get_from_node_display_registry(cls, node_class: Type[NodeType]) -> Type["BaseNodeDisplay"]:
-        return cls._node_display_registry[node_class]
+    def get_from_node_display_registry(cls, node_class: Type[NodeType]) -> Optional[Type["BaseNodeDisplay"]]:
+        return cls._node_display_registry.get(node_class)
 
     @classmethod
     def infer_node_class(cls) -> Type[NodeType]:

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -10,27 +10,23 @@ if TYPE_CHECKING:
 def get_node_display_class(
     base_class: Type["NodeDisplayType"], node_class: Type[NodeType], root_node_class: Optional[Type[NodeType]] = None
 ) -> Type["NodeDisplayType"]:
-    try:
-        node_display_class = base_class.get_from_node_display_registry(node_class)
-    except KeyError:
-        try:
-            base_node_display_class = get_node_display_class(
-                base_class, node_class.__bases__[0], node_class if root_node_class is None else root_node_class
+    node_display_class = base_class.get_from_node_display_registry(node_class)
+    if node_display_class:
+        if not issubclass(node_display_class, base_class):
+            raise TypeError(
+                f"Expected to find a subclass of '{base_class.__name__}' for node class '{node_class.__name__}'"
             )
 
-            # `base_node_display_class` is always a Generic class, so it's safe to index into it
-            NodeDisplayBaseClass = base_node_display_class[node_class]  # type: ignore[index]
-            NodeDisplayClass = types.new_class(
-                f"{node_class.__name__}Display",
-                bases=(NodeDisplayBaseClass,),
-            )
-            return NodeDisplayClass
-        except IndexError:
-            return base_class
+        return node_display_class
 
-    if not issubclass(node_display_class, base_class):
-        raise TypeError(
-            f"Expected to find a subclass of '{base_class.__name__}' for node class '{node_class.__name__}'"
-        )
+    base_node_display_class = get_node_display_class(
+        base_class, node_class.__bases__[0], node_class if root_node_class is None else root_node_class
+    )
 
-    return node_display_class
+    # `base_node_display_class` is always a Generic class, so it's safe to index into it
+    NodeDisplayBaseClass = base_node_display_class[node_class]  # type: ignore[index]
+    NodeDisplayClass = types.new_class(
+        f"{node_class.__name__}Display",
+        bases=(NodeDisplayBaseClass,),
+    )
+    return NodeDisplayClass

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

@@ -10,6 +10,7 @@ from .map_node import BaseMapNodeDisplay
 from .merge_node import BaseMergeNodeDisplay
 from .note_node import BaseNoteNodeDisplay
 from .prompt_deployment_node import BasePromptDeploymentNodeDisplay
+from .retry_node import BaseRetryNodeDisplay
 from .search_node import BaseSearchNodeDisplay
 from .subworkflow_deployment_node import BaseSubworkflowDeploymentNodeDisplay
 from .templating_node import BaseTemplatingNodeDisplay
@@ -17,6 +18,7 @@ from .try_node import BaseTryNodeDisplay
 
 # All node display classes must be imported here to be registered in BaseNodeDisplay's node display registry
 __all__ = [
+    "BaseRetryNodeDisplay",
     "BaseAPINodeDisplay",
     "BaseCodeExecutionNodeDisplay",
     "BaseConditionalNodeDisplay",

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

@@ -0,0 +1,10 @@
+from typing import Generic, TypeVar
+
+from vellum.workflows.nodes.core.retry_node.node import RetryNode
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+
+_RetryNodeType = TypeVar("_RetryNodeType", bound=RetryNode)
+
+
+class BaseRetryNodeDisplay(BaseNodeDisplay[_RetryNodeType], Generic[_RetryNodeType]):
+    pass

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

@@ -7,10 +7,13 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 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.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):
@@ -110,6 +113,26 @@ def test_serialize_node__retry(serialize_node):
     )
 
 
+def test_serialize_node__retry__no_display():  # GIVEN an adornment node
+    @RetryNode.wrap(max_attempts=5)
+    class StartNode(BaseNode):
+        pass
+
+    # AND a workflow that uses the adornment node
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=MyWorkflow,
+    )
+    exec_config = workflow_display.serialize()
+
+    # THEN the workflow display is created successfully
+    assert exec_config is not None
+
+
 @TryNode.wrap()
 class InnerTryGenericNode(BaseNode):
     input = Inputs.input

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

@@ -93,7 +93,7 @@ def test_serialize_workflow_with_filepath():
             "code_input_id": "f2e8a4fa-b54e-41e9-b314-0e5443519ac7",
             "runtime_input_id": "19d64948-f22b-4103-a7f5-3add184b31cc",
             "output_type": "NUMBER",
-            "packages": [],
+            "packages": [{"name": "openai", "version": "1.0.0"}],
             "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
             "log_output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
         },
@@ -566,7 +566,7 @@ def test_serialize_workflow__try_wrapped():
             "code_input_id": "f2e8a4fa-b54e-41e9-b314-0e5443519ac7",
             "runtime_input_id": "19d64948-f22b-4103-a7f5-3add184b31cc",
             "output_type": "NUMBER",
-            "packages": [],
+            "packages": [{"name": "openai", "version": "1.0.0"}],
             "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
             "log_output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
         },