vellum-ai 0.10.9__py3-none-any.whl → 0.11.1__py3-none-any.whl

vellum_cli/tests/test_pull.py

@@ -1,18 +1,17 @@
 import pytest
 import io
+import json
 import os
-import shutil
 import tempfile
 from uuid import uuid4
 import zipfile
-from typing import Generator, Tuple
 
-import tomli_w
+from click.testing import CliRunner
 
-from vellum_cli.pull import pull_command
+from vellum_cli import main as cli_main
 
 
-def zip_file_map(file_map: dict[str, str]) -> bytes:
+def _zip_file_map(file_map: dict[str, str]) -> bytes:
     # Create an in-memory bytes buffer to store the zip
     zip_buffer = io.BytesIO()
 
@@ -28,47 +27,60 @@ def zip_file_map(file_map: dict[str, str]) -> bytes:
     return zip_bytes
 
 
-@pytest.fixture
-def mock_module() -> Generator[Tuple[str, str], None, None]:
-    current_dir = os.getcwd()
-    temp_dir = tempfile.mkdtemp()
-    os.chdir(temp_dir)
-    module = "examples.mock"
-
-    with open(os.path.join(temp_dir, "pyproject.toml"), "wb") as f:
-        tomli_w.dump(
-            {
-                "tool": {
-                    "vellum": {
-                        "workflows": [
-                            {
-                                "module": module,
-                                "workflow_sandbox_id": str(uuid4()),
-                            }
-                        ]
-                    }
-                }
-            },
-            f,
-        )
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["pull"],
+        ["workflows", "pull"],
+    ],
+    ids=["pull", "workflows_pull"],
+)
+def test_pull(vellum_client, mock_module, base_command):
+    # GIVEN a module on the user's filesystem
+    temp_dir, module, _ = mock_module
 
-    yield temp_dir, module
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
 
-    os.chdir(current_dir)
-    shutil.rmtree(temp_dir)
+    # WHEN the user runs the pull command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [module])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the workflow.py file is written to the module directory
+    workflow_py = os.path.join(temp_dir, *module.split("."), "workflow.py")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
+        assert f.read() == "print('hello')"
 
 
-def test_pull(vellum_client, mock_module):
+def test_pull__second_module(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
-    temp_dir, module = mock_module
+    temp_dir, module, set_pyproject_toml = mock_module
 
     # 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')"})])
+
+    # AND the module we're about to pull is configured second
+    set_pyproject_toml(
+        {
+            "workflows": [
+                {"module": "another.module", "workflow_sandbox_id": str(uuid4())},
+                {"module": module, "workflow_sandbox_id": str(uuid4())},
+            ]
+        }
+    )
 
     # WHEN the user runs the pull command
-    pull_command(module)
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module])
 
-    # THEN the workflow.py file is written to the module directory
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the workflow.py file is written to the module directory
     workflow_py = os.path.join(temp_dir, *module.split("."), "workflow.py")
     assert os.path.exists(workflow_py)
     with open(workflow_py) as f:
@@ -80,7 +92,7 @@ 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')"})])
 
     # AND we are currently in a new directory
     current_dir = os.getcwd()
@@ -88,23 +100,73 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
     os.chdir(temp_dir)
 
     # WHEN the user runs the pull command with the workflow sandbox id and no module
-    pull_command(workflow_sandbox_id=workflow_sandbox_id)
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-sandbox-id", workflow_sandbox_id])
     os.chdir(current_dir)
 
-    # THEN the pull api is called with exclude_code=True
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # 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")
     assert os.path.exists(workflow_py)
     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.loads(f.read())
+        assert lock_data == {
+            "version": "1.0",
+            "workflows": [
+                {
+                    "module": "workflow_87654321",
+                    "workflow_sandbox_id": "87654321-0000-0000-0000-000000000000",
+                    "ignore": None,
+                    "deployments": [],
+                }
+            ],
+        }
+
+
+def test_pull__sandbox_id_with_other_workflow_configured(vellum_client, mock_module):
+    # GIVEN a pyproject.toml with a workflow configured
+    temp_dir, _, _ = mock_module
+
+    # AND a different workflow sandbox id
+    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')"})])
+
+    # WHEN the user runs the pull command with the new workflow sandbox id
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-sandbox-id", workflow_sandbox_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the pull api is called with the new workflow sandbox id
+    vellum_client.workflows.pull.assert_called_once()
+    call_args = vellum_client.workflows.pull.call_args.args
+    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")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
+        assert f.read() == "print('hello')"
+
 
 def test_pull__remove_missing_files(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
-    temp_dir, module = mock_module
+    temp_dir, module, _ = mock_module
 
     # 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')"})])
 
     # AND there is already a different file in the module directory
     other_file_path = os.path.join(temp_dir, *module.split("."), "other_file.py")
@@ -113,9 +175,13 @@ def test_pull__remove_missing_files(vellum_client, mock_module):
         f.write("print('hello')")
 
     # WHEN the user runs the pull command
-    pull_command(module)
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
 
-    # THEN the workflow.py file is written to the module directory
+    # AND the workflow.py file is written to the module directory
     assert os.path.exists(os.path.join(temp_dir, *module.split("."), "workflow.py"))
     with open(os.path.join(temp_dir, *module.split("."), "workflow.py")) as f:
         assert f.read() == "print('hello')"
@@ -126,10 +192,10 @@ def test_pull__remove_missing_files(vellum_client, mock_module):
 
 def test_pull__remove_missing_files__ignore_pattern(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
-    temp_dir, module = mock_module
+    temp_dir, module, set_pyproject_toml = mock_module
 
     # 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')"})])
 
     # AND there is already a different file in the module directory
     other_file_path = os.path.join(temp_dir, *module.split("."), "other_file.py")
@@ -144,28 +210,26 @@ def test_pull__remove_missing_files__ignore_pattern(vellum_client, mock_module):
         f.write("print('hello')")
 
     # AND the ignore pattern is set to tests
-    with open(os.path.join(temp_dir, "pyproject.toml"), "wb") as f:
-        tomli_w.dump(
-            {
-                "tool": {
-                    "vellum": {
-                        "workflows": [
-                            {
-                                "module": module,
-                                "workflow_sandbox_id": str(uuid4()),
-                                "ignore": "tests/*",
-                            }
-                        ]
-                    }
+    set_pyproject_toml(
+        {
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": str(uuid4()),
+                    "ignore": "tests/*",
                 }
-            },
-            f,
-        )
+            ]
+        }
+    )
 
     # WHEN the user runs the pull command
-    pull_command(module)
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
 
-    # THEN the workflow.py file is written to the module directory
+    # AND the workflow.py file is written to the module directory
     assert os.path.exists(os.path.join(temp_dir, *module.split("."), "workflow.py"))
     with open(os.path.join(temp_dir, *module.split("."), "workflow.py")) as f:
         assert f.read() == "print('hello')"
@@ -179,17 +243,21 @@ def test_pull__remove_missing_files__ignore_pattern(vellum_client, mock_module):
 
 def test_pull__include_json(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
-    _, module = mock_module
+    _, module, __ = mock_module
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
-        [zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
+        [_zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
     )
 
     # WHEN the user runs the pull command
-    pull_command(module, include_json=True)
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module, "--include-json"])
 
-    # THEN the pull api is called with include_json=True
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the pull api is called with include_json=True
     vellum_client.workflows.pull.assert_called_once()
     call_args = vellum_client.workflows.pull.call_args.kwargs
     assert call_args["request_options"]["additional_query_parameters"] == {"include_json": True}
@@ -197,17 +265,21 @@ def test_pull__include_json(vellum_client, mock_module):
 
 def test_pull__exclude_code(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
-    _, module = mock_module
+    _, module, __ = mock_module
 
     # AND the workflow pull API call returns a zip file
     vellum_client.workflows.pull.return_value = iter(
-        [zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
+        [_zip_file_map({"workflow.py": "print('hello')", "workflow.json": "{}"})]
     )
 
     # WHEN the user runs the pull command
-    pull_command(module, exclude_code=True)
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module, "--exclude-code"])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
 
-    # THEN the pull api is called with exclude_code=True
+    # AND the pull api is called with exclude_code=True
     vellum_client.workflows.pull.assert_called_once()
     call_args = vellum_client.workflows.pull.call_args.kwargs
     assert call_args["request_options"]["additional_query_parameters"] == {"exclude_code": True}

vellum_cli/tests/test_push.py

@@ -0,0 +1,173 @@
+import pytest
+import io
+import json
+import os
+import tarfile
+from uuid import uuid4
+
+from click.testing import CliRunner
+
+from vellum.client.types.workflow_push_response import WorkflowPushResponse
+from vellum.evaluations.utils.uuid import is_valid_uuid
+from vellum_cli import main as cli_main
+
+
+def _extract_tar_gz(tar_gz_bytes: bytes) -> dict[str, str]:
+    files = {}
+    with tarfile.open(fileobj=io.BytesIO(tar_gz_bytes), mode="r:gz") as tar:
+        for member in tar.getmembers():
+            if not member.isfile():
+                continue
+            content = tar.extractfile(member)
+            if content is None:
+                continue
+
+            files[member.name] = content.read().decode("latin-1")
+
+    return files
+
+
+def test_push__no_config(mock_module):
+    # GIVEN no config file set
+    _, _, set_pyproject_toml = mock_module
+    set_pyproject_toml({"workflows": []})
+
+    # WHEN calling `vellum push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push"])
+
+    # THEN it should fail
+    assert result.exit_code == 1
+    assert result.exception
+    assert str(result.exception) == "No Workflows found in project to push."
+
+
+def test_push__multiple_workflows_configured__no_module_specified(mock_module):
+    # GIVEN multiple workflows configured
+    _, _, set_pyproject_toml = mock_module
+    set_pyproject_toml({"workflows": [{"module": "examples.mock"}, {"module": "examples.mock2"}]})
+
+    # WHEN calling `vellum push` without a module specified
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push"])
+
+    # 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."
+    )
+
+
+def test_push__multiple_workflows_configured__not_found_module(mock_module):
+    # GIVEN multiple workflows configured
+    _, module, set_pyproject_toml = mock_module
+    set_pyproject_toml({"workflows": [{"module": "examples.mock2"}, {"module": "examples.mock3"}]})
+
+    # WHEN calling `vellum push` with a module that doesn't exist
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module])
+
+    # THEN it should fail
+    assert result.exit_code == 1
+    assert result.exception
+    assert str(result.exception) == f"No workflow config for '{module}' found in project to push."
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["push"],
+        ["workflows", "push"],
+    ],
+    ids=["push", "workflows_push"],
+)
+def test_push__happy_path(mock_module, vellum_client, base_command):
+    # GIVEN a single workflow configured
+    temp_dir, module, _ = mock_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)
+
+    # AND the push API call returns successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+    )
+
+    # WHEN calling `vellum push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [module])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # 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"] == "Mock"
+    assert is_valid_uuid(call_args["workflow_sandbox_id"])
+    assert call_args["artifact"].name == "examples__mock.tar.gz"
+    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
+
+
+@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, module, _ = mock_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)
+
+    # AND the push API call returns successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=str(uuid4()),
+    )
+
+    # WHEN calling `vellum push`
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [module, "--deploy"])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # 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"] == "Mock"
+    assert is_valid_uuid(call_args["workflow_sandbox_id"])
+    assert call_args["artifact"].name == "examples__mock.tar.gz"
+    assert call_args["deployment_config"] == "{}"
+
+    extracted_files = _extract_tar_gz(call_args["artifact"].read())
+    assert extracted_files["workflow.py"] == workflow_py_file_content

vellum_ee/workflows/display/base.py

@@ -20,6 +20,7 @@ WorkflowMetaDisplayOverridesType = TypeVar("WorkflowMetaDisplayOverridesType", b
 @dataclass
 class WorkflowInputsDisplayOverrides:
     id: UUID
+    required: Optional[bool] = None
     color: Optional[str] = None
 
 

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -100,8 +100,9 @@ class BaseNodeDisplay(Generic[NodeType]):
             if origin is dict and isinstance(node_display_attribute, dict):
                 if len(args) == 2:
                     key_type, value_type = args
-                    if all(isinstance(k, key_type) and isinstance(v, value_type) for k, v in
-                           node_display_attribute.items()):
+                    if all(
+                        isinstance(k, key_type) and isinstance(v, value_type) for k, v in node_display_attribute.items()
+                    ):
                         return cast(_NodeDisplayAttrType, node_display_attribute)
                 raise ValueError(f"Node {cls.__name__} must define an explicit {attribute} of type {attribute_type}.")
 

vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -1,12 +1,12 @@
 from uuid import UUID
 from typing import ClassVar, Dict, Optional
 
+from vellum.workflows.ports import Port
+from vellum.workflows.types.generics import NodeType
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import PortDisplay
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.vellum import NodeDisplayData
-from vellum.workflows.ports import Port
-from vellum.workflows.types.generics import NodeType
 
 
 class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,7 +1,7 @@
 from typing import Optional, Type
 
-from vellum_ee.workflows.display.types import NodeDisplayType
 from vellum.workflows.types.generics import NodeType
+from vellum_ee.workflows.display.types import NodeDisplayType
 
 
 def get_node_display_class(

vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -1,8 +1,8 @@
 import pytest
 from uuid import UUID
 
-from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum.workflows.nodes.bases import BaseNode
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 
 
 @pytest.fixture

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

@@ -32,5 +32,5 @@ __all__ = [
     "BaseSearchNodeDisplay",
     "BaseSubworkflowDeploymentNodeDisplay",
     "BaseTemplatingNodeDisplay",
-    "BaseTryNodeDisplay"
+    "BaseTryNodeDisplay",
 ]