vellum-ai 0.14.6__py3-none-any.whl → 0.14.7__py3-none-any.whl

vellum_cli/__init__.py

@@ -4,6 +4,7 @@ import click
 
 from vellum_cli.aliased_group import ClickAliasedGroup
 from vellum_cli.image_push import image_push_command
+from vellum_cli.init import init_command
 from vellum_cli.ping import ping_command
 from vellum_cli.pull import pull_command
 from vellum_cli.push import push_command
@@ -193,12 +194,20 @@ Should only be used for debugging purposes.""",
     help="""Generates a runnable sandbox.py file containing test data from the Resource's sandbox. \
 Helpful for running and debugging workflows locally.""",
 )
+@click.option(
+    "--target-dir",
+    "target_directory",  # Internal parameter name is target_directory
+    type=str,
+    help="""Directory to pull the workflow into. If not specified, \
+the workflow will be pulled into the current working directory.""",
+)
 def pull(
     ctx: click.Context,
     include_json: Optional[bool],
     exclude_code: Optional[bool],
     strict: Optional[bool],
     include_sandbox: Optional[bool],
+    target_directory: Optional[str],
 ) -> None:
     """Pull Resources from Vellum"""
 
@@ -208,6 +217,7 @@ def pull(
             exclude_code=exclude_code,
             strict=strict,
             include_sandbox=include_sandbox,
+            target_directory=target_directory,
         )
 
 
@@ -242,6 +252,13 @@ Should only be used for debugging purposes.""",
     help="""Generates a runnable sandbox.py file containing test data from the Resource's sandbox. \
 Helpful for running and debugging workflows locally.""",
 )
+@click.option(
+    "--target-dir",
+    "target_directory",  # Internal parameter name is target_directory
+    type=str,
+    help="""Directory to pull the workflow into. If not specified, \
+the workflow will be pulled into the current working directory.""",
+)
 def workflows_pull(
     module: Optional[str],
     include_json: Optional[bool],
@@ -250,6 +267,7 @@ def workflows_pull(
     exclude_code: Optional[bool],
     strict: Optional[bool],
     include_sandbox: Optional[bool],
+    target_directory: Optional[str],
 ) -> None:
     """
     Pull Workflows from Vellum. If a module is provided, only the Workflow for that module will be pulled.
@@ -264,6 +282,7 @@ def workflows_pull(
         exclude_code=exclude_code,
         strict=strict,
         include_sandbox=include_sandbox,
+        target_directory=target_directory,
     )
 
 
@@ -292,12 +311,20 @@ Should only be used for debugging purposes.""",
     help="""Generates a runnable sandbox.py file containing test data from the Resource's sandbox. \
 Helpful for running and debugging resources locally.""",
 )
+@click.option(
+    "--target-dir",
+    "target_directory",  # Internal parameter name is target_directory
+    type=str,
+    help="""Directory to pull the workflow into. If not specified, \
+the workflow will be pulled into the current working directory.""",
+)
 def pull_module(
     ctx: click.Context,
     include_json: Optional[bool],
     exclude_code: Optional[bool],
     strict: Optional[bool],
     include_sandbox: Optional[bool],
+    target_directory: Optional[str],
 ) -> None:
     """Pull a specific module from Vellum"""
 
@@ -308,6 +335,7 @@ def pull_module(
             exclude_code=exclude_code,
             strict=strict,
             include_sandbox=include_sandbox,
+            target_directory=target_directory,
         )
 
 
@@ -331,5 +359,13 @@ def image_push(image: str, tag: Optional[List[str]] = None) -> None:
     image_push_command(image, tag)
 
 
+@workflows.command(name="init")
+@click.argument("template_name", required=False)
+def workflows_init(template_name: Optional[str] = None) -> None:
+    """Initialize a new Vellum Workflow using a predefined template"""
+
+    init_command(template_name=template_name)
+
+
 if __name__ == "__main__":
     main()

vellum_cli/init.py

@@ -0,0 +1,128 @@
+import io
+import json
+import os
+import zipfile
+from typing import Optional, cast
+
+import click
+from dotenv import load_dotenv
+from pydash import snake_case
+
+from vellum.client.types.workflow_sandbox_example import WorkflowSandboxExample
+from vellum.workflows.vellum_client import create_vellum_client
+from vellum_cli.config import WorkflowConfig, load_vellum_cli_config
+from vellum_cli.logger import load_cli_logger
+from vellum_cli.pull import PullContentsMetadata, WorkflowConfigResolutionResult
+
+ERROR_LOG_FILE_NAME = "error.log"
+METADATA_FILE_NAME = "metadata.json"
+
+
+def init_command(template_name: Optional[str] = None):
+    load_dotenv()
+    logger = load_cli_logger()
+    config = load_vellum_cli_config()
+
+    client = create_vellum_client()
+    templates_response = client.workflow_sandboxes.list_workflow_sandbox_examples(tag="TEMPLATES")
+
+    templates = templates_response.results
+    if not templates:
+        logger.error("No templates available")
+        return
+
+    if template_name:
+        selected_template = next((t for t in templates if snake_case(t.label) == template_name), None)
+        if not selected_template:
+            logger.error(f"Template {template_name} not found")
+            return
+    else:
+        click.echo(click.style("Available Templates", bold=True, fg="green"))
+        for idx, template in enumerate(templates, 1):
+            click.echo(f"{idx}. {template.label}")
+
+        choice = click.prompt(
+            f"Please select a template number (1-{len(templates)})", type=click.IntRange(1, len(templates))
+        )
+        selected_template = cast(WorkflowSandboxExample, templates[choice - 1])
+
+        click.echo(click.style(f"\nYou selected: {selected_template.label}\n", bold=True, fg="cyan"))
+
+    # Create workflow config with module name from template label
+    workflow_config = WorkflowConfig(
+        workflow_sandbox_id=selected_template.id,
+        module=snake_case(selected_template.label),  # Set module name directly from template
+    )
+    config.workflows.append(workflow_config)
+
+    workflow_config_result = WorkflowConfigResolutionResult(
+        workflow_config=workflow_config,
+        pk=selected_template.id,
+    )
+
+    pk = workflow_config_result.pk
+    if not pk:
+        raise ValueError("No workflow sandbox ID found in project to pull from.")
+
+    target_dir = os.path.join(os.getcwd(), workflow_config.module)
+    if os.path.exists(target_dir):
+        click.echo(click.style(f"{target_dir} already exists.", fg="red"))
+        return
+
+    logger.info(f"Pulling workflow into {workflow_config.module}...")
+
+    query_parameters = {
+        "include_sandbox": True,
+    }
+
+    response = client.workflows.pull(
+        pk,
+        request_options={"additional_query_parameters": query_parameters},
+    )
+
+    zip_bytes = b"".join(response)
+    zip_buffer = io.BytesIO(zip_bytes)
+
+    error_content = ""
+
+    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}")
+
+        for file_name in zip_file.namelist():
+            with zip_file.open(file_name) as source:
+                content = source.read().decode("utf-8")
+                if file_name == ERROR_LOG_FILE_NAME:
+                    error_content = content
+                    continue
+                if file_name == METADATA_FILE_NAME:
+                    continue
+
+                target_file = os.path.join(target_dir, file_name)
+                os.makedirs(os.path.dirname(target_file), exist_ok=True)
+                with open(target_file, "w") as target:
+                    logger.info(f"Writing to {target_file}...")
+                    target.write(content)
+
+    config.save()
+
+    if error_content:
+        logger.error(error_content)
+    else:
+        logger.info(f"Successfully pulled Workflow into {target_dir}")

vellum_cli/pull.py

@@ -108,6 +108,7 @@ def pull_command(
     exclude_code: Optional[bool] = None,
     strict: Optional[bool] = None,
     include_sandbox: Optional[bool] = None,
+    target_directory: Optional[str] = None,
 ) -> None:
     load_dotenv()
     logger = load_cli_logger()
@@ -129,7 +130,7 @@ def pull_command(
         raise ValueError("No workflow sandbox ID found in project to pull from.")
 
     if workflow_config.module:
-        logger.info(f"Pulling workflow into {workflow_config.module}...")
+        logger.info(f"Pulling workflow {workflow_config.module}...")
     else:
         logger.info(f"Pulling workflow from {pk}...")
 
@@ -175,7 +176,9 @@ def pull_command(
         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("."))
+        # Use target_directory if provided, otherwise use current working directory
+        base_dir = os.path.join(os.getcwd(), target_directory) if target_directory else os.getcwd()
+        target_dir = os.path.join(base_dir, *workflow_config.module.split("."))
 
         # Delete files in target_dir that aren't in the zip file
         if os.path.exists(target_dir):
@@ -233,4 +236,4 @@ Its schema should be considered unstable and subject to change at any time."""
     if error_content:
         logger.error(error_content)
     else:
-        logger.info(f"Successfully pulled Workflow into {workflow_config.module}")
+        logger.info(f"Successfully pulled Workflow into {target_dir}")

vellum_cli/tests/test_init.py

@@ -0,0 +1,355 @@
+import pytest
+import io
+import json
+import os
+from unittest.mock import patch
+import zipfile
+
+from click.testing import CliRunner
+from pydash import snake_case
+
+from vellum_cli import main as cli_main
+
+
+def _zip_file_map(file_map: dict[str, str]) -> bytes:
+    # Create an in-memory bytes buffer to store the zip
+    zip_buffer = io.BytesIO()
+
+    # Create zip file and add files from file_map
+    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+        for filename, content in file_map.items():
+            zip_file.writestr(filename, content)
+
+    # Get the bytes from the buffer
+    zip_bytes = zip_buffer.getvalue()
+    zip_buffer.close()
+
+    return zip_bytes
+
+
+class MockTemplate:
+    def __init__(self, id, label):
+        self.id = id
+        self.label = label
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["workflows", "init"],
+    ],
+    ids=["workflows_init"],
+)
+def test_init_command(vellum_client, mock_module, base_command):
+    # GIVEN a module on the user's filesystem
+    temp_dir = mock_module.temp_dir
+    mock_module.set_pyproject_toml({"workflows": []})
+    # GIVEN the vellum client returns a list of template workflows
+    fake_templates = [
+        MockTemplate(id="template-1", label="Example Workflow"),
+        MockTemplate(id="template-2", label="Another Workflow"),
+    ]
+    vellum_client.workflow_sandboxes.list_workflow_sandbox_examples.return_value.results = fake_templates
+
+    # 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 `init` command and selects the first template
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command, input="1\n")
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND `vellum_client.workflows.pull` is called with the selected template ID
+    vellum_client.workflows.pull.assert_called_once_with(
+        "template-1",
+        request_options={"additional_query_parameters": {"include_sandbox": True}},
+    )
+
+    # AND the `workflow.py` file should be created in the correct module directory
+    workflow_py = os.path.join(temp_dir, "example_workflow", "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 should be 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["workflows"] == [
+            {
+                "module": "example_workflow",
+                "workflow_sandbox_id": "template-1",
+                "ignore": None,
+                "deployments": [],
+                "container_image_name": None,
+                "container_image_tag": None,
+                "workspace": "default",
+            }
+        ]
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["workflows", "init"],
+    ],
+    ids=["workflows_init"],
+)
+def test_init_command__invalid_template_id(vellum_client, mock_module, base_command):
+    # GIVEN a module on the user's filesystem
+    temp_dir = mock_module.temp_dir
+    mock_module.set_pyproject_toml({"workflows": []})
+    # GIVEN the vellum client returns a list of template workflows
+    fake_templates = [
+        MockTemplate(id="template-1", label="Example Workflow"),
+        MockTemplate(id="template-2", label="Another Workflow"),
+    ]
+    vellum_client.workflow_sandboxes.list_workflow_sandbox_examples.return_value.results = fake_templates
+
+    # WHEN the user runs the `init` command, enters invalid input and then cancels
+    runner = CliRunner()
+    # Mock click.prompt to raise a KeyboardInterrupt (simulating Ctrl+C)
+    with patch("click.prompt", side_effect=KeyboardInterrupt):
+        runner = CliRunner()
+        result = runner.invoke(cli_main, base_command)
+
+    # THEN the command is aborted
+    assert result.exit_code != 0
+    assert "Aborted!" in result.output  # Click shows this message on Ctrl+C
+
+    # AND `vellum_client.workflows.pull` is not called
+    vellum_client.workflows.pull.assert_not_called()
+
+    # AND no workflow files are created
+    workflow_py = os.path.join(temp_dir, "example_workflow", "workflow.py")
+    assert not os.path.exists(workflow_py)
+
+    # AND the lock file remains empty
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    if os.path.exists(vellum_lock_json):
+        with open(vellum_lock_json) as f:
+            lock_data = json.load(f)
+            assert lock_data["workflows"] == []
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["workflows", "init"],
+    ],
+    ids=["workflows_init"],
+)
+def test_init_command__no_templates(vellum_client, mock_module, base_command):
+    # GIVEN a module on the user's filesystem
+    temp_dir = mock_module.temp_dir
+    mock_module.set_pyproject_toml({"workflows": []})
+    # GIVEN the vellum client returns no template workflows
+    vellum_client.workflow_sandboxes.list_workflow_sandbox_examples.return_value.results = []
+
+    # WHEN the user runs the `init` command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command)
+
+    # THEN the command gracefully exits
+    assert result.exit_code == 0
+    assert "No templates available" in result.output
+
+    # AND `vellum_client.workflows.pull` is not called
+    vellum_client.workflows.pull.assert_not_called()
+
+    # AND no workflow files are created
+    workflow_py = os.path.join(temp_dir, "example_workflow", "workflow.py")
+    assert not os.path.exists(workflow_py)
+
+    # AND the lock file remains empty
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    if os.path.exists(vellum_lock_json):
+        with open(vellum_lock_json) as f:
+            lock_data = json.load(f)
+            assert lock_data["workflows"] == []
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["workflows", "init"],
+    ],
+    ids=["workflows_init"],
+)
+def test_init_command_target_directory_exists(vellum_client, mock_module, base_command):
+    """
+    GIVEN a target directory already exists
+    WHEN the user tries to run the `init` command
+    THEN the command should fail and exit without modifying existing files
+    """
+    temp_dir = mock_module.temp_dir
+    existing_workflow_dir = os.path.join(temp_dir, "example_workflow")
+
+    # Create the target directory to simulate it already existing
+    os.makedirs(existing_workflow_dir, exist_ok=True)
+
+    # Ensure directory exists before command execution
+    assert os.path.exists(existing_workflow_dir)
+
+    # GIVEN the vellum client returns a list of template workflows
+    fake_templates = [
+        MockTemplate(id="template-1", label="Example Workflow"),
+    ]
+    vellum_client.workflow_sandboxes.list_workflow_sandbox_examples.return_value.results = fake_templates
+
+    # 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 `init` command and selects the template
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command, input="1\n")
+
+    # THEN the command should detect the existing directory and abort
+    assert result.exit_code == 0
+    assert f"{existing_workflow_dir} already exists." in result.output
+
+    # Ensure the directory still exists (wasn't deleted or modified)
+    assert os.path.exists(existing_workflow_dir)
+
+    # AND `vellum_client.workflows.pull` is not called
+    vellum_client.workflows.pull.assert_not_called()
+
+    # AND no workflow files are created
+    workflow_py = os.path.join(temp_dir, "example_workflow", "workflow.py")
+    assert not os.path.exists(workflow_py)
+
+    # AND the lock file remains empty
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    if os.path.exists(vellum_lock_json):
+        with open(vellum_lock_json) as f:
+            lock_data = json.load(f)
+            assert lock_data["workflows"] == []
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["workflows", "init"],
+    ],
+    ids=["workflows_init"],
+)
+def test_init_command_with_template_name(vellum_client, mock_module, base_command):
+    # GIVEN a module on the user's filesystem
+    temp_dir = mock_module.temp_dir
+    mock_module.set_pyproject_toml({"workflows": []})
+
+    # GIVEN the vellum client returns a list of template workflows
+    fake_templates = [
+        MockTemplate(id="template-1", label="Example Workflow"),
+        MockTemplate(id="template-2", label="Another Workflow"),
+    ]
+    vellum_client.workflow_sandboxes.list_workflow_sandbox_examples.return_value.results = fake_templates
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter(
+        [_zip_file_map({"workflow.py": "print('hello')", "README.md": "# Another Workflow\nThis is a test template."})]
+    )
+
+    # WHEN the user runs the `init` command with a specific template name
+    template_name = snake_case("Another Workflow")
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [template_name])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND `vellum_client.workflows.pull` is called with the correct template ID
+    vellum_client.workflows.pull.assert_called_once_with(
+        "template-2",  # ID of "Another Workflow"
+        request_options={"additional_query_parameters": {"include_sandbox": True}},
+    )
+
+    # AND the workflow files should be created in the correct module directory
+    workflow_py = os.path.join(temp_dir, "another_workflow", "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 should be created with the correct data
+    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["workflows"] == [
+            {
+                "module": "another_workflow",
+                "workflow_sandbox_id": "template-2",
+                "ignore": None,
+                "deployments": [],
+                "container_image_name": None,
+                "container_image_tag": None,
+                "workspace": "default",
+            }
+        ]
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["workflows", "init"],
+    ],
+    ids=["workflows_init"],
+)
+def test_init_command_with_nonexistent_template_name(vellum_client, mock_module, base_command):
+    # GIVEN a module on the user's filesystem
+    temp_dir = mock_module.temp_dir
+    mock_module.set_pyproject_toml({"workflows": []})
+
+    # GIVEN the vellum client returns a list of template workflows
+    fake_templates = [
+        MockTemplate(id="template-1", label="Example Workflow"),
+        MockTemplate(id="template-2", label="Another Workflow"),
+    ]
+    vellum_client.workflow_sandboxes.list_workflow_sandbox_examples.return_value.results = fake_templates
+
+    # WHEN the user runs the `init` command with a non-existent template name
+    nonexistent_template = "nonexistent_template"
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [nonexistent_template])
+
+    # THEN the command should indicate the template was not found
+    assert result.exit_code == 0
+    assert f"Template {nonexistent_template} not found" in result.output
+
+    # AND `vellum_client.workflows.pull` is not called
+    vellum_client.workflows.pull.assert_not_called()
+
+    # AND no workflow files are created
+    example_workflow_dir = os.path.join(temp_dir, "example_workflow")
+    nonexistent_workflow_dir = os.path.join(temp_dir, nonexistent_template)
+
+    assert not os.path.exists(example_workflow_dir)
+    assert not os.path.exists(nonexistent_workflow_dir)
+
+    # AND the lock file remains empty
+    vellum_lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    if os.path.exists(vellum_lock_json):
+        with open(vellum_lock_json) as f:
+            lock_data = json.load(f)
+            assert lock_data["workflows"] == []