vellum-ai 0.14.7__tar.gz → 0.14.8__tar.gz

{vellum_ai-0.14.7 → vellum_ai-0.14.8}/PKG-INFO

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

{vellum_ai-0.14.7 → vellum_ai-0.14.8}/ee/vellum_cli/__init__.py

@@ -361,10 +361,17 @@ def image_push(image: str, tag: Optional[List[str]] = None) -> None:
 
 @workflows.command(name="init")
 @click.argument("template_name", required=False)
-def workflows_init(template_name: Optional[str] = None) -> None:
+@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)
+    init_command(template_name=template_name, target_directory=target_directory)
 
 
 if __name__ == "__main__":

{vellum_ai-0.14.7 → vellum_ai-0.14.8}/ee/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_ai-0.14.7 → vellum_ai-0.14.8}/ee/vellum_cli/init.py

@@ -18,7 +18,7 @@ ERROR_LOG_FILE_NAME = "error.log"
 METADATA_FILE_NAME = "metadata.json"
 
 
-def init_command(template_name: Optional[str] = None):
+def init_command(template_name: Optional[str] = None, target_directory: Optional[str] = None):
     load_dotenv()
     logger = load_cli_logger()
     config = load_vellum_cli_config()
@@ -64,7 +64,11 @@ def init_command(template_name: Optional[str] = None):
     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)
+    # 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

{vellum_ai-0.14.7 → vellum_ai-0.14.8}/ee/vellum_cli/pull.py

@@ -179,6 +179,7 @@ def pull_command(
         # 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):

{vellum_ai-0.14.7 → vellum_ai-0.14.8}/ee/vellum_cli/tests/test_init.py

@@ -1,4 +1,3 @@
-import pytest
 import io
 import json
 import os
@@ -33,14 +32,7 @@ class MockTemplate:
         self.label = label
 
 
-@pytest.mark.parametrize(
-    "base_command",
-    [
-        ["workflows", "init"],
-    ],
-    ids=["workflows_init"],
-)
-def test_init_command(vellum_client, mock_module, base_command):
+def test_init_command(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
     temp_dir = mock_module.temp_dir
     mock_module.set_pyproject_toml({"workflows": []})
@@ -52,18 +44,11 @@ def test_init_command(vellum_client, mock_module, base_command):
     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')",
-                }
-            )
-        ]
-    )
+    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")
+    result = runner.invoke(cli_main, ["workflows", "init"], input="1\n")
 
     # THEN the command returns successfully
     assert result.exit_code == 0
@@ -94,18 +79,12 @@ def test_init_command(vellum_client, mock_module, base_command):
                 "container_image_name": None,
                 "container_image_tag": None,
                 "workspace": "default",
+                "target_directory": None,
             }
         ]
 
 
-@pytest.mark.parametrize(
-    "base_command",
-    [
-        ["workflows", "init"],
-    ],
-    ids=["workflows_init"],
-)
-def test_init_command__invalid_template_id(vellum_client, mock_module, base_command):
+def test_init_command__invalid_template_id(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
     temp_dir = mock_module.temp_dir
     mock_module.set_pyproject_toml({"workflows": []})
@@ -121,7 +100,7 @@ def test_init_command__invalid_template_id(vellum_client, mock_module, base_comm
     # 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)
+        result = runner.invoke(cli_main, ["workflows", "init"])
 
     # THEN the command is aborted
     assert result.exit_code != 0
@@ -142,14 +121,7 @@ def test_init_command__invalid_template_id(vellum_client, mock_module, base_comm
             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):
+def test_init_command__no_templates(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
     temp_dir = mock_module.temp_dir
     mock_module.set_pyproject_toml({"workflows": []})
@@ -158,7 +130,7 @@ def test_init_command__no_templates(vellum_client, mock_module, base_command):
 
     # WHEN the user runs the `init` command
     runner = CliRunner()
-    result = runner.invoke(cli_main, base_command)
+    result = runner.invoke(cli_main, ["workflows", "init"])
 
     # THEN the command gracefully exits
     assert result.exit_code == 0
@@ -179,14 +151,7 @@ def test_init_command__no_templates(vellum_client, mock_module, base_command):
             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):
+def test_init_command_target_directory_exists(vellum_client, mock_module):
     """
     GIVEN a target directory already exists
     WHEN the user tries to run the `init` command
@@ -208,19 +173,11 @@ def test_init_command_target_directory_exists(vellum_client, mock_module, base_c
     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')",
-                }
-            )
-        ]
-    )
+    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")
+    result = runner.invoke(cli_main, ["workflows", "init"], input="1\n")
 
     # THEN the command should detect the existing directory and abort
     assert result.exit_code == 0
@@ -244,14 +201,7 @@ def test_init_command_target_directory_exists(vellum_client, mock_module, base_c
             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):
+def test_init_command_with_template_name(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
     temp_dir = mock_module.temp_dir
     mock_module.set_pyproject_toml({"workflows": []})
@@ -264,14 +214,12 @@ def test_init_command_with_template_name(vellum_client, mock_module, base_comman
     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."})]
-    )
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
 
     # 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])
+    result = runner.invoke(cli_main, ["workflows", "init", template_name])
 
     # THEN the command returns successfully
     assert result.exit_code == 0
@@ -305,18 +253,12 @@ def test_init_command_with_template_name(vellum_client, mock_module, base_comman
                 "container_image_name": None,
                 "container_image_tag": None,
                 "workspace": "default",
+                "target_directory": None,
             }
         ]
 
 
-@pytest.mark.parametrize(
-    "base_command",
-    [
-        ["workflows", "init"],
-    ],
-    ids=["workflows_init"],
-)
-def test_init_command_with_nonexistent_template_name(vellum_client, mock_module, base_command):
+def test_init_command_with_nonexistent_template_name(vellum_client, mock_module):
     # GIVEN a module on the user's filesystem
     temp_dir = mock_module.temp_dir
     mock_module.set_pyproject_toml({"workflows": []})
@@ -331,7 +273,7 @@ def test_init_command_with_nonexistent_template_name(vellum_client, mock_module,
     # 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])
+    result = runner.invoke(cli_main, ["workflows", "init", nonexistent_template])
 
     # THEN the command should indicate the template was not found
     assert result.exit_code == 0
@@ -353,3 +295,179 @@ def test_init_command_with_nonexistent_template_name(vellum_client, mock_module,
         with open(vellum_lock_json) as f:
             lock_data = json.load(f)
             assert lock_data["workflows"] == []
+
+
+def test_init__with_target_dir(vellum_client, mock_module):
+    # 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"),
+    ]
+    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')"})])
+
+    # AND a target directory
+    target_dir = os.path.join(temp_dir, "dir")
+    os.makedirs(target_dir, exist_ok=True)
+
+    # WHEN the user runs the init command with target-dir
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "init", "--target-dir", target_dir], input="1\n")
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the `workflow.py` file should be created in the target directory
+    module_path = os.path.join(target_dir, "example_workflow")
+    workflow_py = os.path.join(module_path, "workflow.py")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
+        assert f.read() == "print('hello')"
+
+    # AND the files are not in the default module directory
+    default_module_path = os.path.join(temp_dir, "example_workflow", "workflow.py")
+    assert not os.path.exists(default_module_path)
+
+    # AND the vellum.lock.json file should be created in the original directory
+    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",
+                "target_directory": module_path,
+            }
+        ]
+
+
+def test_init__with_nested_target_dir(vellum_client, mock_module):
+    # 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"),
+    ]
+    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')"})])
+
+    # AND a nested target directory that doesn't exist yet
+    nested_target_dir = os.path.join(temp_dir, "dir-1", "dir-2")
+
+    # WHEN the user runs the init command with nested target-dir
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "init", "--target-dir", nested_target_dir], input="1\n")
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the nested directory with module subdirectory should be created
+    module_path = os.path.join(nested_target_dir, "example_workflow")
+    assert os.path.exists(module_path)
+
+    # AND the workflow.py file is written to the nested target directory
+    workflow_py = os.path.join(module_path, "workflow.py")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
+        assert f.read() == "print('hello')"
+
+    # AND the files are not in the default module directory
+    default_module_path = os.path.join(temp_dir, "example_workflow", "workflow.py")
+    assert not os.path.exists(default_module_path)
+
+    # AND the vellum.lock.json file is still updated
+    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",
+                "target_directory": module_path,
+            }
+        ]
+
+
+def test_init__with_template_name_and_target_dir(vellum_client, mock_module):
+    # 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')"})])
+
+    # AND a target directory
+    target_dir = os.path.join(temp_dir, "dir")
+    os.makedirs(target_dir, exist_ok=True)
+
+    # WHEN the user runs the init command with a specific template name and target-dir
+    template_name = snake_case("Another Workflow")
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "init", template_name, "--target-dir", target_dir])
+
+    # 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 target directory with the correct module subdirectory
+    module_path = os.path.join(target_dir, "another_workflow")
+    workflow_py = os.path.join(module_path, "workflow.py")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
+        assert f.read() == "print('hello')"
+
+    # AND the files are not in the default module directory
+    default_module_path = os.path.join(temp_dir, "another_workflow", "workflow.py")
+    assert not os.path.exists(default_module_path)
+
+    # 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",
+                "target_directory": module_path,
+            }
+        ]

{vellum_ai-0.14.7 → vellum_ai-0.14.8}/ee/vellum_cli/tests/test_pull.py

@@ -73,6 +73,7 @@ def test_pull(vellum_client, mock_module, base_command):
                     "ignore": None,
                     "deployments": [],
                     "workspace": "default",
+                    "target_directory": None,
                 }
             ],
             "workspaces": [],
@@ -167,6 +168,7 @@ def test_pull__with_target_dir(vellum_client, mock_module, base_command):
                     "ignore": None,
                     "deployments": [],
                     "workspace": "default",
+                    "target_directory": module_path,
                 }
             ],
             "workspaces": [],
@@ -233,6 +235,7 @@ def test_pull__with_nested_target_dir(vellum_client, mock_module, base_command):
                     "ignore": None,
                     "deployments": [],
                     "workspace": "default",
+                    "target_directory": module_path,
                 }
             ],
             "workspaces": [],
@@ -289,6 +292,7 @@ def test_pull__sandbox_id_with_no_config(vellum_client):
                     "container_image_tag": None,
                     "container_image_name": None,
                     "workspace": "default",
+                    "target_directory": None,
                 }
             ],
         }
@@ -372,6 +376,7 @@ def test_pull__workflow_deployment_with_no_config(vellum_client):
                     "container_image_tag": None,
                     "container_image_name": None,
                     "workspace": "default",
+                    "target_directory": None,
                 }
             ],
             "workspaces": [],
@@ -619,6 +624,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "container_image_name": None,
                 "container_image_tag": None,
                 "workspace": "default",
+                "target_directory": None,
             },
             {
                 "module": "super_cool_workflow",
@@ -628,6 +634,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "container_image_name": "test",
                 "container_image_tag": "1.0",
                 "workspace": "default",
+                "target_directory": None,
             },
         ]
 
@@ -771,6 +778,7 @@ def test_pull__module_not_in_config(vellum_client, mock_module):
                 "container_image_name": None,
                 "container_image_tag": None,
                 "workspace": "default",
+                "target_directory": None,
             }
         ]
 

{vellum_ai-0.14.7 → vellum_ai-0.14.8}/ee/vellum_cli/tests/test_push.py

@@ -508,6 +508,7 @@ MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
             "container_image_tag": None,
             "deployments": [],
             "ignore": None,
+            "target_directory": None,
         }
 
 

{vellum_ai-0.14.7 → vellum_ai-0.14.8}/ee/vellum_ee/workflows/display/tests/test_vellum_workflow_display.py

@@ -1,4 +1,5 @@
 from uuid import UUID
+from typing import Dict
 
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import BaseNode
@@ -145,3 +146,85 @@ def test_vellum_workflow_display_serialize_valid_handle_ids_for_base_nodes():
         assert (
             node["trigger"]["id"] in edge_target_handle_ids
         ), f"Trigger {node['trigger']['id']} from node {node['label']} not found in edge target handle ids"
+
+
+def test_vellum_workflow_display__serialize_with_unused_nodes_and_edges():
+    # GIVEN a workflow with active and unused nodes
+    class NodeA(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class NodeB(BaseNode):
+        pass
+
+    class NodeC(BaseNode):
+        pass
+
+    # AND A workflow that uses them correctly
+    class Workflow(BaseWorkflow):
+        graph = NodeA
+        unused_graphs = {NodeB >> NodeC}
+
+        class Outputs(BaseWorkflow.Outputs):
+            final = NodeA.Outputs.result
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=Workflow,
+    )
+
+    # WHEN we serialize the workflow
+    exec_config = workflow_display.serialize()
+
+    # THEN the serialized workflow contains the expected nodes and edges
+    raw_data = exec_config["workflow_raw_data"]
+    assert isinstance(raw_data, dict)
+
+    nodes = raw_data["nodes"]
+    edges = raw_data["edges"]
+
+    assert isinstance(nodes, list)
+    assert isinstance(edges, list)
+
+    # Find nodes by their definition name
+    node_ids: Dict[str, str] = {}
+
+    for node in nodes:
+        assert isinstance(node, dict)
+        definition = node.get("definition")
+        if definition is None:
+            continue
+
+        assert isinstance(definition, dict)
+        name = definition.get("name")
+        if not isinstance(name, str):
+            continue
+
+        if name in ["NodeA", "NodeB", "NodeC"]:
+            node_id = node.get("id")
+            if isinstance(node_id, str):
+                node_ids[name] = node_id
+
+    # Verify all nodes are present
+    assert "NodeA" in node_ids, "Active node NodeA not found in serialized output"
+    assert "NodeB" in node_ids, "Unused node NodeB not found in serialized output"
+    assert "NodeC" in node_ids, "Unused node NodeC not found in serialized output"
+
+    # Verify the edge between NodeB and NodeC is present
+    edge_found = False
+    for edge in edges:
+        assert isinstance(edge, dict)
+        source_id = edge.get("source_node_id")
+        target_id = edge.get("target_node_id")
+
+        if (
+            isinstance(source_id, str)
+            and isinstance(target_id, str)
+            and source_id == node_ids["NodeB"]
+            and target_id == node_ids["NodeC"]
+        ):
+            edge_found = True
+            break
+
+    assert edge_found, "Edge between unused nodes NodeB and NodeC not found in serialized output"