vellum-ai 0.10.8__tar.gz → 0.11.0__tar.gz

{vellum_ai-0.10.8 → vellum_ai-0.11.0}/PKG-INFO

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

vellum_ai-0.11.0/ee/vellum_cli/__init__.py

@@ -0,0 +1,226 @@
+from typing import List, Optional
+
+import click
+
+from vellum_cli.aliased_group import ClickAliasedGroup
+from vellum_cli.image_push import image_push_command
+from vellum_cli.pull import pull_command
+from vellum_cli.push import push_command
+
+
+@click.group(cls=ClickAliasedGroup)
+def main() -> None:
+    """Vellum SDK CLI"""
+    pass
+
+
+class PushGroup(ClickAliasedGroup):
+    def get_command(self, ctx, cmd_name):
+        # First try to get the command normally
+        cmd = super().get_command(ctx, cmd_name)
+        if cmd is not None:
+            return cmd
+
+        # If no command found, call our catch-all
+        return push_module
+
+
+@main.group(invoke_without_command=True, cls=PushGroup)
+@click.pass_context
+def push(
+    ctx: click.Context,
+) -> None:
+    """Push Resources to Vellum"""
+
+    if ctx.invoked_subcommand is None:
+        push_command()
+
+
+@main.group()
+def workflows():
+    """Operations related to Vellum Workflows"""
+    pass
+
+
+@workflows.command(name="push")
+@click.argument("module", required=False)
+@click.option("--deploy", is_flag=True, help="Deploy the Workflow after pushing it to Vellum")
+@click.option("--deployment-label", type=str, help="Label to use for the Deployment")
+@click.option("--deployment-name", type=str, help="Unique name for the Deployment")
+@click.option("--deployment-description", type=str, help="Description for the Deployment")
+@click.option("--release-tag", type=list, help="Release Tag for the Deployment", multiple=True)
+def workflows_push(
+    module: Optional[str],
+    deploy: Optional[bool],
+    deployment_label: Optional[str],
+    deployment_name: Optional[str],
+    deployment_description: Optional[str],
+    release_tag: Optional[List[str]],
+) -> None:
+    """
+    Push Workflows to Vellum. If a module is provided, only the Workflow for that module will be pushed.
+    If no module is provided, the first configured Workflow will be pushed.
+    """
+
+    push_command(
+        module=module,
+        deploy=deploy,
+        deployment_label=deployment_label,
+        deployment_name=deployment_name,
+        deployment_description=deployment_description,
+        release_tags=release_tag,
+    )
+
+
+@push.command(name="*", hidden=True)
+@click.pass_context
+@click.option("--deploy", is_flag=True, help="Deploy the Resource after pushing it to Vellum")
+@click.option("--deployment-label", type=str, help="Label to use for the Deployment")
+@click.option("--deployment-name", type=str, help="Unique name for the Deployment")
+@click.option("--deployment-description", type=str, help="Description for the Deployment")
+@click.option("--release-tag", type=list, help="Release Tag for the Deployment", multiple=True)
+def push_module(
+    ctx: click.Context,
+    deploy: Optional[bool],
+    deployment_label: Optional[str],
+    deployment_name: Optional[str],
+    deployment_description: Optional[str],
+    release_tag: Optional[List[str]],
+) -> None:
+    """Push a specific module to Vellum"""
+
+    if ctx.parent:
+        push_command(
+            module=ctx.parent.invoked_subcommand,
+            deploy=deploy,
+            deployment_label=deployment_label,
+            deployment_name=deployment_name,
+            deployment_description=deployment_description,
+            release_tags=release_tag,
+        )
+
+
+class PullGroup(ClickAliasedGroup):
+    def get_command(self, ctx, cmd_name):
+        # First try to get the command normally
+        cmd = super().get_command(ctx, cmd_name)
+        if cmd is not None:
+            return cmd
+
+        # If no command found, call our catch-all
+        return pull_module
+
+
+@main.group(invoke_without_command=True, cls=PullGroup)
+@click.pass_context
+@click.option(
+    "--include-json",
+    is_flag=True,
+    help="""Include the JSON representation of each Resource in the pull response. \
+Should only be used for debugging purposes.""",
+)
+@click.option(
+    "--exclude-code",
+    is_flag=True,
+    help="""Exclude the code definition of each Resource from the pull response. \
+Should only be used for debugging purposes.""",
+)
+def pull(
+    ctx: click.Context,
+    include_json: Optional[bool],
+    exclude_code: Optional[bool],
+) -> None:
+    """Pull Resources from Vellum"""
+
+    if ctx.invoked_subcommand is None:
+        pull_command(
+            include_json=include_json,
+            exclude_code=exclude_code,
+        )
+
+
+@workflows.command(name="pull")
+@click.argument("module", required=False)
+@click.option(
+    "--include-json",
+    is_flag=True,
+    help="""Include the JSON representation of the Workflow in the pull response. \
+Should only be used for debugging purposes.""",
+)
+@click.option("--workflow-sandbox-id", type=str, help="Pull the Workflow from a specific Sandbox ID")
+@click.option(
+    "--exclude-code",
+    is_flag=True,
+    help="""Exclude the code definition of the Workflow from the pull response. \
+Should only be used for debugging purposes.""",
+)
+def workflows_pull(
+    module: Optional[str],
+    include_json: Optional[bool],
+    workflow_sandbox_id: Optional[str],
+    exclude_code: Optional[bool],
+) -> None:
+    """
+    Pull Workflows from Vellum. If a module is provided, only the Workflow for that module will be pulled.
+    If no module is provided, the first configured Workflow will be pulled.
+    """
+
+    pull_command(
+        module=module,
+        include_json=include_json,
+        workflow_sandbox_id=workflow_sandbox_id,
+        exclude_code=exclude_code,
+    )
+
+
+@pull.command(name="*")
+@click.pass_context
+@click.option(
+    "--include-json",
+    is_flag=True,
+    help="""Include the JSON representation of the Resource in the pull response. \
+Should only be used for debugging purposes.""",
+)
+@click.option(
+    "--exclude-code",
+    is_flag=True,
+    help="""Exclude the code definition of the Resource from the pull response. \
+Should only be used for debugging purposes.""",
+)
+def pull_module(
+    ctx: click.Context,
+    include_json: Optional[bool],
+    exclude_code: Optional[bool],
+) -> None:
+    """Pull a specific module from Vellum"""
+
+    if ctx.parent:
+        pull_command(
+            module=ctx.parent.invoked_subcommand,
+            include_json=include_json,
+            exclude_code=exclude_code,
+        )
+
+
+@main.group(aliases=["images", "image"])
+def images() -> None:
+    """Vellum Docker Images"""
+    pass
+
+
+@images.command(name="push")
+@click.argument("image", required=True)
+@click.option(
+    "--tag",
+    "-t",
+    multiple=True,
+    help="Tags the provided image inside of Vellum's repo. "
+    "This field does not push multiple local tags of the passed in image.",
+)
+def image_push(image: str, tag: Optional[List[str]] = None) -> None:
+    """Push Docker image to Vellum"""
+    image_push_command(image, tag)
+
+
+if __name__ == "__main__":
+    main()

{vellum_ai-0.10.8 → vellum_ai-0.11.0}/ee/vellum_cli/config.py

@@ -7,7 +7,6 @@ from typing import Dict, List, Literal, Optional, Union
 import tomli
 
 from vellum.core.pydantic_utilities import UniversalBaseModel
-
 from vellum.workflows.state.encoder import DefaultStateEncoder
 
 LOCKFILE_PATH = "vellum.lock.json"

{vellum_ai-0.10.8 → vellum_ai-0.11.0}/ee/vellum_cli/pull.py

@@ -7,10 +7,28 @@ from typing import Optional
 from dotenv import load_dotenv
 
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum_cli.config import WorkflowConfig, load_vellum_cli_config
+from vellum_cli.config import VellumCliConfig, WorkflowConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 
 
+def resolve_workflow_config(
+    config: VellumCliConfig,
+    module: Optional[str] = None,
+    workflow_sandbox_id: Optional[str] = None,
+) -> Optional[WorkflowConfig]:
+    if module:
+        return next((w for w in config.workflows if w.module == module), None)
+    elif workflow_sandbox_id:
+        return WorkflowConfig(
+            workflow_sandbox_id=workflow_sandbox_id,
+            module=f"workflow_{workflow_sandbox_id.split('-')[0]}",
+        )
+    elif config.workflows:
+        return config.workflows[0]
+
+    return None
+
+
 def pull_command(
     module: Optional[str] = None,
     workflow_sandbox_id: Optional[str] = None,
@@ -21,24 +39,15 @@ def pull_command(
     logger = load_cli_logger()
     config = load_vellum_cli_config()
 
-    workflow_config = (
-        next((w for w in config.workflows if w.module == module), None)
-        if module
-        else (config.workflows[0] if config.workflows else None)
+    workflow_config = resolve_workflow_config(
+        config,
+        module,
+        workflow_sandbox_id,
     )
-    save_lock_file = False
-    if workflow_config is None:
-        if module:
-            raise ValueError(f"No workflow config for '{module}' found in project to pull.")
-        elif workflow_sandbox_id:
-            workflow_config = WorkflowConfig(
-                workflow_sandbox_id=workflow_sandbox_id,
-                module=f"workflow_{workflow_sandbox_id.split('-')[0]}",
-            )
-            config.workflows.append(workflow_config)
-            save_lock_file = True
-        else:
-            raise ValueError("No workflow config found in project to pull from.")
+    save_lock_file = not module
+
+    if not workflow_config:
+        raise ValueError("No workflow config found in project to pull from.")
 
     if not workflow_config.workflow_sandbox_id:
         raise ValueError("No workflow sandbox ID found in project to pull from.")
@@ -92,7 +101,8 @@ def pull_command(
 
     if include_json:
         logger.warning(
-            "The pulled JSON representation of the Workflow should be used for debugging purposely only. Its schema should be considered unstable and subject to change at any time."
+            """The pulled JSON representation of the Workflow should be used for debugging purposely only. \
+Its schema should be considered unstable and subject to change at any time."""
         )
 
     if save_lock_file:

{vellum_ai-0.10.8 → vellum_ai-0.11.0}/ee/vellum_cli/push.py

@@ -10,23 +10,22 @@ from dotenv import load_dotenv
 
 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 WorkflowDeploymentConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 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
-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
 
 
 def push_command(
-    module: Optional[str],
-    deploy: Optional[bool],
-    deployment_label: Optional[str],
-    deployment_name: Optional[str],
-    deployment_description: Optional[str],
-    release_tags: Optional[List[str]],
+    module: Optional[str] = None,
+    deploy: Optional[bool] = None,
+    deployment_label: Optional[str] = None,
+    deployment_name: Optional[str] = None,
+    deployment_description: Optional[str] = None,
+    release_tags: Optional[List[str]] = None,
 ) -> None:
     load_dotenv()
     logger = load_cli_logger()

vellum_ai-0.11.0/ee/vellum_cli/tests/conftest.py

@@ -0,0 +1,40 @@
+import pytest
+import os
+import shutil
+import tempfile
+from uuid import uuid4
+from typing import Any, Callable, Dict, Generator, Tuple
+
+import tomli_w
+
+
+@pytest.fixture
+def mock_module() -> Generator[Tuple[str, str, Callable[[Dict[str, Any]], None]], None, None]:
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+    module = "examples.mock"
+
+    def set_pyproject_toml(vellum_config: Dict[str, Any]) -> None:
+        pyproject_toml_path = os.path.join(temp_dir, "pyproject.toml")
+        with open(pyproject_toml_path, "wb") as f:
+            tomli_w.dump(
+                {"tool": {"vellum": vellum_config}},
+                f,
+            )
+
+    set_pyproject_toml(
+        {
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": str(uuid4()),
+                }
+            ]
+        }
+    )
+
+    yield temp_dir, module, set_pyproject_toml
+
+    os.chdir(current_dir)
+    shutil.rmtree(temp_dir)

vellum_ai-0.11.0/ee/vellum_cli/tests/test_main.py

@@ -0,0 +1,11 @@
+import subprocess
+
+
+def test_cli__root():
+    """
+    A test sanity ensuring that the CLI is accessible
+    """
+
+    result = subprocess.run(["vellum", "--help"], capture_output=True)
+    assert result.returncode == 0
+    assert result.stdout.startswith(b"Usage: vellum")

vellum_ai-0.11.0/ee/vellum_cli/tests/test_pull.py

@@ -0,0 +1,267 @@
+import pytest
+import io
+import os
+import tempfile
+from uuid import uuid4
+import zipfile
+
+from click.testing import CliRunner
+
+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
+
+
+@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
+
+    # 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
+    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__second_module(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    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')"})])
+
+    # 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
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", 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__sandbox_id_with_no_config(vellum_client):
+    # GIVEN a 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')"})])
+
+    # AND we are currently in a new directory
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    # WHEN the user runs the pull command with the workflow sandbox id and no module
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-sandbox-id", workflow_sandbox_id])
+    os.chdir(current_dir)
+
+    # 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')"
+
+
+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
+
+    # AND the workflow pull API call returns a zip file
+    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")
+    os.makedirs(os.path.dirname(other_file_path), exist_ok=True)
+    with open(other_file_path, "w") as f:
+        f.write("print('hello')")
+
+    # WHEN the user runs the pull command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # 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')"
+
+    # AND the other_file.py file is deleted
+    assert not os.path.exists(other_file_path)
+
+
+def test_pull__remove_missing_files__ignore_pattern(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    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')"})])
+
+    # AND there is already a different file in the module directory
+    other_file_path = os.path.join(temp_dir, *module.split("."), "other_file.py")
+    os.makedirs(os.path.dirname(other_file_path), exist_ok=True)
+    with open(other_file_path, "w") as f:
+        f.write("print('hello')")
+
+    # AND there is already a test file
+    test_file_path = os.path.join(temp_dir, *module.split("."), "tests", "test_workflow.py")
+    os.makedirs(os.path.dirname(test_file_path), exist_ok=True)
+    with open(test_file_path, "w") as f:
+        f.write("print('hello')")
+
+    # AND the ignore pattern is set to tests
+    set_pyproject_toml(
+        {
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": str(uuid4()),
+                    "ignore": "tests/*",
+                }
+            ]
+        }
+    )
+
+    # WHEN the user runs the pull command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # 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')"
+
+    # AND the other_file.py file is deleted
+    assert not os.path.exists(other_file_path)
+
+    # AND the tests/test_workflow.py file is untouched
+    assert os.path.exists(test_file_path)
+
+
+def test_pull__include_json(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    _, 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": "{}"})]
+    )
+
+    # WHEN the user runs the pull command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module, "--include-json"])
+
+    # 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}
+
+
+def test_pull__exclude_code(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    _, 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": "{}"})]
+    )
+
+    # WHEN the user runs the pull command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["pull", module, "--exclude-code"])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # 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}