vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

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

@@ -1,4 +1,5 @@
 from uuid import uuid4
+from typing import Any, Dict, List, cast
 
 from deepdiff import DeepDiff
 
@@ -45,8 +46,8 @@ def test_serialize_node__retry(serialize_node):
     )
     assert not DeepDiff(
         {
-            "id": "188b50aa-e518-4b7b-a5e0-e2585fb1d7b5",
-            "label": "test_serialize_node__retry.<locals>.InnerRetryGenericNode",
+            "id": "93572a37-c28e-41ec-8ed0-328200d3fc5e",
+            "label": "Inner Retry Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -62,12 +63,12 @@ def test_serialize_node__retry(serialize_node):
                     "test_adornments_serialization",
                 ],
             },
-            "trigger": {"id": "75fbe874-c00b-4fc2-9ade-52f4fe9209fa", "merge_behavior": "AWAIT_ATTRIBUTES"},
-            "ports": [{"id": "078650c9-f775-4cd0-a08c-23af9983a361", "name": "default", "type": "DEFAULT"}],
+            "trigger": {"id": "f642862d-2f8a-47b1-b23f-21f26ccfb2cc", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "66374ff9-7b8a-4255-ab48-8971cf4101ba", "name": "default", "type": "DEFAULT"}],
             "adornments": [
                 {
                     "id": "5be7d260-74f7-4734-b31b-a46a94539586",
-                    "label": "RetryNode",
+                    "label": "Retry Node",
                     "base": {
                         "name": "RetryNode",
                         "module": ["vellum", "workflows", "nodes", "core", "retry_node", "node"],
@@ -98,13 +99,13 @@ def test_serialize_node__retry(serialize_node):
             ],
             "attributes": [
                 {
-                    "id": "278df25e-58b5-43c3-b346-cf6444d893a5",
+                    "id": "016fa09c-3b6f-49c3-a177-5f1bb1afbeb2",
                     "name": "input",
                     "value": {"type": "WORKFLOW_INPUT", "input_variable_id": str(input_id)},
                 }
             ],
             "outputs": [
-                {"id": "dc89dc0d-c0bd-47fd-88aa-ec7b262aa2f1", "name": "output", "type": "STRING", "value": None}
+                {"id": "78eece53-8a20-40a1-8a86-ffebe256282b", "name": "output", "type": "STRING", "value": None}
             ],
         },
         serialized_node,
@@ -153,7 +154,7 @@ def test_serialize_node__try(serialize_node):
     assert not DeepDiff(
         {
             "id": str(InnerTryGenericNode.__wrapped_node__.__id__),
-            "label": "test_serialize_node__try.<locals>.InnerTryGenericNode",
+            "label": "Inner Try Generic Node",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -169,12 +170,12 @@ def test_serialize_node__try(serialize_node):
                     "test_adornments_serialization",
                 ],
             },
-            "trigger": {"id": "bbb343ff-2b7a-4793-a8cf-fb05132ca46a", "merge_behavior": "AWAIT_ATTRIBUTES"},
-            "ports": [{"id": "8d25f244-4b12-4f8b-b202-8948698679a0", "name": "default", "type": "DEFAULT"}],
+            "trigger": {"id": "596c162f-98fd-499a-a2ff-aa576cd8458a", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "ee8c7891-0125-4f5c-b5b0-6e3f51bc0e09", "name": "default", "type": "DEFAULT"}],
             "adornments": [
                 {
                     "id": "3344083c-a32c-4a32-920b-0fb5093448fa",
-                    "label": "TryNode",
+                    "label": "Try Node",
                     "base": {
                         "name": "TryNode",
                         "module": ["vellum", "workflows", "nodes", "core", "try_node", "node"],
@@ -190,13 +191,13 @@ def test_serialize_node__try(serialize_node):
             ],
             "attributes": [
                 {
-                    "id": "51aa0077-4060-496b-8e2e-e79d56ee6a32",
+                    "id": "66a1a015-948e-4b9d-8746-f766fa70a445",
                     "name": "input",
                     "value": {"type": "WORKFLOW_INPUT", "input_variable_id": str(input_id)},
                 }
             ],
             "outputs": [
-                {"id": "ce9f8b86-6d26-4c03-8bfa-a31aa2cd97f1", "name": "output", "type": "STRING", "value": None}
+                {"id": "d8d0c9a8-0804-4b43-a874-28a7e7d6aec8", "name": "output", "type": "STRING", "value": None}
             ],
         },
         serialized_node,
@@ -237,124 +238,148 @@ def test_serialize_node__stacked():
     exec_config = workflow_display.serialize()
 
     # THEN the workflow display is created successfully
+    assert isinstance(exec_config["workflow_raw_data"], dict)
+    assert isinstance(exec_config["workflow_raw_data"]["nodes"], list)
+    inner_stacked_generic_node = [
+        node
+        for node in exec_config["workflow_raw_data"]["nodes"]
+        if isinstance(node, dict) and node["type"] == "GENERIC"
+    ][0]
     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": "957df57c-42bc-44af-9a31-ce7ac3bb1b8a",
+            "label": "Inner Stacked Generic Node",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 200.0, "y": -50.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": "65595aa4-f5d7-4462-8074-2bd8eb9dc6ee",
+                "merge_behavior": "AWAIT_ATTRIBUTES",
+            },
+            "ports": [{"id": "aebdab62-838e-427a-9de9-301c98472bd4", "name": "default", "type": "DEFAULT"}],
+            "adornments": [
+                {
+                    "id": "2e22c747-9b17-4029-b2ee-e22e22056e1f",
+                    "label": "Try Node",
+                    "base": {
+                        "name": "TryNode",
+                        "module": ["vellum", "workflows", "nodes", "core", "try_node", "node"],
                     },
-                    {
-                        "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",
-                            ],
+                    "attributes": [
+                        {
+                            "id": "205b2e6c-1f49-459c-8ed6-58e1b012e5e7",
+                            "name": "on_error_code",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        }
+                    ],
+                },
+                {
+                    "id": "3149e684-69cf-4909-ae68-c14d101b689f",
+                    "label": "Retry Node",
+                    "base": {
+                        "name": "RetryNode",
+                        "module": ["vellum", "workflows", "nodes", "core", "retry_node", "node"],
+                    },
+                    "attributes": [
+                        {
+                            "id": "0b10be3f-c221-4f92-9930-7a382390963f",
+                            "name": "max_attempts",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "NUMBER", "value": 5.0}},
                         },
-                        "trigger": {
-                            "id": "6e4af17f-bbee-4777-b10d-af042cd6e16a",
-                            "merge_behavior": "AWAIT_ATTRIBUTES",
+                        {
+                            "id": "24f2a261-45fc-41bb-b700-e2346a0b785b",
+                            "name": "delay",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
+                        {
+                            "id": "64c268c4-3fed-4a9f-b22b-959f7921e165",
+                            "name": "retry_on_error_code",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
+                        {
+                            "id": "d17f61e6-0b43-45e4-b204-39f82a562fee",
+                            "name": "retry_on_condition",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
                         },
-                        "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": "6e4af17f-bbee-4777-b10d-af042cd6e16a",
-                        "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": [],
+            ],
+            "attributes": [],
+            "outputs": [],
         },
-        exec_config,
-        ignore_order=True,
+        inner_stacked_generic_node,
     )
+
+
+def test_serialize_node__adornment_order_matches_decorator_order():
+    """
+    Tests that adornments are serialized in the same order as decorators are applied.
+    """
+
+    @TryNode.wrap()
+    @RetryNode.wrap(max_attempts=3)
+    class MyNode(BaseNode):
+        pass
+
+    # AND a workflow that uses the decorated node
+    class MyWorkflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    exec_config = cast(Dict[str, Any], workflow_display.serialize())
+
+    # THEN the workflow should serialize successfully
+    assert isinstance(exec_config["workflow_raw_data"], dict)
+    assert isinstance(exec_config["workflow_raw_data"]["nodes"], list)
+
+    # AND we should find our decorated node
+    nodes = cast(List[Dict[str, Any]], exec_config["workflow_raw_data"]["nodes"])
+    my_node = [node for node in nodes if isinstance(node, dict) and node["type"] == "GENERIC"][0]
+
+    adornments = cast(List[Dict[str, Any]], my_node["adornments"])
+    assert len(adornments) == 2
+    assert adornments[0]["label"] == "Try Node"
+    assert adornments[1]["label"] == "Retry Node"
+
+
+def test_serialize_workflow__retry_node_edges():
+    """
+    Tests that both retry-adorned nodes are correctly serialized in the nodes array.
+    """
+
+    @RetryNode.wrap(max_attempts=3, delay=60)
+    class FirstNode(BaseNode):
+        class Outputs(BaseOutputs):
+            value: str
+
+    @RetryNode.wrap(max_attempts=5, delay=120)
+    class SecondNode(BaseNode):
+        class Outputs(BaseOutputs):
+            result: str
+
+    class MyWorkflow(BaseWorkflow):
+        graph = FirstNode >> SecondNode
+
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    exec_config = cast(Dict[str, Any], workflow_display.serialize())
+
+    assert isinstance(exec_config["workflow_raw_data"], dict)
+    assert isinstance(exec_config["workflow_raw_data"]["nodes"], list)
+
+    nodes = cast(List[Dict[str, Any]], exec_config["workflow_raw_data"]["nodes"])
+
+    generic_nodes = [node for node in nodes if node["type"] == "GENERIC"]
+    assert len(generic_nodes) == 2