vellum-ai 0.13.15__tar.gz → 0.13.19__tar.gz

{vellum_ai-0.13.15 → vellum_ai-0.13.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.13.15
+Version: 0.13.19
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-0.13.15 → vellum_ai-0.13.19}/ee/vellum_cli/config.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from dataclasses import field
 import json
 import os
@@ -88,6 +89,50 @@ class WorkflowConfig(UniversalBaseModel):
         )
 
 
+def merge_workflows_by_sandbox_id(
+    workflows: List[WorkflowConfig], other_workflows: List[WorkflowConfig]
+) -> List[WorkflowConfig]:
+    merged_workflows: List[WorkflowConfig] = []
+    for self_workflow in workflows:
+        if self_workflow.workflow_sandbox_id is None:
+            # If the user defines a workflow in the pyproject.toml without a sandbox_id,
+            # we merge the workflow with one of the ones in the lockfile.
+            other_workflow = next(
+                (
+                    other_workflow
+                    for other_workflow in other_workflows
+                    if self_workflow.workspace == other_workflow.workspace
+                ),
+                None,
+            )
+            if other_workflow is not None:
+                merged_workflows.append(self_workflow.merge(other_workflow))
+            else:
+                merged_workflows.append(self_workflow)
+        else:
+            # If the user defines a workflow in the pyproject.toml with a sandbox_id,
+            # we merge the workflow with one of the ones in the lockfile with the same sandbox_id.
+            other_workflow = next(
+                (
+                    other_workflow
+                    for other_workflow in other_workflows
+                    if self_workflow.workflow_sandbox_id == other_workflow.workflow_sandbox_id
+                ),
+                None,
+            )
+            if other_workflow is not None:
+                merged_workflows.append(self_workflow.merge(other_workflow))
+            else:
+                merged_workflows.append(self_workflow)
+
+    workflow_sandbox_ids_so_far = {workflow.workflow_sandbox_id for workflow in merged_workflows}
+    for other_workflow in other_workflows:
+        if other_workflow.workflow_sandbox_id not in workflow_sandbox_ids_so_far:
+            merged_workflows.append(other_workflow)
+
+    return merged_workflows
+
+
 class VellumCliConfig(UniversalBaseModel):
     version: Literal["1.0"] = "1.0"
     workflows: List[WorkflowConfig] = field(default_factory=list)
@@ -97,24 +142,31 @@ class VellumCliConfig(UniversalBaseModel):
         lockfile_path = os.path.join(os.getcwd(), LOCKFILE_PATH)
         with open(lockfile_path, "w") as f:
             json.dump(self.model_dump(), f, indent=2, cls=DefaultStateEncoder)
+            # Makes sure the file ends with a newline, consistent with most autoformatters
+            f.write("\n")
 
     def merge(self, other: "VellumCliConfig") -> "VellumCliConfig":
         if other.version != self.version:
             raise ValueError("Lockfile version mismatch")
 
-        self_workflow_by_module = {workflow.module: workflow for workflow in self.workflows}
-        other_workflow_by_module = {workflow.module: workflow for workflow in other.workflows}
-        all_modules = sorted(set(self_workflow_by_module.keys()).union(set(other_workflow_by_module.keys())))
+        self_workflows_by_module = self.get_workflows_by_module_mapping()
+        other_workflows_by_module = other.get_workflows_by_module_mapping()
+        all_modules = sorted(set(self_workflows_by_module.keys()).union(set(other_workflows_by_module.keys())))
         merged_workflows = []
         for module in all_modules:
-            self_workflow = self_workflow_by_module.get(module)
-            other_workflow = other_workflow_by_module.get(module)
-            if self_workflow and other_workflow:
-                merged_workflows.append(self_workflow.merge(other_workflow))
-            elif self_workflow:
-                merged_workflows.append(self_workflow)
-            elif other_workflow:
-                merged_workflows.append(other_workflow)
+            self_workflows = self_workflows_by_module.get(module)
+            other_workflows = other_workflows_by_module.get(module)
+            if self_workflows and other_workflows:
+                merged_workflows.extend(
+                    merge_workflows_by_sandbox_id(
+                        self_workflows,
+                        other_workflows,
+                    )
+                )
+            elif self_workflows:
+                merged_workflows.extend(self_workflows)
+            elif other_workflows:
+                merged_workflows.extend(other_workflows)
 
         self_workspace_by_name = {workspace.name: workspace for workspace in self.workspaces}
         other_workspace_by_name = {workspace.name: workspace for workspace in other.workspaces}
@@ -132,6 +184,12 @@ class VellumCliConfig(UniversalBaseModel):
 
         return VellumCliConfig(workflows=merged_workflows, workspaces=merged_workspaces, version=self.version)
 
+    def get_workflows_by_module_mapping(self) -> Dict[str, List[WorkflowConfig]]:
+        workflows_by_module = defaultdict(list)
+        for workflow in self.workflows:
+            workflows_by_module[workflow.module].append(workflow)
+        return workflows_by_module
+
 
 def load_vellum_cli_config(root_dir: Optional[str] = None) -> VellumCliConfig:
     if root_dir is None:

{vellum_ai-0.13.15 → vellum_ai-0.13.19}/ee/vellum_cli/pull.py

@@ -31,6 +31,16 @@ class WorkflowConfigResolutionResult(UniversalBaseModel):
     pk: Optional[str] = None
 
 
+class RunnerConfig(UniversalBaseModel):
+    container_image_name: Optional[str] = None
+    container_image_tag: Optional[str] = None
+
+
+class PullContentsMetadata(UniversalBaseModel):
+    label: Optional[str] = None
+    runner_config: Optional[RunnerConfig] = None
+
+
 def _resolve_workflow_config(
     config: VellumCliConfig,
     module: Optional[str] = None,
@@ -42,14 +52,33 @@ def _resolve_workflow_config(
 
     if module:
         workflow_config = next((w for w in config.workflows if w.module == module), None)
+        if not workflow_config and workflow_sandbox_id:
+            workflow_config = WorkflowConfig(
+                workflow_sandbox_id=workflow_sandbox_id,
+                module=module,
+            )
+            config.workflows.append(workflow_config)
+            return WorkflowConfigResolutionResult(
+                workflow_config=workflow_config,
+                pk=workflow_sandbox_id,
+            )
+
         return WorkflowConfigResolutionResult(
             workflow_config=workflow_config,
             pk=workflow_config.workflow_sandbox_id if workflow_config else None,
         )
     elif workflow_sandbox_id:
+        workflow_config = next((w for w in config.workflows if w.workflow_sandbox_id == workflow_sandbox_id), None)
+        if workflow_config:
+            return WorkflowConfigResolutionResult(
+                workflow_config=workflow_config,
+                pk=workflow_sandbox_id,
+            )
+
+        # We use an empty module name to indicate that we want to backfill it once we have the Workflow Sandbox Label
         workflow_config = WorkflowConfig(
             workflow_sandbox_id=workflow_sandbox_id,
-            module=f"workflow_{workflow_sandbox_id.split('-')[0]}",
+            module="",
         )
         config.workflows.append(workflow_config)
         return WorkflowConfigResolutionResult(
@@ -98,7 +127,6 @@ def pull_command(
         workflow_sandbox_id=workflow_sandbox_id,
         workflow_deployment=workflow_deployment,
     )
-    save_lock_file = not module
 
     workflow_config = workflow_config_result.workflow_config
     if not workflow_config:
@@ -108,7 +136,11 @@ def pull_command(
     if not pk:
         raise ValueError("No workflow sandbox ID found in project to pull from.")
 
-    logger.info(f"Pulling workflow into {workflow_config.module}")
+    if workflow_config.module:
+        logger.info(f"Pulling workflow into {workflow_config.module}...")
+    else:
+        logger.info(f"Pulling workflow from {pk}...")
+
     client = create_vellum_client()
     query_parameters = {}
 
@@ -129,11 +161,30 @@ def pull_command(
     zip_bytes = b"".join(response)
     zip_buffer = io.BytesIO(zip_bytes)
 
-    target_dir = os.path.join(os.getcwd(), *workflow_config.module.split("."))
     error_content = ""
-    metadata_json: Optional[dict] = None
 
     with zipfile.ZipFile(zip_buffer) as zip_file:
+        if METADATA_FILE_NAME in zip_file.namelist():
+            metadata_json: Optional[dict] = None
+            with zip_file.open(METADATA_FILE_NAME) as source:
+                metadata_json = json.load(source)
+
+            pull_contents_metadata = PullContentsMetadata.model_validate(metadata_json)
+
+            if pull_contents_metadata.runner_config:
+                workflow_config.container_image_name = pull_contents_metadata.runner_config.container_image_name
+                workflow_config.container_image_tag = pull_contents_metadata.runner_config.container_image_tag
+                if workflow_config.container_image_name and not workflow_config.container_image_tag:
+                    workflow_config.container_image_tag = "latest"
+
+            if not workflow_config.module and pull_contents_metadata.label:
+                workflow_config.module = snake_case(pull_contents_metadata.label)
+
+        if not workflow_config.module:
+            raise ValueError(f"Failed to resolve a module name for Workflow {pk}")
+
+        target_dir = os.path.join(os.getcwd(), *workflow_config.module.split("."))
+
         # Delete files in target_dir that aren't in the zip file
         if os.path.exists(target_dir):
             ignore_patterns = (
@@ -163,7 +214,6 @@ def pull_command(
                     error_content = content
                     continue
                 if file_name == METADATA_FILE_NAME:
-                    metadata_json = json.loads(content)
                     continue
 
                 target_file = os.path.join(target_dir, file_name)
@@ -172,15 +222,6 @@ def pull_command(
                     logger.info(f"Writing to {target_file}...")
                     target.write(content)
 
-    if metadata_json:
-        runner_config = metadata_json.get("runner_config")
-
-        if runner_config:
-            workflow_config.container_image_name = runner_config.get("container_image_name")
-            workflow_config.container_image_tag = runner_config.get("container_image_tag")
-            if workflow_config.container_image_name and not workflow_config.container_image_tag:
-                workflow_config.container_image_tag = "latest"
-
     if include_json:
         logger.warning(
             """The pulled JSON representation of the Workflow should be used for debugging purposely only. \
@@ -190,16 +231,12 @@ Its schema should be considered unstable and subject to change at any time."""
     if include_sandbox:
         if not workflow_config.ignore:
             workflow_config.ignore = "sandbox.py"
-            save_lock_file = True
         elif isinstance(workflow_config.ignore, str) and "sandbox.py" != workflow_config.ignore:
             workflow_config.ignore = [workflow_config.ignore, "sandbox.py"]
-            save_lock_file = True
         elif isinstance(workflow_config.ignore, list) and "sandbox.py" not in workflow_config.ignore:
             workflow_config.ignore.append("sandbox.py")
-            save_lock_file = True
 
-    if save_lock_file:
-        config.save()
+    config.save()
 
     if error_content:
         logger.error(error_content)

{vellum_ai-0.13.15 → vellum_ai-0.13.19}/ee/vellum_cli/push.py

@@ -12,7 +12,6 @@ from dotenv import load_dotenv
 from vellum.client.core.api_error import ApiError
 from vellum.resources.workflows.client import OMIT
 from vellum.types import WorkflowPushDeploymentConfigRequest
-from vellum.workflows.utils.names import snake_to_title_case
 from vellum.workflows.vellum_client import create_vellum_client
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
@@ -111,8 +110,6 @@ def push_command(
         "container_image_name": workflow_config.container_image_name,
     }
 
-    label = snake_to_title_case(workflow_config.module.split(".")[-1])
-
     deployment_config: WorkflowPushDeploymentConfigRequest = OMIT
     deployment_config_serialized: str = OMIT
     if deploy:
@@ -158,7 +155,6 @@ def push_command(
             # Remove this once we could serialize using the artifact in Vembda
             # https://app.shortcut.com/vellum/story/5585
             exec_config=json.dumps(exec_config),
-            label=label,
             workflow_sandbox_id=workflow_config.workflow_sandbox_id or workflow_sandbox_id,
             artifact=artifact,
             # We should check with fern if we could auto-serialize typed object fields for us
@@ -230,7 +226,7 @@ def push_command(
         )  # type: ignore[attr-defined]
     else:
         logger.info(
-            f"""Successfully pushed {label} to Vellum!
+            f"""Successfully pushed {workflow_config.module} to Vellum!
 Visit at: https://app.vellum.ai/workflow-sandboxes/{response.workflow_sandbox_id}"""
         )
 

{vellum_ai-0.13.15 → vellum_ai-0.13.19}/ee/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')"
@@ -394,7 +428,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "workflows": [
                     {
                         "module": module,
-                        "workflow_sandbox_id": "0edc07cd-45b9-43e8-99bc-1f181972a857",
+                        "workflow_sandbox_id": workflow_sandbox_id,
                         "ignore": "tests/*",
                         "deployments": [
                             {
@@ -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,113 @@ 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",
+            }
+        ]
+
+
+def test_pull__multiple_instances_of_same_module__keep_when_pulling_another_module(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 vellum lock file has two instances of some other module
+    lock_data = {
+        "workflows": [
+            {
+                "module": "some_other_module",
+                "workflow_sandbox_id": str(uuid4()),
+                "workspace": "default",
+            },
+            {
+                "module": "some_other_module",
+                "workflow_sandbox_id": str(uuid4()),
+                "workspace": "other",
+            },
+        ]
+    }
+    lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(lock_json, "w") as f:
+        json.dump(lock_data, f)
+
+    # 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 on the new 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 all three entries
+    with open(lock_json) as f:
+        lock_data = json.load(f)
+        assert len(lock_data["workflows"]) == 3

{vellum_ai-0.13.15 → vellum_ai-0.13.19}/ee/vellum_cli/tests/test_push.py

@@ -119,14 +119,12 @@ def test_push__happy_path(mock_module, vellum_client, base_command):
     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
@@ -165,14 +163,12 @@ def test_push__workflow_sandbox_option__existing_id(mock_module, vellum_client,
     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 call_args["workflow_sandbox_id"] == existing_workflow_sandbox_id
     assert call_args["artifact"].name == expected_artifact_name
     assert "deplyment_config" not in call_args
@@ -216,14 +212,12 @@ def test_push__workflow_sandbox_option__existing_no_module(mock_module, vellum_c
     assert result.exit_code == 0
 
     # Get the last part of the module path and format it
-    expected_label = second_module.split(".")[-1].replace("_", " ").title()
     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["label"] == expected_label
     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
@@ -297,14 +291,12 @@ def test_push__deployment(mock_module, vellum_client, base_command):
     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"] == "{}"

{vellum_ai-0.13.15 → vellum_ai-0.13.19}/ee/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_ai-0.13.19/ee/vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -0,0 +1,32 @@
+import types
+from typing import TYPE_CHECKING, Optional, Type
+
+from vellum.workflows.types.generics import NodeType
+
+if TYPE_CHECKING:
+    from vellum_ee.workflows.display.types import NodeDisplayType
+
+
+def get_node_display_class(
+    base_class: Type["NodeDisplayType"], node_class: Type[NodeType], root_node_class: Optional[Type[NodeType]] = None
+) -> Type["NodeDisplayType"]:
+    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__}'"
+            )
+
+        return node_display_class
+
+    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
+    )
+
+    # `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_ai-0.13.15 → vellum_ai-0.13.19}/ee/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
@@ -29,6 +30,7 @@ __all__ = [
     "BaseMergeNodeDisplay",
     "BaseNoteNodeDisplay",
     "BasePromptDeploymentNodeDisplay",
+    "BaseRetryNodeDisplay",
     "BaseSearchNodeDisplay",
     "BaseSubworkflowDeploymentNodeDisplay",
     "BaseTemplatingNodeDisplay",

vellum_ai-0.13.19/ee/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