vellum-ai 0.13.15__py3-none-any.whl → 0.13.19__py3-none-any.whl

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_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_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_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
@@ -29,6 +30,7 @@ __all__ = [
     "BaseMergeNodeDisplay",
     "BaseNoteNodeDisplay",
     "BasePromptDeploymentNodeDisplay",
+    "BaseRetryNodeDisplay",
     "BaseSearchNodeDisplay",
     "BaseSubworkflowDeploymentNodeDisplay",
     "BaseTemplatingNodeDisplay",

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):
@@ -77,6 +80,11 @@ def test_serialize_node__retry(serialize_node):
                         "module": ["vellum", "workflows", "nodes", "core", "retry_node", "node"],
                     },
                     "attributes": [
+                        {
+                            "id": "8a07dc58-3fed-41d4-8ca6-31ee0bb86c61",
+                            "name": "delay",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
                         {
                             "id": "f388e93b-8c68-4f54-8577-bbd0c9091557",
                             "name": "max_attempts",
@@ -110,6 +118,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",
         },