vellum-ai 0.14.7__py3-none-any.whl → 0.14.9__py3-none-any.whl

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_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_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_ee/workflows/display/nodes/base_node_display.py

@@ -27,6 +27,7 @@ from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpressi
 from vellum.workflows.expressions.is_null import IsNullExpression
 from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
 from vellum.workflows.expressions.not_between import NotBetweenExpression
+from vellum.workflows.expressions.parse_json import ParseJsonExpression
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.utils import get_wrapped_node
 from vellum.workflows.ports import Port
@@ -386,6 +387,9 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 "node_id": str(node_class_display.node_id),
             }
 
+        if isinstance(value, ParseJsonExpression):
+            raise ValueError("ParseJsonExpression is not supported in the UI")
+
         if not isinstance(value, BaseDescriptor):
             vellum_value = primitive_to_vellum_value(value)
             return {

vellum_ee/workflows/display/tests/test_vellum_workflow_display.py

@@ -1,4 +1,6 @@
+import pytest
 from uuid import UUID
+from typing import Dict
 
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import BaseNode
@@ -145,3 +147,115 @@ 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"
+
+
+def test_parse_json_not_supported_in_ui():
+    """
+    Test that verifies ParseJsonExpression is not yet supported in the UI.
+    This test should fail once UI support is added, at which point it should be updated.
+    """
+    # GIVEN a workflow that uses the parse_json function
+    from vellum.workflows.references.constant import ConstantValueReference
+
+    class JsonNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            json_result = ConstantValueReference('{"key": "value"}').parse_json()
+
+    class Workflow(BaseWorkflow):
+        graph = JsonNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final = JsonNode.Outputs.json_result
+
+    # WHEN we attempt to serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=Workflow,
+    )
+
+    with pytest.raises(ValueError) as exc_info:
+        workflow_display.serialize()
+
+    assert "ParseJsonExpression is not supported in the UI" == str(exc_info.value)