vellum-ai 0.14.6__py3-none-any.whl → 0.14.8__py3-none-any.whl

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"

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py

@@ -1,4 +1,3 @@
-import pytest
 from uuid import uuid4
 
 from deepdiff import DeepDiff
@@ -217,7 +216,6 @@ def test_serialize_node__try(serialize_node):
     )
 
 
-@pytest.mark.skip(reason="Not implemented")
 def test_serialize_node__stacked():
     @TryNode.wrap()
     @RetryNode.wrap(max_attempts=5)
@@ -236,4 +234,121 @@ def test_serialize_node__stacked():
     exec_config = workflow_display.serialize()
 
     # THEN the workflow display is created successfully
-    assert exec_config is not None
+    assert not DeepDiff(
+        {
+            "workflow_raw_data": {
+                "nodes": [
+                    {
+                        "id": "c14c1c9b-a7a4-4d2c-84fb-c940cfb09525",
+                        "type": "ENTRYPOINT",
+                        "inputs": [],
+                        "data": {
+                            "label": "Entrypoint Node",
+                            "source_handle_id": "51a5eb25-af14-4bee-9ced-d2aa534ea8e9",
+                        },
+                        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                        "base": None,
+                        "definition": None,
+                    },
+                    {
+                        "id": "074833b0-e142-4bbc-8dec-209a35e178a3",
+                        "label": "test_serialize_node__stacked.<locals>.InnerStackedGenericNode",
+                        "type": "GENERIC",
+                        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                        "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+                        "definition": {
+                            "name": "InnerStackedGenericNode",
+                            "module": [
+                                "vellum_ee",
+                                "workflows",
+                                "display",
+                                "tests",
+                                "workflow_serialization",
+                                "generic_nodes",
+                                "test_adornments_serialization",
+                            ],
+                        },
+                        "trigger": {"id": "f206358d-04a5-41c9-beee-0871a074fa48", "merge_behavior": "AWAIT_ATTRIBUTES"},
+                        "ports": [{"id": "408cd5fb-3a3e-4eb2-9889-61111bd6a129", "name": "default", "type": "DEFAULT"}],
+                        "adornments": [
+                            {
+                                "id": "5be7d260-74f7-4734-b31b-a46a94539586",
+                                "label": "RetryNode",
+                                "base": {
+                                    "name": "RetryNode",
+                                    "module": ["vellum", "workflows", "nodes", "core", "retry_node", "node"],
+                                },
+                                "attributes": [
+                                    {
+                                        "id": "c91782e3-140f-4938-9c23-d2a7b85dcdd8",
+                                        "name": "retry_on_error_code",
+                                        "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                                    },
+                                    {
+                                        "id": "f388e93b-8c68-4f54-8577-bbd0c9091557",
+                                        "name": "max_attempts",
+                                        "value": {"type": "CONSTANT_VALUE", "value": {"type": "NUMBER", "value": 5}},
+                                    },
+                                    {
+                                        "id": "8a07dc58-3fed-41d4-8ca6-31ee0bb86c61",
+                                        "name": "delay",
+                                        "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                                    },
+                                    {
+                                        "id": "73a02e62-4535-4e1f-97b5-1264ca8b1d71",
+                                        "name": "retry_on_condition",
+                                        "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                                    },
+                                ],
+                            },
+                            {
+                                "id": "3344083c-a32c-4a32-920b-0fb5093448fa",
+                                "label": "TryNode",
+                                "base": {
+                                    "name": "TryNode",
+                                    "module": ["vellum", "workflows", "nodes", "core", "try_node", "node"],
+                                },
+                                "attributes": [
+                                    {
+                                        "id": "ab2fbab0-e2a0-419b-b1ef-ce11ecf11e90",
+                                        "name": "on_error_code",
+                                        "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                                    }
+                                ],
+                            },
+                        ],
+                        "attributes": [],
+                        "outputs": [],
+                    },
+                ],
+                "edges": [
+                    {
+                        "id": "e8bd50dd-37a0-49b0-8b7b-f1dd8eb478b9",
+                        "source_node_id": "c14c1c9b-a7a4-4d2c-84fb-c940cfb09525",
+                        "source_handle_id": "51a5eb25-af14-4bee-9ced-d2aa534ea8e9",
+                        "target_node_id": "074833b0-e142-4bbc-8dec-209a35e178a3",
+                        "target_handle_id": "f206358d-04a5-41c9-beee-0871a074fa48",
+                        "type": "DEFAULT",
+                    }
+                ],
+                "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
+                "definition": {
+                    "name": "StackedWorkflow",
+                    "module": [
+                        "vellum_ee",
+                        "workflows",
+                        "display",
+                        "tests",
+                        "workflow_serialization",
+                        "generic_nodes",
+                        "test_adornments_serialization",
+                    ],
+                },
+            },
+            "input_variables": [],
+            "state_variables": [],
+            "output_variables": [],
+        },
+        exec_config,
+        ignore_order=True,
+    )

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/base_workflow_display.py

@@ -217,23 +217,30 @@ class BaseWorkflowDisplay(
         # TODO: We should still serialize nodes that are in the workflow's directory but aren't used in the graph.
         # https://app.shortcut.com/vellum/story/5394
         for node in self._workflow.get_nodes():
-            node_display = self._get_node_display(node)
+            extracted_node_displays = self._extract_node_displays(node)
 
-            if node not in node_displays:
-                node_displays[node] = node_display
+            for extracted_node, extracted_node_display in extracted_node_displays.items():
+                if extracted_node not in node_displays:
+                    node_displays[extracted_node] = extracted_node_display
 
-            if node not in global_node_displays:
-                global_node_displays[node] = node_display
+                if extracted_node not in global_node_displays:
+                    global_node_displays[extracted_node] = extracted_node_display
 
-            # Nodes wrapped in a decorator need to be in our node display dictionary for later retrieval
-            inner_node = get_wrapped_node(node)
-            if inner_node:
-                inner_node_display = self._get_node_display(inner_node)
-                node_displays[inner_node] = inner_node_display
-                global_node_displays[inner_node] = inner_node_display
+            self._enrich_global_node_output_displays(node, extracted_node_displays[node], global_node_output_displays)
+            self._enrich_node_port_displays(node, extracted_node_displays[node], port_displays)
 
-            self._enrich_global_node_output_displays(node, node_display, global_node_output_displays)
-            self._enrich_node_port_displays(node, node_display, port_displays)
+        for node in self._workflow.get_unused_nodes():
+            extracted_node_displays = self._extract_node_displays(node)
+
+            for extracted_node, extracted_node_display in extracted_node_displays.items():
+                if extracted_node not in node_displays:
+                    node_displays[extracted_node] = extracted_node_display
+
+                if extracted_node not in global_node_displays:
+                    global_node_displays[extracted_node] = extracted_node_display
+
+            self._enrich_global_node_output_displays(node, extracted_node_displays[node], global_node_output_displays)
+            self._enrich_node_port_displays(node, extracted_node_displays[node], port_displays)
 
         workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {}
         # If we're dealing with a nested workflow, then it should have access to the inputs of its parents.
@@ -280,6 +287,15 @@ class BaseWorkflowDisplay(
                 edge, node_displays, port_displays, overrides=edge_display_overrides
             )
 
+        for edge in self._workflow.get_unused_edges():
+            if edge in edge_displays:
+                continue
+
+            edge_display_overrides = self.edge_displays.get((edge.from_port, edge.to_node))
+            edge_displays[(edge.from_port, edge.to_node)] = self._generate_edge_display(
+                edge, node_displays, port_displays, overrides=edge_display_overrides
+            )
+
         workflow_output_displays: Dict[BaseDescriptor, WorkflowOutputDisplay] = {}
         for workflow_output in self._workflow.Outputs:
             if workflow_output in workflow_output_displays:
@@ -409,3 +425,20 @@ class BaseWorkflowDisplay(
             node_displays=temp_node_displays,
         )
         return display_meta
+
+    def _extract_node_displays(self, node: Type[BaseNode]) -> Dict[Type[BaseNode], NodeDisplayType]:
+        node_display = self._get_node_display(node)
+        additional_node_displays: Dict[Type[BaseNode], NodeDisplayType] = {
+            node: node_display,
+        }
+
+        # Nodes wrapped in a decorator need to be in our node display dictionary for later retrieval
+        inner_node = get_wrapped_node(node)
+        if inner_node:
+            inner_node_displays = self._extract_node_displays(inner_node)
+
+            for node, display in inner_node_displays.items():
+                if node not in additional_node_displays:
+                    additional_node_displays[node] = display
+
+        return additional_node_displays

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

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -126,6 +126,18 @@ class VellumWorkflowDisplay(
 
             nodes.append(serialized_node)
 
+        # Add all unused nodes in the workflow
+        for node in self._workflow.get_unused_nodes():
+            node_display = self.display_context.node_displays[node]
+
+            try:
+                serialized_node = node_display.serialize(self.display_context)
+            except NotImplementedError as e:
+                self.add_error(e)
+                continue
+
+            nodes.append(serialized_node)
+
         synthetic_output_edges: JsonArray = []
         output_variables: JsonArray = []
         final_output_nodes = [