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

vellum_cli/tests/test_pull.py

@@ -112,6 +112,133 @@ def test_pull__second_module(vellum_client, mock_module):
         assert f.read() == "print('hello')"
 
 
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["pull"],
+        ["workflows", "pull"],
+    ],
+    ids=["pull", "workflows_pull"],
+)
+def test_pull__with_target_dir(vellum_client, mock_module, base_command):
+    # GIVEN a module on the user's filesystem
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+
+    # AND a target directory
+    target_dir = os.path.join(temp_dir, "dir")
+    os.makedirs(target_dir, exist_ok=True)
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+
+    # WHEN the user runs the pull command with target-dir
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [module, "--target-dir", target_dir])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the workflow.py file is written to the target directory
+    module_path = os.path.join(target_dir, *module.split("."))
+    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, *module.split("."), "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 == {
+            "version": "1.0",
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": workflow_sandbox_id,
+                    "container_image_name": None,
+                    "container_image_tag": None,
+                    "ignore": None,
+                    "deployments": [],
+                    "workspace": "default",
+                }
+            ],
+            "workspaces": [],
+        }
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["pull"],
+        ["workflows", "pull"],
+    ],
+    ids=["pull", "workflows_pull"],
+)
+def test_pull__with_nested_target_dir(vellum_client, mock_module, base_command):
+    # GIVEN a module on the user's filesystem
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+
+    # AND a nested target directory that doesn't exist yet
+    nested_target_dir = os.path.join(temp_dir, "dir-1", "dir-2")
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+
+    # WHEN the user runs the pull command with nested target-dir
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [module, "--target-dir", nested_target_dir])
+
+    # 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, *module.split("."))
+    assert os.path.exists(module_path)
+
+    # AND the nested directory should be created
+    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, *module.split("."), "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 == {
+            "version": "1.0",
+            "workflows": [
+                {
+                    "module": module,
+                    "workflow_sandbox_id": workflow_sandbox_id,
+                    "container_image_name": None,
+                    "container_image_tag": None,
+                    "ignore": None,
+                    "deployments": [],
+                    "workspace": "default",
+                }
+            ],
+            "workspaces": [],
+        }
+
+
 def test_pull__sandbox_id_with_no_config(vellum_client):
     # GIVEN a workflow sandbox id
     workflow_sandbox_id = "87654321-0000-0000-0000-000000000000"

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -47,7 +47,7 @@ from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_di
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
 from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator, primitive_to_vellum_value
-from vellum_ee.workflows.display.vellum import CodeResourceDefinition, GenericNodeDisplayData
+from vellum_ee.workflows.display.vellum import CodeResourceDefinition, NodeDisplayData
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -295,9 +295,9 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         cls._node_display_registry[node_class] = cls
 
-    def _get_generic_node_display_data(self) -> GenericNodeDisplayData:
-        explicit_value = self._get_explicit_node_display_attr("display_data", GenericNodeDisplayData)
-        return explicit_value if explicit_value else GenericNodeDisplayData()
+    def _get_generic_node_display_data(self) -> NodeDisplayData:
+        explicit_value = self._get_explicit_node_display_attr("display_data", NodeDisplayData)
+        return explicit_value if explicit_value else NodeDisplayData()
 
     def serialize_condition(self, display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
         if isinstance(

vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -7,6 +7,7 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.references import LazyReference
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input_value_pointer_rules
@@ -43,6 +44,10 @@ class MyNodeADisplay(BaseNodeVellumDisplay[MyNodeA]):
 class MyNodeB(BaseNode):
     example = MyNodeA.Outputs.output
     fallback_example = MyNodeA.Outputs.output.coalesce(Inputs.example_workflow_input).coalesce("fallback")
+    constant_coalesce = Inputs.example_workflow_input.coalesce("default_value")
+    lazy_coalesce: BaseDescriptor[str] = LazyReference(
+        lambda: MyNodeA.Outputs.output.coalesce(Inputs.example_workflow_input)
+    )
 
 
 @pytest.mark.parametrize(
@@ -76,6 +81,32 @@ class MyNodeB(BaseNode):
                 ConstantValuePointer(type="CONSTANT_VALUE", data=StringVellumValue(value="fallback")),
             ],
         ),
+        (
+            MyNodeB.constant_coalesce,
+            [
+                InputVariablePointer(
+                    type="INPUT_VARIABLE",
+                    data=InputVariableData(input_variable_id="a154c29d-fac0-4cd0-ba88-bc52034f5470"),
+                ),
+                ConstantValuePointer(type="CONSTANT_VALUE", data=StringVellumValue(value="default_value")),
+            ],
+        ),
+        (
+            MyNodeB.lazy_coalesce,
+            [
+                NodeOutputPointer(
+                    type="NODE_OUTPUT",
+                    data=NodeOutputData(
+                        node_id="b48fa5e0-d7d3-4fe3-ae48-615415011cc5",
+                        output_id="4b16a629-11a1-4b3f-a965-a57b872d13b8",
+                    ),
+                ),
+                InputVariablePointer(
+                    type="INPUT_VARIABLE",
+                    data=InputVariableData(input_variable_id="a154c29d-fac0-4cd0-ba88-bc52034f5470"),
+                ),
+            ],
+        ),
     ],
 )
 def test_create_node_input_value_pointer_rules(

vellum_ee/workflows/display/nodes/vellum/utils.py

@@ -5,8 +5,10 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.references import NodeReference
+from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule
 from vellum_ee.workflows.display.vellum import (
     ConstantValuePointer,
@@ -56,6 +58,12 @@ def create_node_input_value_pointer_rules(
                 raise ValueError(f"Expected NodeReference {value.name} to have an instance")
             value = cast(BaseDescriptor, value.instance)
 
+        if isinstance(value, LazyReference):
+            child_descriptor = get_child_descriptor(value, display_context)
+            return create_node_input_value_pointer_rules(
+                child_descriptor, display_context, [], pointer_type=pointer_type
+            )
+
         if isinstance(value, CoalesceExpression):
             # Recursively handle the left-hand side
             lhs_rules = create_node_input_value_pointer_rules(value.lhs, display_context, [], pointer_type=pointer_type)

vellum_ee/workflows/display/vellum.py

@@ -41,10 +41,6 @@ class NodeDisplayData(UniversalBaseModel):
     comment: Optional[NodeDisplayComment] = None
 
 
-class GenericNodeDisplayData(UniversalBaseModel):
-    position: NodeDisplayPosition = Field(default_factory=NodeDisplayPosition)
-
-
 class CodeResourceDefinition(UniversalBaseModel):
     name: str
     module: List[str]

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -2,6 +2,8 @@ import pytest
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.nodes import BaseNodeDisplay
+from vellum_ee.workflows.display.vellum import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -63,3 +65,30 @@ def test_serialize_workflow__workflow_outputs_reference_non_node_outputs():
         == """Failed to serialize output 'final': Reference to outputs \
 'test_serialize_workflow__workflow_outputs_reference_non_node_outputs.<locals>.FirstWorkflow.Outputs' is invalid."""
     )
+
+
+def test_serialize_workflow__node_display_class_not_registered():
+    # GIVEN a workflow with a node that has a display class referencing display data
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class StartNodeDisplay(BaseNodeDisplay[StartNode]):
+        node_input_ids_by_name = {}
+        display_data = NodeDisplayData(position=NodeDisplayPosition(x=0, y=0), width=None, height=None)
+
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            answer = StartNode.Outputs.result
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=MyWorkflow,
+    )
+    data = workflow_display.serialize()
+
+    # THEN it should should succeed
+    assert data is not None