vellum-ai 0.14.6__py3-none-any.whl → 0.14.8__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,20 @@ 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)
+@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_init(template_name: Optional[str] = None, target_directory: Optional[str] = None) -> None:
+    """Initialize a new Vellum Workflow using a predefined template"""
+
+    init_command(template_name=template_name, target_directory=target_directory)
+
+
 if __name__ == "__main__":
     main()

vellum_cli/config.py

@@ -50,6 +50,7 @@ class WorkflowConfig(UniversalBaseModel):
     container_image_name: Optional[str] = None
     container_image_tag: Optional[str] = None
     workspace: str = DEFAULT_WORKSPACE_CONFIG.name
+    target_directory: Optional[str] = None
 
     def merge(self, other: "WorkflowConfig") -> "WorkflowConfig":
         self_deployment_by_id = {

vellum_cli/init.py

@@ -0,0 +1,132 @@
+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, target_directory: 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.")
+
+    # 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("."))
+    workflow_config.target_directory = target_dir if target_directory else None
+
+    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,10 @@ 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("."))
+        workflow_config.target_directory = target_dir if target_directory else None
 
         # Delete files in target_dir that aren't in the zip file
         if os.path.exists(target_dir):
@@ -233,4 +237,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}")