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

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

@@ -1,9 +1,8 @@
 from dataclasses import dataclass, field
-from uuid import UUID
 from typing import TYPE_CHECKING, Dict, Generic, Tuple, Type, TypeVar
 
-from vellum.client.core import UniversalBaseModel
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.events.workflow import WorkflowEventDisplayContext  # noqa: F401
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, StateValueReference, WorkflowInputReference
@@ -25,18 +24,6 @@ NodeDisplayType = TypeVar("NodeDisplayType", bound="BaseNodeDisplay")
 WorkflowDisplayType = TypeVar("WorkflowDisplayType", bound="BaseWorkflowDisplay")
 
 
-class NodeDisplay(UniversalBaseModel):
-    input_display: Dict[str, UUID]
-    output_display: Dict[str, UUID]
-    port_display: Dict[str, UUID]
-
-
-class WorkflowEventDisplayContext(UniversalBaseModel):
-    node_displays: Dict[str, NodeDisplay]
-    workflow_inputs: Dict[str, UUID]
-    workflow_outputs: Dict[str, UUID]
-
-
 @dataclass
 class WorkflowDisplayContext(
     Generic[

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -9,6 +9,7 @@ from typing import Any, Dict, Generic, Iterator, List, Optional, Tuple, Type, Un
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
+from vellum.workflows.events.workflow import NodeDisplay, WorkflowEventDisplayContext
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.utils import get_wrapped_node
@@ -33,12 +34,7 @@ from vellum_ee.workflows.display.base import (
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
-from vellum_ee.workflows.display.types import (
-    NodeDisplay,
-    NodeDisplayType,
-    WorkflowDisplayContext,
-    WorkflowEventDisplayContext,
-)
+from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext
 
 logger = logging.getLogger(__name__)
 
@@ -217,23 +213,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)
+
+            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
 
-            if node not in node_displays:
-                node_displays[node] = 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)
 
-            if node not in global_node_displays:
-                global_node_displays[node] = node_display
+        for node in self._workflow.get_unused_nodes():
+            extracted_node_displays = self._extract_node_displays(node)
 
-            # 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
+            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
 
-            self._enrich_global_node_output_displays(node, node_display, global_node_output_displays)
-            self._enrich_node_port_displays(node, node_display, port_displays)
+                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 +283,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 +421,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/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 = [

vellum_ee/workflows/tests/test_server.py

@@ -2,6 +2,7 @@ from vellum.client.core.pydantic_utilities import UniversalBaseModel
 
 
 def test_load_workflow_event_display_context():
+    # DEPRECATED: Use `vellum.workflows.events.workflow.WorkflowEventDisplayContext` instead. Will be removed in 0.15.0
     from vellum_ee.workflows.display.types import WorkflowEventDisplayContext
 
     # We are actually just ensuring there are no circular dependencies when