vellum-ai 0.9.16rc2__py3-none-any.whl → 0.9.16rc4__py3-none-any.whl

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py

@@ -0,0 +1,372 @@
+from unittest import mock
+
+from deepdiff import DeepDiff
+
+from tests.workflows.basic_map_node.workflow import SimpleMapExample
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_serialize_workflow():
+    # GIVEN a Workflow that uses a MapNode
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleMapExample)
+
+    # TODO: Support serialization of BaseNode
+    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
+    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
+        mocked_serialize.return_value = {"type": "MOCKED"}
+        serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "db2eb237-38e4-417a-8bfc-5bda0f3165ca",
+                "key": "fruits",
+                "type": "JSON",
+                "required": None,
+                "default": None,
+                "extensions": None,
+            },
+        ],
+        input_variables,
+        ignore_order=True,
+    )
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "145b0b68-224b-4f83-90e6-eea3457e6c3e",
+                "key": "final_value",
+                "type": "JSON",
+            },
+        ],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    assert workflow_raw_data.keys() == {"edges", "nodes", "display_data", "definition"}
+    assert len(workflow_raw_data["edges"]) == 2
+    assert len(workflow_raw_data["nodes"]) == 3
+
+    # AND each node should be serialized correctly
+    entrypoint_node = workflow_raw_data["nodes"][0]
+    assert entrypoint_node == {
+        "id": "c0aa464d-1685-4f15-a051-31b426fec92e",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "data": {
+            "label": "Entrypoint Node",
+            "source_handle_id": "844d992e-60ab-4af2-a8ff-52cd858386f7",
+        },
+        "definition": {
+            "bases": [],
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "name": "BaseNode",
+        },
+        "display_data": {
+            "position": {"x": 0.0, "y": 0.0},
+        },
+    }
+
+    map_node = workflow_raw_data["nodes"][1]
+    assert not DeepDiff(
+        {
+            "id": "bf83099a-40df-4445-b90d-1f6f1067ebe3",
+            "type": "MAP",
+            "inputs": [
+                {
+                    "id": "4c0a109f-599e-4d04-a396-a51474fc2996",
+                    "key": "items",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "INPUT_VARIABLE",
+                                "data": {"input_variable_id": "db2eb237-38e4-417a-8bfc-5bda0f3165ca"},
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "data": {
+                "label": "Map Fruits Node",
+                "error_output_id": None,
+                "source_handle_id": "a2171a61-0657-43ad-b6d9-cf93ce3270d0",
+                "target_handle_id": "b5e8182e-20c5-482b-b4c5-4dde48c01472",
+                "variant": "INLINE",
+                "workflow_raw_data": {
+                    "nodes": [
+                        {
+                            "id": "ff9bfe6e-839d-4d40-b8fc-313b3bbd0ab0",
+                            "type": "ENTRYPOINT",
+                            "inputs": [],
+                            "data": {
+                                "label": "Entrypoint Node",
+                                "source_handle_id": "520d3616-8369-4e79-9da5-3febae299c2a",
+                            },
+                            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                            "definition": {
+                                "name": "BaseNode",
+                                "module": [
+                                    "vellum",
+                                    "workflows",
+                                    "nodes",
+                                    "bases",
+                                    "base",
+                                ],
+                                "bases": [],
+                            },
+                        },
+                        {"type": "MOCKED"},
+                        {
+                            "id": "6f4883b2-70b1-4e1c-ae15-7d0f5aec810b",
+                            "type": "TERMINAL",
+                            "data": {
+                                "label": "Final Output",
+                                "name": "count",
+                                "target_handle_id": "9d74571f-b7f5-4c1d-8b7c-b9c648738a4d",
+                                "output_id": "2a957315-fae0-4366-8a35-f0b315c5eade",
+                                "output_type": "NUMBER",
+                                "node_input_id": "ae65cf7e-3db6-410f-a556-75ee11ce7e84",
+                            },
+                            "inputs": [
+                                {
+                                    "id": "ae65cf7e-3db6-410f-a556-75ee11ce7e84",
+                                    "key": "node_input",
+                                    "value": {
+                                        "rules": [
+                                            {
+                                                "type": "NODE_OUTPUT",
+                                                "data": {
+                                                    "node_id": "baf6d316-dc75-41e8-96c0-015aede96309",
+                                                    "output_id": "a7bcb362-a2b8-4476-b0de-a361efeec204",
+                                                },
+                                            }
+                                        ],
+                                        "combinator": "OR",
+                                    },
+                                }
+                            ],
+                            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                            "definition": {
+                                "name": "FinalOutputNode",
+                                "module": [
+                                    "vellum",
+                                    "workflows",
+                                    "nodes",
+                                    "displayable",
+                                    "final_output_node",
+                                    "node",
+                                ],
+                                "bases": [
+                                    {
+                                        "name": "BaseNode",
+                                        "module": [
+                                            "vellum",
+                                            "workflows",
+                                            "nodes",
+                                            "bases",
+                                            "base",
+                                        ],
+                                        "bases": [],
+                                    }
+                                ],
+                            },
+                        },
+                    ],
+                    "edges": [
+                        {
+                            "id": "b59c050b-0f6a-4153-ab58-fa051222fa05",
+                            "source_node_id": "ff9bfe6e-839d-4d40-b8fc-313b3bbd0ab0",
+                            "source_handle_id": "520d3616-8369-4e79-9da5-3febae299c2a",
+                            "target_node_id": "baf6d316-dc75-41e8-96c0-015aede96309",
+                            "target_handle_id": "551d5528-f4e1-42ea-bde0-9de4b4968253",
+                            "type": "DEFAULT",
+                        },
+                        {
+                            "id": "14152688-6996-4d64-9231-a6e66a4827eb",
+                            "source_node_id": "baf6d316-dc75-41e8-96c0-015aede96309",
+                            "source_handle_id": "71ada606-d791-4a59-a252-0795c5faeeaf",
+                            "target_node_id": "6f4883b2-70b1-4e1c-ae15-7d0f5aec810b",
+                            "target_handle_id": "9d74571f-b7f5-4c1d-8b7c-b9c648738a4d",
+                            "type": "DEFAULT",
+                        },
+                    ],
+                    "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
+                    "definition": {
+                        "name": "IterationSubworkflow",
+                        "module": [
+                            "tests",
+                            "workflows",
+                            "basic_map_node",
+                            "workflow",
+                        ],
+                    },
+                },
+                "input_variables": [
+                    {"id": "0fe9dff5-7595-4ff3-a420-416f3fc6f809", "key": "item", "type": "JSON"},
+                    {"id": "27dfe8ba-8be7-4de9-a930-2ec3e375b149", "key": "index", "type": "NUMBER"},
+                    {"id": "4c0a109f-599e-4d04-a396-a51474fc2996", "key": "items", "type": "JSON"},
+                ],
+                "output_variables": [{"id": "56b34d4a-d3a6-4bdb-80a4-dbf644791274", "key": "count", "type": "NUMBER"}],
+                "concurrency": None,
+                "items_input_id": "4c0a109f-599e-4d04-a396-a51474fc2996",
+                "item_input_id": "0fe9dff5-7595-4ff3-a420-416f3fc6f809",
+                "index_input_id": "27dfe8ba-8be7-4de9-a930-2ec3e375b149",
+            },
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "MapFruitsNode",
+                "module": [
+                    "tests",
+                    "workflows",
+                    "basic_map_node",
+                    "workflow",
+                ],
+                "bases": [
+                    {
+                        "name": "MapNode",
+                        "module": [
+                            "vellum",
+                            "workflows",
+                            "nodes",
+                            "core",
+                            "map_node",
+                            "node",
+                        ],
+                    }
+                ],
+            },
+        },
+        map_node,
+        ignore_order=True,
+    )
+
+    assert not DeepDiff(
+        {
+            "id": "bacc5d55-07d4-4a0a-a69e-831524480de5",
+            "type": "TERMINAL",
+            "data": {
+                "label": "Final Output",
+                "name": "final_value",
+                "target_handle_id": "720dd872-2f3d-47b9-8245-89387f04f300",
+                "output_id": "145b0b68-224b-4f83-90e6-eea3457e6c3e",
+                "output_type": "JSON",
+                "node_input_id": "d1b01eac-23a9-43ce-beb3-e13f83bd7d18",
+            },
+            "inputs": [
+                {
+                    "id": "d1b01eac-23a9-43ce-beb3-e13f83bd7d18",
+                    "key": "node_input",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "NODE_OUTPUT",
+                                "data": {
+                                    "node_id": "bf83099a-40df-4445-b90d-1f6f1067ebe3",
+                                    "output_id": "56b34d4a-d3a6-4bdb-80a4-dbf644791274",
+                                },
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "bases": [
+                    {
+                        "bases": [],
+                        "module": [
+                            "vellum",
+                            "workflows",
+                            "nodes",
+                            "bases",
+                            "base",
+                        ],
+                        "name": "BaseNode",
+                    },
+                ],
+                "module": [
+                    "vellum",
+                    "workflows",
+                    "nodes",
+                    "displayable",
+                    "final_output_node",
+                    "node",
+                ],
+                "name": "FinalOutputNode",
+            },
+        },
+        workflow_raw_data["nodes"][2],
+    )
+
+    # AND each edge should be serialized correctly
+    serialized_edges = workflow_raw_data["edges"]
+    assert not DeepDiff(
+        [
+            {
+                "id": "f39477f3-445a-412e-a8ab-371baa8ae990",
+                "source_node_id": "c0aa464d-1685-4f15-a051-31b426fec92e",
+                "source_handle_id": "844d992e-60ab-4af2-a8ff-52cd858386f7",
+                "target_node_id": "bf83099a-40df-4445-b90d-1f6f1067ebe3",
+                "target_handle_id": "b5e8182e-20c5-482b-b4c5-4dde48c01472",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "47a34f6e-d139-4702-aa46-6212bb8a150f",
+                "source_node_id": "bf83099a-40df-4445-b90d-1f6f1067ebe3",
+                "source_handle_id": "a2171a61-0657-43ad-b6d9-cf93ce3270d0",
+                "target_node_id": "bacc5d55-07d4-4a0a-a69e-831524480de5",
+                "target_handle_id": "720dd872-2f3d-47b9-8245-89387f04f300",
+                "type": "DEFAULT",
+            },
+        ],
+        serialized_edges,
+        ignore_order=True,
+    )
+
+    # AND the display data should be what we expect
+    display_data = workflow_raw_data["display_data"]
+    assert display_data == {
+        "viewport": {
+            "x": 0.0,
+            "y": 0.0,
+            "zoom": 1.0,
+        }
+    }
+
+    # AND the definition should be what we expect
+    definition = workflow_raw_data["definition"]
+    assert definition == {
+        "name": "SimpleMapExample",
+        "module": [
+            "tests",
+            "workflows",
+            "basic_map_node",
+            "workflow",
+        ],
+    }

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py

@@ -0,0 +1,272 @@
+from unittest import mock
+
+from deepdiff import DeepDiff
+
+from tests.workflows.basic_merge_node.await_all_workflow import AwaitAllPassingWorkflow
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_serialize_workflow__await_all():
+    # GIVEN a Workflow that uses an await all merge node
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=AwaitAllPassingWorkflow
+    )
+
+    # TODO: Support serialization of BaseNode
+    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
+    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
+        mocked_serialize.return_value = {"type": "MOCKED"}
+        serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 0
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 1
+    assert not DeepDiff(
+        [
+            {"id": "959ba00d-d30b-402e-8676-76efc4c3d2ae", "key": "value", "type": "STRING"},
+        ],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    assert workflow_raw_data.keys() == {"edges", "nodes", "display_data", "definition"}
+    assert len(workflow_raw_data["edges"]) == 6
+    assert len(workflow_raw_data["nodes"]) == 6
+
+    # AND each node should be serialized correctly
+    entrypoint_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "ENTRYPOINT")
+    assert entrypoint_node == {
+        "id": "dc8aecd0-49ba-4464-a45f-29d3bfd686e4",
+        "type": "ENTRYPOINT",
+        "definition": {
+            "bases": [],
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "name": "BaseNode",
+        },
+        "inputs": [],
+        "data": {
+            "label": "Entrypoint Node",
+            "source_handle_id": "017d40f5-8326-4e42-a409-b08995defaa8",
+        },
+        "display_data": {
+            "position": {"x": 0.0, "y": 0.0},
+        },
+    }
+
+    passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "MOCKED"]
+    assert not DeepDiff(
+        [
+            {
+                "type": "MOCKED",
+            },
+            {
+                "type": "MOCKED",
+            },
+            {
+                "type": "MOCKED",
+            },
+        ],
+        passthrough_nodes,
+        ignore_order=True,
+    )
+
+    merge_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "MERGE")
+    assert not DeepDiff(
+        {
+            "id": "37c10e8a-771b-432b-a767-31f5007851f0",
+            "type": "MERGE",
+            "inputs": [],
+            "data": {
+                "label": "Await All Merge Node",
+                "merge_strategy": "AWAIT_ALL",
+                "target_handles": [
+                    {"id": "f40ff7fb-de1b-4aa4-ba3c-7630f7357cbf"},
+                    {"id": "42eeb66c-9792-4609-8c71-3a56f668f4dc"},
+                ],
+                "source_handle_id": "3bbc469f-0fb0-4b3d-a28b-746fefec2818",
+            },
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "bases": [
+                    {
+                        "module": [
+                            "vellum",
+                            "workflows",
+                            "nodes",
+                            "displayable",
+                            "merge_node",
+                            "node",
+                        ],
+                        "name": "MergeNode",
+                    }
+                ],
+                "module": [
+                    "tests",
+                    "workflows",
+                    "basic_merge_node",
+                    "await_all_workflow",
+                ],
+                "name": "AwaitAllMergeNode",
+            },
+        },
+        merge_node,
+        ignore_order_func=lambda x: x.path() == "root['data']['target_handles']",
+    )
+
+    final_output_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "TERMINAL")
+    assert final_output_node == {
+        "id": "8187ce10-62b7-4a2c-8c0f-297387915467",
+        "type": "TERMINAL",
+        "data": {
+            "label": "Final Output",
+            "name": "value",
+            "target_handle_id": "ff55701c-16d3-4348-a633-6a298e71377d",
+            "output_id": "959ba00d-d30b-402e-8676-76efc4c3d2ae",
+            "output_type": "STRING",
+            "node_input_id": "7f950be4-2fab-44e0-87a3-b1631aadd0e3",
+        },
+        "definition": {
+            "bases": [
+                {
+                    "bases": [],
+                    "module": [
+                        "vellum",
+                        "workflows",
+                        "nodes",
+                        "bases",
+                        "base",
+                    ],
+                    "name": "BaseNode",
+                },
+            ],
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "displayable",
+                "final_output_node",
+                "node",
+            ],
+            "name": "FinalOutputNode",
+        },
+        "inputs": [
+            {
+                "id": "7f950be4-2fab-44e0-87a3-b1631aadd0e3",
+                "key": "node_input",
+                "value": {
+                    "rules": [
+                        {
+                            "type": "NODE_OUTPUT",
+                            "data": {
+                                "node_id": "634f0202-9ea9-4c62-b152-1a58c595cffb",
+                                "output_id": "d4266640-9718-4a74-b24b-500448d87871",
+                            },
+                        }
+                    ],
+                    "combinator": "OR",
+                },
+            }
+        ],
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+    }
+
+    # AND each edge should be serialized correctly
+    serialized_edges = workflow_raw_data["edges"]
+    assert not DeepDiff(
+        [
+            {
+                "id": "9a65dd52-c3eb-496e-9d34-46b39534a261",
+                "source_node_id": "dc8aecd0-49ba-4464-a45f-29d3bfd686e4",
+                "source_handle_id": "017d40f5-8326-4e42-a409-b08995defaa8",
+                "target_node_id": "59243c65-053f-4ea6-9157-3f3edb1477bf",
+                "target_handle_id": "e622fe61-3bca-4aff-86e1-25dad7bdf9d4",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "e5598f3c-fb00-4f25-a0a6-9fb6af9d69a8",
+                "source_node_id": "dc8aecd0-49ba-4464-a45f-29d3bfd686e4",
+                "source_handle_id": "017d40f5-8326-4e42-a409-b08995defaa8",
+                "target_node_id": "127ef456-91bc-43c6-bd8b-1772db5e3cb5",
+                "target_handle_id": "e5cc41cb-71db-43ec-b3f0-c78706af3351",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "8ff20817-974e-4a3a-bb65-f0ad73557649",
+                "source_node_id": "59243c65-053f-4ea6-9157-3f3edb1477bf",
+                "source_handle_id": "b9c5f52b-b714-46e8-a09c-38b4e770dd36",
+                "target_node_id": "37c10e8a-771b-432b-a767-31f5007851f0",
+                "target_handle_id": "0efd256f-f5f6-45fe-9adb-651780f5e63d",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "0d8c801c-d76a-437a-831a-530885b75f96",
+                "source_node_id": "127ef456-91bc-43c6-bd8b-1772db5e3cb5",
+                "source_handle_id": "b0bd17f3-4ce6-4232-9666-ec8afa161bf2",
+                "target_node_id": "37c10e8a-771b-432b-a767-31f5007851f0",
+                "target_handle_id": "0efd256f-f5f6-45fe-9adb-651780f5e63d",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "70c1005d-339a-41bc-b6c2-10bc30a0281c",
+                "source_node_id": "37c10e8a-771b-432b-a767-31f5007851f0",
+                "source_handle_id": "3bbc469f-0fb0-4b3d-a28b-746fefec2818",
+                "target_node_id": "634f0202-9ea9-4c62-b152-1a58c595cffb",
+                "target_handle_id": "acd48f48-54fb-4b2b-ab37-96d336f6dfb3",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "3d031c93-09b1-4937-9f98-c30a7ba6823d",
+                "source_node_id": "634f0202-9ea9-4c62-b152-1a58c595cffb",
+                "source_handle_id": "de32167b-cf53-4df5-a344-1b9be852e9ff",
+                "target_node_id": "8187ce10-62b7-4a2c-8c0f-297387915467",
+                "target_handle_id": "ff55701c-16d3-4348-a633-6a298e71377d",
+                "type": "DEFAULT",
+            },
+        ],
+        serialized_edges,
+        ignore_order=True,
+    )
+
+    # AND the display data should be what we expect
+    display_data = workflow_raw_data["display_data"]
+    assert display_data == {
+        "viewport": {
+            "x": 0.0,
+            "y": 0.0,
+            "zoom": 1.0,
+        }
+    }
+
+    # AND the definition should be what we expect
+    definition = workflow_raw_data["definition"]
+    assert definition == {
+        "name": "AwaitAllPassingWorkflow",
+        "module": [
+            "tests",
+            "workflows",
+            "basic_merge_node",
+            "await_all_workflow",
+        ],
+    }