vellum-ai 0.14.68__py3-none-any.whl → 0.14.70__py3-none-any.whl

vellum_ee/workflows/display/utils/tests/test_auto_layout.py

@@ -0,0 +1,56 @@
+from uuid import uuid4
+
+from vellum_ee.workflows.display.base import EdgeDisplay
+from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
+from vellum_ee.workflows.display.utils.auto_layout import auto_layout_nodes
+
+
+def test_auto_layout_basic():
+    """Test basic auto layout functionality with a simple linear graph."""
+    nodes = [
+        ("node1", NodeDisplayData(position=NodeDisplayPosition(x=0, y=0), width=100, height=80)),
+        ("node2", NodeDisplayData(position=NodeDisplayPosition(x=0, y=0), width=120, height=90)),
+        ("node3", NodeDisplayData(position=NodeDisplayPosition(x=0, y=0), width=110, height=85)),
+    ]
+
+    edges = [
+        ("node1", "node2", EdgeDisplay(id=uuid4())),
+        ("node2", "node3", EdgeDisplay(id=uuid4())),
+    ]
+
+    positioned_nodes = auto_layout_nodes(nodes, edges)
+
+    node_positions = {node_id: data.position for node_id, data in positioned_nodes}
+
+    assert node_positions["node1"].x < node_positions["node2"].x
+    assert node_positions["node2"].x < node_positions["node3"].x
+
+
+def test_auto_layout_parallel_branches():
+    """Test auto layout with parallel branches."""
+    nodes = [
+        ("start", NodeDisplayData(position=NodeDisplayPosition(x=0, y=0), width=100, height=80)),
+        ("branch1", NodeDisplayData(position=NodeDisplayPosition(x=0, y=0), width=100, height=80)),
+        ("branch2", NodeDisplayData(position=NodeDisplayPosition(x=0, y=0), width=100, height=80)),
+        ("end", NodeDisplayData(position=NodeDisplayPosition(x=0, y=0), width=100, height=80)),
+    ]
+
+    edges = [
+        ("start", "branch1", EdgeDisplay(id=uuid4())),
+        ("start", "branch2", EdgeDisplay(id=uuid4())),
+        ("branch1", "end", EdgeDisplay(id=uuid4())),
+        ("branch2", "end", EdgeDisplay(id=uuid4())),
+    ]
+
+    positioned_nodes = auto_layout_nodes(nodes, edges)
+
+    node_positions = {node_id: data.position for node_id, data in positioned_nodes}
+
+    assert node_positions["start"].x < node_positions["branch1"].x
+    assert node_positions["start"].x < node_positions["branch2"].x
+
+    assert node_positions["branch1"].x == node_positions["branch2"].x
+    assert node_positions["branch1"].x < node_positions["end"].x
+
+    assert node_positions["branch1"].x < node_positions["end"].x
+    assert node_positions["branch2"].x < node_positions["end"].x

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -200,16 +200,21 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                     raise ValueError(f"Failed to serialize output '{workflow_output.name}': {str(e)}") from e
 
                 source_node_display: Optional[BaseNodeDisplay]
-                first_rule = node_input.value.rules[0]
-                if first_rule.type == "NODE_OUTPUT":
-                    source_node_id = UUID(first_rule.data.node_id)
-                    try:
-                        source_node_display = [
-                            node_display
-                            for node_display in self.display_context.node_displays.values()
-                            if node_display.node_id == source_node_id
-                        ][0]
-                    except IndexError:
+                if not node_input.value.rules:
+                    source_node_display = None
+                else:
+                    first_rule = node_input.value.rules[0]
+                    if first_rule.type == "NODE_OUTPUT":
+                        source_node_id = UUID(first_rule.data.node_id)
+                        try:
+                            source_node_display = [
+                                node_display
+                                for node_display in self.display_context.node_displays.values()
+                                if node_display.node_id == source_node_id
+                            ][0]
+                        except IndexError:
+                            source_node_display = None
+                    else:
                         source_node_display = None
 
                 synthetic_target_handle_id = str(

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

@@ -787,3 +787,44 @@ def test_serialize_workflow__dict_reference():
             },
         ],
     }
+
+
+def test_serialize_workflow__empty_rules_indexerror():
+    """Test that workflow serialization handles dictionary key access correctly."""
+
+    # GIVEN a node with dictionary output
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            data: dict = {"key": "value"}
+
+    # AND a workflow that references dictionary key access in its outputs
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            # This dictionary key access should be handled gracefully
+            problematic_output = StartNode.Outputs.data["bar"]
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    result: dict = workflow_display.serialize()
+
+    assert result is not None
+    assert "output_variables" in result
+    assert "workflow_raw_data" in result
+
+    # AND the workflow output should contain the dictionary key access
+    output_variables = result["output_variables"]
+    assert len(output_variables) == 1
+    assert output_variables[0]["key"] == "problematic_output"
+
+    # AND the workflow raw data should contain nodes including terminal node
+    workflow_raw_data = result["workflow_raw_data"]
+    assert "nodes" in workflow_raw_data
+    nodes = workflow_raw_data["nodes"]
+
+    assert len(nodes) >= 3
+
+    terminal_nodes = [node for node in nodes if node.get("type") == "TERMINAL"]
+    assert len(terminal_nodes) == 1
+    assert terminal_nodes[0]["data"]["name"] == "problematic_output"