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

vellum_cli/__init__.py

@@ -14,14 +14,42 @@ def main() -> None:
     pass
 
 
-@main.command()
-@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)
+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],
@@ -29,7 +57,11 @@ def push(
     deployment_description: Optional[str],
     release_tag: Optional[List[str]],
 ) -> None:
-    """Push Workflow to Vellum"""
+    """
+    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,
@@ -40,26 +72,99 @@ def push(
     )
 
 
-@main.command()
+@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.",
+    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.",
+    help="""Exclude the code definition of the Workflow from the pull response. \
+Should only be used for debugging purposes.""",
 )
-def pull(
+def workflows_pull(
     module: Optional[str],
     include_json: Optional[bool],
     workflow_sandbox_id: Optional[str],
     exclude_code: Optional[bool],
 ) -> None:
-    """Pull Workflow from Vellum"""
+    """
+    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,
@@ -68,6 +173,35 @@ def pull(
     )
 
 
+@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"""

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_cli/image_push.py

@@ -7,8 +7,8 @@ import docker
 from docker import DockerClient
 from dotenv import load_dotenv
 
-from vellum_cli.logger import load_cli_logger
 from vellum.workflows.vellum_client import create_vellum_client
+from vellum_cli.logger import load_cli_logger
 
 _SUPPORTED_ARCHITECTURE = "amd64"
 

vellum_cli/pull.py

@@ -7,10 +7,30 @@ 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:
+        workflow_config = WorkflowConfig(
+            workflow_sandbox_id=workflow_sandbox_id,
+            module=f"workflow_{workflow_sandbox_id.split('-')[0]}",
+        )
+        config.workflows.append(workflow_config)
+        return workflow_config
+    elif config.workflows:
+        return config.workflows[0]
+
+    return None
+
+
 def pull_command(
     module: Optional[str] = None,
     workflow_sandbox_id: Optional[str] = None,
@@ -21,24 +41,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 +103,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_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_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_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")