vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

vellum_ai-1.11.2.dist-info/entry_points.txt

@@ -0,0 +1,4 @@
+[console_scripts]
+generate_node_definitions=vellum_ee.scripts.generate_node_definitions:main
+vellum=vellum_cli:main
+

vellum_cli/__init__.py

@@ -5,6 +5,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.move import move_command
 from vellum_cli.ping import ping_command
 from vellum_cli.pull import pull_command
 from vellum_cli.push import push_command
@@ -67,7 +68,8 @@ with the provided module or be available for use. The Workflow Sandbox must also
 @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)
+@click.option("--release-tag", help="Release Tag for the Deployment", multiple=True)
+@click.option("--release-description", type=str, help="Description for the Release")
 @click.option(
     "--dry-run",
     is_flag=True,
@@ -87,6 +89,7 @@ def workflows_push(
     deployment_name: Optional[str],
     deployment_description: Optional[str],
     release_tag: Optional[List[str]],
+    release_description: Optional[str],
     dry_run: Optional[bool],
     strict: Optional[bool],
     workspace: Optional[str],
@@ -104,6 +107,7 @@ def workflows_push(
         deployment_name=deployment_name,
         deployment_description=deployment_description,
         release_tags=release_tag,
+        release_description=release_description,
         dry_run=dry_run,
         strict=strict,
         workspace=workspace,
@@ -117,7 +121,8 @@ def workflows_push(
 @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)
+@click.option("--release-tag", help="Release Tag for the Deployment", multiple=True)
+@click.option("--release-description", type=str, help="Description for the Release")
 @click.option(
     "--dry-run",
     is_flag=True,
@@ -137,6 +142,7 @@ def push_module(
     deployment_name: Optional[str],
     deployment_description: Optional[str],
     release_tag: Optional[List[str]],
+    release_description: Optional[str],
     dry_run: Optional[bool],
     strict: Optional[bool],
     workspace: Optional[str],
@@ -152,6 +158,7 @@ def push_module(
             deployment_name=deployment_name,
             deployment_description=deployment_description,
             release_tags=release_tag,
+            release_description=release_description,
             dry_run=dry_run,
             strict=strict,
             workspace=workspace,
@@ -259,6 +266,11 @@ Helpful for running and debugging workflows locally.""",
     help="""Directory to pull the workflow into. If not specified, \
 the workflow will be pulled into the current working directory.""",
 )
+@click.option(
+    "--workspace",
+    type=str,
+    help="The specific Workspace config to use when pulling",
+)
 def workflows_pull(
     module: Optional[str],
     include_json: Optional[bool],
@@ -268,6 +280,7 @@ def workflows_pull(
     strict: Optional[bool],
     include_sandbox: Optional[bool],
     target_directory: Optional[str],
+    workspace: Optional[str],
 ) -> None:
     """
     Pull Workflows from Vellum. If a module is provided, only the Workflow for that module will be pulled.
@@ -283,6 +296,7 @@ def workflows_pull(
         strict=strict,
         include_sandbox=include_sandbox,
         target_directory=target_directory,
+        workspace=workspace,
     )
 
 
@@ -354,9 +368,33 @@ def images() -> None:
     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:
+@click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
+@click.option("--source", type=str, help="Path to Dockerfile to build before pushing")
+def image_push(
+    image: str, tag: Optional[List[str]] = None, workspace: Optional[str] = None, source: Optional[str] = None
+) -> None:
     """Push Docker image to Vellum"""
-    image_push_command(image, tag)
+    image_push_command(image, tag, workspace, source)
+
+
+@workflows.command(name="move")
+@click.argument("old_module", required=True)
+@click.argument("new_module", required=True)
+@click.option("--workspace", type=str, help="The specific Workspace config to use when moving")
+def workflows_move(
+    old_module: str,
+    new_module: str,
+    workspace: Optional[str],
+) -> None:
+    """
+    Move/rename a Workflow module. Updates both the filesystem structure and configuration.
+    """
+
+    move_command(
+        old_module=old_module,
+        new_module=new_module,
+        workspace=workspace,
+    )
 
 
 @workflows.command(name="init")

vellum_cli/config.py

@@ -7,8 +7,8 @@ from typing import Dict, List, Literal, Optional, Union
 
 import tomli
 
-from vellum.core.pydantic_utilities import UniversalBaseModel
-from vellum.workflows.state.encoder import DefaultStateEncoder
+from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.utils.json_encoder import VellumJsonEncoder
 
 LOCKFILE_PATH = "vellum.lock.json"
 PYPROJECT_TOML_PATH = "pyproject.toml"
@@ -16,10 +16,15 @@ PYPROJECT_TOML_PATH = "pyproject.toml"
 
 class WorkspaceConfig(UniversalBaseModel):
     name: str
-    api_key: str
+    api_key: str = "VELLUM_API_KEY"
+    api_url: Optional[str] = None
 
     def merge(self, other: "WorkspaceConfig") -> "WorkspaceConfig":
-        return WorkspaceConfig(name=self.name or other.name, api_key=self.api_key or other.api_key)
+        return WorkspaceConfig(
+            name=self.name or other.name,
+            api_key=self.api_key or other.api_key,
+            api_url=self.api_url or other.api_url,
+        )
 
 
 DEFAULT_WORKSPACE_CONFIG = WorkspaceConfig(name="default", api_key="VELLUM_API_KEY")
@@ -31,6 +36,7 @@ class WorkflowDeploymentConfig(UniversalBaseModel):
     name: Optional[str] = None
     description: Optional[str] = None
     release_tags: Optional[List[str]] = None
+    release_description: Optional[str] = None
 
     def merge(self, other: "WorkflowDeploymentConfig") -> "WorkflowDeploymentConfig":
         return WorkflowDeploymentConfig(
@@ -39,6 +45,7 @@ class WorkflowDeploymentConfig(UniversalBaseModel):
             name=self.name or other.name,
             description=self.description or other.description,
             release_tags=self.release_tags or other.release_tags,
+            release_description=self.release_description or other.release_description,
         )
 
 
@@ -142,7 +149,7 @@ class VellumCliConfig(UniversalBaseModel):
     def save(self) -> None:
         lockfile_path = os.path.join(os.getcwd(), LOCKFILE_PATH)
         with open(lockfile_path, "w") as f:
-            json.dump(self.model_dump(), f, indent=2, cls=DefaultStateEncoder)
+            json.dump(self.model_dump(), f, indent=2, cls=VellumJsonEncoder)
             # Makes sure the file ends with a newline, consistent with most autoformatters
             f.write("\n")
 

vellum_cli/image_push.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 import re
 import subprocess
 from typing import List, Optional
@@ -8,16 +9,57 @@ import docker
 from docker import DockerClient
 from dotenv import load_dotenv
 
+from vellum.client.core.api_error import ApiError
 from vellum.workflows.vellum_client import create_vellum_client, create_vellum_environment
-from vellum_cli.logger import load_cli_logger
+from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, load_vellum_cli_config
+from vellum_cli.logger import handle_cli_error, load_cli_logger
 
 _SUPPORTED_ARCHITECTURE = "amd64"
 
 
-def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
-    load_dotenv()
+def image_push_command(
+    image: str, tags: Optional[List[str]] = None, workspace: Optional[str] = None, source: Optional[str] = None
+) -> None:
+    load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
     logger = load_cli_logger()
-    vellum_client = create_vellum_client()
+    config = load_vellum_cli_config()
+
+    logger.info("Cleaning up unused Docker images...")
+    try:
+        subprocess.run(
+            ["docker", "image", "prune", "--all", "-f", "--filter", "label=image-type=python-workflow-runtime"],
+            check=True,
+        )
+        logger.info("Docker image pruning completed successfully")
+    except subprocess.CalledProcessError as e:
+        logger.warning(f"Docker image pruning failed: {e}")
+
+    if source:
+        logger.info(f"Building Docker image from Dockerfile: {source}")
+
+        if not os.path.exists(source):
+            handle_cli_error(logger, "Dockerfile not found", f"Dockerfile does not exist: {source}")
+
+        build_result = subprocess.run(
+            ["docker", "buildx", "build", "-f", source, "--platform=linux/amd64", "-t", image, "."],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        if build_result.returncode != 0:
+            handle_cli_error(logger, "Docker build failed", build_result.stderr.decode("utf-8"))
+
+        logger.info("Docker build completed successfully")
+    workspace_config = next((w for w in config.workspaces if w.name == workspace), DEFAULT_WORKSPACE_CONFIG)
+
+    api_key = os.getenv(workspace_config.api_key, None)
+    if not api_key:
+        raise ValueError(f"API key {workspace_config.api_key} for workspace {workspace} not found")
+
+    vellum_client = create_vellum_client(
+        api_key=api_key,
+        api_url=workspace_config.api_url,
+    )
 
     # Check if we are self hosted by looking at our base url
     api_url = create_vellum_environment().default
@@ -70,7 +112,41 @@ def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
             exit(1)
     else:
         logger.info("Authenticating...")
-        auth = vellum_client.container_images.docker_service_token()
+        try:
+            auth = vellum_client.container_images.docker_service_token()
+        except ApiError as e:
+            if e.status_code == 401 or e.status_code == 403:
+                handle_cli_error(
+                    logger,
+                    title="Authentication failed",
+                    message="Unable to authenticate with Vellum API. Please check your API key.",
+                    suggestion="Make sure your VELLUM_API_KEY environment variable is set correctly.",
+                )
+                return
+            elif e.status_code == 400 and isinstance(e.body, dict) and "detail" in e.body:
+                handle_cli_error(
+                    logger,
+                    title="API request failed",
+                    message=e.body["detail"],
+                )
+                return
+            elif e.status_code == 500:
+                handle_cli_error(
+                    logger,
+                    title="Server error",
+                    message="The Vellum API failed with an unexpected server error.",
+                    suggestion="Please try again later and contact support if the problem persists.",
+                )
+                return
+            else:
+                handle_cli_error(
+                    logger,
+                    title="API request failed",
+                    message=f"Failed to get Docker service token from Vellum API (HTTP {e.status_code}).",
+                    suggestion="Please check your configuration and try again. If the problem persists, "
+                    "contact support.",
+                )
+                return
 
         docker_client.login(
             username="oauth2accesstoken",
@@ -131,11 +207,45 @@ def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
 
     logger.info(f"Updating Vellum metadata and validating image works in our system with image digest: {sha}...")
 
-    vellum_client.container_images.push_container_image(
-        name=image_name,
-        sha=sha,
-        tags=all_tags,
-    )
+    try:
+        vellum_client.container_images.push_container_image(
+            name=image_name,
+            sha=sha,
+            tags=all_tags,
+        )
+    except ApiError as e:
+        if e.status_code == 401 or e.status_code == 403:
+            handle_cli_error(
+                logger,
+                title="Authentication failed",
+                message="Unable to push container image metadata to Vellum API. Please check your API key.",
+                suggestion="Make sure your VELLUM_API_KEY environment variable is set correctly.",
+            )
+            return
+        elif e.status_code == 400 and isinstance(e.body, dict) and "detail" in e.body:
+            handle_cli_error(
+                logger,
+                title="API request failed",
+                message=e.body["detail"],
+            )
+            return
+        elif e.status_code == 500:
+            handle_cli_error(
+                logger,
+                title="Server error",
+                message="The Vellum API failed with an unexpected server error while pushing container image metadata.",
+                suggestion="Please try again later and contact support if the problem persists.",
+            )
+            return
+        else:
+            handle_cli_error(
+                logger,
+                title="API request failed",
+                message=f"Failed to push container image metadata to Vellum API (HTTP {e.status_code}).",
+                suggestion="Please check your configuration and try again. If the problem persists, contact support.",
+            )
+            return
+
     logger.info(f"Image successfully pushed as {image_name} to vellum with tags: {all_tags}.")
 
 

vellum_cli/logger.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from typing import Optional
 
 
 class CLIFormatter(logging.Formatter):
@@ -34,3 +35,13 @@ def load_cli_logger() -> logging.Logger:
     logger.addHandler(handler)
 
     return logger
+
+
+def handle_cli_error(logger: logging.Logger, title: str, message: str, suggestion: Optional[str] = None):
+    logger.error("\n\033[1;31m✖ {}\033[0m".format(title))  # Bold red X and title
+    logger.error("\n\033[1m{}\033[0m".format(message))  # Bold message
+    if suggestion:
+        logger.error("\n\033[1;34mNext steps:\033[0m")  # Bold blue
+        logger.error(suggestion)
+
+    exit(1)

vellum_cli/move.py

@@ -0,0 +1,56 @@
+import os
+import shutil
+from typing import Optional
+
+from dotenv import load_dotenv
+
+from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, load_vellum_cli_config
+from vellum_cli.logger import load_cli_logger
+from vellum_cli.push import module_exists
+
+
+def move_command(
+    old_module: str,
+    new_module: str,
+    workspace: Optional[str] = None,
+) -> None:
+    load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
+    logger = load_cli_logger()
+    config = load_vellum_cli_config()
+
+    if not module_exists(old_module):
+        raise ValueError(f"Module '{old_module}' does not exist in the filesystem.")
+
+    if module_exists(new_module):
+        raise ValueError(f"Module '{new_module}' already exists. Cannot move to existing module.")
+
+    matching_configs = [w for w in config.workflows if w.module == old_module]
+
+    if workspace:
+        matching_configs = [w for w in matching_configs if w.workspace == workspace]
+    else:
+        matching_configs = [w for w in matching_configs if w.workspace == DEFAULT_WORKSPACE_CONFIG.name]
+
+    if not matching_configs:
+        workspace_msg = (
+            f" in workspace '{workspace}'" if workspace else f" in workspace '{DEFAULT_WORKSPACE_CONFIG.name}'"
+        )
+        raise ValueError(f"No workflow configuration found for module '{old_module}'{workspace_msg}.")
+
+    logger.info(f"Moving module from '{old_module}' to '{new_module}'...")
+
+    old_path = os.path.join(os.getcwd(), *old_module.split("."))
+    new_path = os.path.join(os.getcwd(), *new_module.split("."))
+
+    os.makedirs(os.path.dirname(new_path), exist_ok=True)
+
+    shutil.move(old_path, new_path)
+    logger.info(f"Moved filesystem directory from '{old_path}' to '{new_path}'")
+
+    for workflow_config in matching_configs:
+        workflow_config.module = new_module
+        logger.info(f"Updated workflow configuration: {workflow_config.workflow_sandbox_id}")
+
+    config.save()
+    logger.info("Updated vellum.lock.json file.")
+    logger.info(f"Successfully moved module from '{old_module}' to '{new_module}'")

vellum_cli/ping.py

@@ -14,6 +14,7 @@ def ping_command():
     try:
         workspace = client.workspaces.workspace_identity()
         organization = client.organizations.organization_identity()
+        environment = client.environments.environment_identity()
     except ApiError as e:
         # If user did not provide an API key, we will get a 403 error
         if e.status_code == 401 or e.status_code == 403:
@@ -35,5 +36,10 @@ Organization:
 Workspace:
     ID: {workspace.id}
     Name: {workspace.name}
+
+Environment:
+    ID: {environment.id}
+    Name: {environment.name}
+    Label: {environment.label}
 """
     )

vellum_cli/pull.py

@@ -7,14 +7,20 @@ import zipfile
 from typing import Optional
 
 from dotenv import load_dotenv
-from pydash import snake_case
 
 from vellum.client.core.api_error import ApiError
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.utils.uuid import is_valid_uuid
+from vellum.workflows.utils.names import create_module_name
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum_cli.config import VellumCliConfig, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
-from vellum_cli.logger import load_cli_logger
+from vellum_cli.config import (
+    DEFAULT_WORKSPACE_CONFIG,
+    VellumCliConfig,
+    WorkflowConfig,
+    WorkflowDeploymentConfig,
+    load_vellum_cli_config,
+)
+from vellum_cli.logger import handle_cli_error, load_cli_logger
 
 ERROR_LOG_FILE_NAME = "error.log"
 METADATA_FILE_NAME = "metadata.json"
@@ -35,7 +41,11 @@ class PullContentsMetadata(UniversalBaseModel):
     runner_config: Optional[RunnerConfig] = None
     deployment_id: Optional[UUID] = None
     deployment_name: Optional[str] = None
+    deployment_history_item_id: Optional[UUID] = None
+    release_tag_id: Optional[UUID] = None
+    release_tag_name: Optional[str] = None
     workflow_sandbox_id: Optional[UUID] = None
+    workflow_version_id: Optional[UUID] = None
 
 
 def _resolve_workflow_config(
@@ -43,6 +53,7 @@ def _resolve_workflow_config(
     module: Optional[str] = None,
     workflow_sandbox_id: Optional[str] = None,
     workflow_deployment: Optional[str] = None,
+    workspace: Optional[str] = None,
 ) -> WorkflowConfigResolutionResult:
     if workflow_sandbox_id and workflow_deployment:
         raise ValueError("Cannot specify both workflow_sandbox_id and workflow_deployment")
@@ -53,6 +64,7 @@ def _resolve_workflow_config(
             workflow_config = WorkflowConfig(
                 workflow_sandbox_id=workflow_sandbox_id,
                 module=module,
+                workspace=workspace or DEFAULT_WORKSPACE_CONFIG.name,
             )
             config.workflows.append(workflow_config)
             return WorkflowConfigResolutionResult(
@@ -132,8 +144,9 @@ def pull_command(
     strict: Optional[bool] = None,
     include_sandbox: Optional[bool] = None,
     target_directory: Optional[str] = None,
+    workspace: Optional[str] = None,
 ) -> None:
-    load_dotenv()
+    load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
     logger = load_cli_logger()
     config = load_vellum_cli_config()
 
@@ -142,6 +155,7 @@ def pull_command(
         module=module,
         workflow_sandbox_id=workflow_sandbox_id,
         workflow_deployment=workflow_deployment,
+        workspace=workspace,
     )
 
     workflow_config = workflow_config_result.workflow_config
@@ -157,7 +171,20 @@ def pull_command(
     else:
         logger.info(f"Pulling workflow from {pk}...")
 
-    client = create_vellum_client()
+    resolved_workspace = workspace or workflow_config.workspace or DEFAULT_WORKSPACE_CONFIG.name
+    workspace_config = (
+        next((w for w in config.workspaces if w.name == resolved_workspace), DEFAULT_WORKSPACE_CONFIG)
+        if workspace
+        else DEFAULT_WORKSPACE_CONFIG
+    )
+    api_key = os.getenv(workspace_config.api_key)
+    if not api_key:
+        raise ValueError(f"No API key value found in environment for workspace '{workspace_config.name}'.")
+
+    client = create_vellum_client(
+        api_key=api_key,
+        api_url=workspace_config.api_url,
+    )
     query_parameters = {}
 
     if include_json:
@@ -178,9 +205,42 @@ def pull_command(
         zip_bytes = b"".join(response)
     except ApiError as e:
         if e.status_code == 401 or e.status_code == 403:
-            raise Exception("Please make sure your `VELLUM_API_KEY` environment variable is set correctly.")
-        raise Exception(
-            "The API we tried to pull is invalid. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+            handle_cli_error(
+                logger,
+                title="Authentication failed",
+                message="Unable to authenticate with the Vellum API.",
+                suggestion="Please make sure your `VELLUM_API_KEY` environment variable is set correctly and that you have access to this workflow.",  # noqa: E501
+            )
+
+        if e.status_code == 404:
+            handle_cli_error(
+                logger,
+                title="Workflow not found",
+                message=f"The workflow with ID '{pk}' could not be found.",
+                suggestion="Please verify the workflow ID is correct and that you have access to it in your workspace.",
+            )
+
+        if e.status_code == 500:
+            handle_cli_error(
+                logger,
+                title="Server error occurred",
+                message="The Vellum API encountered an internal server error while processing your request.",
+                suggestion="Please try again in a few moments. If the problem persists, contact Vellum support with the workflow ID and timestamp.",  # noqa: E501
+            )
+
+        if e.status_code == 502 or e.status_code == 503 or e.status_code == 504:
+            handle_cli_error(
+                logger,
+                title="Service temporarily unavailable",
+                message="The Vellum API is temporarily unavailable or experiencing high load.",
+                suggestion="Please wait a moment and try again. If the issue continues, check the Vellum status page or contact support.",  # noqa: E501
+            )
+
+        handle_cli_error(
+            logger,
+            title="API request failed",
+            message=f"The API request failed with status code {e.status_code}.",
+            suggestion="Please verify your `VELLUM_API_URL` environment variable is set correctly and try again.",
         )
 
     zip_buffer = io.BytesIO(zip_bytes)
@@ -201,10 +261,11 @@ def pull_command(
                         workflow_config.container_image_tag = "latest"
                 if not workflow_config.workflow_sandbox_id and pull_contents_metadata.workflow_sandbox_id:
                     workflow_config.workflow_sandbox_id = str(pull_contents_metadata.workflow_sandbox_id)
-                if not workflow_config.module and workflow_deployment and pull_contents_metadata.deployment_name:
-                    workflow_config.module = snake_case(pull_contents_metadata.deployment_name)
-                if not workflow_config.module and pull_contents_metadata.label:
-                    workflow_config.module = snake_case(pull_contents_metadata.label)
+                if not workflow_config.module:
+                    deployment_name = pull_contents_metadata.deployment_name if workflow_deployment else None
+                    workflow_config.module = create_module_name(
+                        deployment_name=deployment_name, label=pull_contents_metadata.label
+                    )
 
                 # Save or update the deployment info when pulling with --workflow-deployment
                 if workflow_deployment:
@@ -269,8 +330,11 @@ def pull_command(
                         logger.info(f"Writing to {target_file}...")
                         target.write(content)
     except zipfile.BadZipFile:
-        raise Exception(
-            "The API we tried to pull from returned an invalid zip file. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        handle_cli_error(
+            logger,
+            title="Invalid response format",
+            message="The API returned an invalid zip file format.",
+            suggestion="Please verify your `VELLUM_API_URL` environment variable is set correctly and try again.",
         )
 
     if include_json: