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

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

@@ -40,22 +40,12 @@ def test_serialize_workflow_with_filepath():
 
     # AND its raw data should be what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    assert len(workflow_raw_data["edges"]) == 3
-    assert len(workflow_raw_data["nodes"]) == 4
 
     # AND each node should be serialized correctly
-    entrypoint_node = workflow_raw_data["nodes"][0]
-    assert entrypoint_node == {
-        "id": "bd18f11c-5f7a-45d5-9970-0b1cf10d3761",
-        "type": "ENTRYPOINT",
-        "inputs": [],
-        "data": {"label": "Entrypoint Node", "source_handle_id": "118e4298-aa79-467c-b8b4-2df540905e86"},
-        "display_data": {"position": {"x": 0.0, "y": -50.0}},
-        "base": None,
-        "definition": None,
-    }
 
-    code_execution_node = workflow_raw_data["nodes"][1]
+    code_execution_node = next(
+        n for n in workflow_raw_data["nodes"] if (n.get("base") or {}).get("name") == "CodeExecutionNode"
+    )
     assert code_execution_node == {
         "id": "b4a0526b-0e31-4f63-892a-6d5197d09acf",
         "type": "CODE_EXECUTION",
@@ -115,127 +105,22 @@ def test_serialize_workflow_with_filepath():
             "merge_behavior": "AWAIT_ANY",
         },
         "ports": [{"id": "6aed9e19-9d26-457b-966d-0a9112f84070", "name": "default", "type": "DEFAULT"}],
-    }
-    assert not DeepDiff(
-        [
+        "outputs": [
             {
-                "id": "994d5c2e-00d2-4dff-9a9d-804766d03698",
-                "type": "TERMINAL",
-                "data": {
-                    "label": "Final Output",
-                    "name": "result",
-                    "target_handle_id": "30fb0f4a-61c3-49de-a0aa-7dfdcee6ea07",
-                    "output_id": "f6a3e3e0-f83f-4491-8b7a-b20fddd7160c",
-                    "output_type": "NUMBER",
-                    "node_input_id": "529bdd20-0985-4c99-87dc-590907937c1d",
-                },
-                "inputs": [
-                    {
-                        "id": "529bdd20-0985-4c99-87dc-590907937c1d",
-                        "key": "node_input",
-                        "value": {
-                            "rules": [
-                                {
-                                    "type": "NODE_OUTPUT",
-                                    "data": {
-                                        "node_id": "b4a0526b-0e31-4f63-892a-6d5197d09acf",
-                                        "output_id": "730c8f17-891c-40b4-b43d-26672bd38eef",
-                                    },
-                                }
-                            ],
-                            "combinator": "OR",
-                        },
-                    }
-                ],
-                "display_data": {"position": {"x": 400.0, "y": -175.0}},
-                "base": {
-                    "name": "FinalOutputNode",
-                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
-                },
-                "definition": None,
+                "id": "730c8f17-891c-40b4-b43d-26672bd38eef",
+                "name": "result",
+                "schema": {"type": "integer"},
+                "type": "NUMBER",
+                "value": None,
             },
             {
-                "id": "c6e3aced-1fc9-48d2-ae55-d2a880e359cb",
-                "type": "TERMINAL",
-                "data": {
-                    "label": "Final Output",
-                    "name": "log",
-                    "target_handle_id": "1e126004-9de7-42c0-b1e1-87f9eb0642e2",
-                    "output_id": "1cee930f-342f-421c-89fc-ff212b3764bb",
-                    "output_type": "STRING",
-                    "node_input_id": "09501b65-d9b3-4920-81d4-96f93c840667",
-                },
-                "inputs": [
-                    {
-                        "id": "09501b65-d9b3-4920-81d4-96f93c840667",
-                        "key": "node_input",
-                        "value": {
-                            "rules": [
-                                {
-                                    "type": "NODE_OUTPUT",
-                                    "data": {
-                                        "node_id": "b4a0526b-0e31-4f63-892a-6d5197d09acf",
-                                        "output_id": "8eea2893-1e79-4d1c-ba51-2e045968abfb",
-                                    },
-                                }
-                            ],
-                            "combinator": "OR",
-                        },
-                    }
-                ],
-                "display_data": {"position": {"x": 400.0, "y": 75.0}},
-                "base": {
-                    "name": "FinalOutputNode",
-                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
-                },
-                "definition": None,
+                "id": "8eea2893-1e79-4d1c-ba51-2e045968abfb",
+                "name": "log",
+                "schema": {"type": "string"},
+                "type": "STRING",
+                "value": None,
             },
         ],
-        workflow_raw_data["nodes"][2:],
-        ignore_order=True,
-    )
-
-    # AND each edge should be serialized correctly
-    serialized_edges = workflow_raw_data["edges"]
-    assert not DeepDiff(
-        [
-            {
-                "id": "48e2c971-b780-436a-873a-3f775fc04ab3",
-                "source_node_id": "bd18f11c-5f7a-45d5-9970-0b1cf10d3761",
-                "source_handle_id": "118e4298-aa79-467c-b8b4-2df540905e86",
-                "target_node_id": "b4a0526b-0e31-4f63-892a-6d5197d09acf",
-                "target_handle_id": "f1ea9f65-e225-49cc-a779-6bd0797ba22a",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "d1e66711-75b3-41c3-beb6-424894fdd307",
-                "source_node_id": "b4a0526b-0e31-4f63-892a-6d5197d09acf",
-                "source_handle_id": "6aed9e19-9d26-457b-966d-0a9112f84070",
-                "target_node_id": "994d5c2e-00d2-4dff-9a9d-804766d03698",
-                "target_handle_id": "30fb0f4a-61c3-49de-a0aa-7dfdcee6ea07",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "67d4c43e-80f9-4875-b6ab-9ecbba19fc7a",
-                "source_node_id": "b4a0526b-0e31-4f63-892a-6d5197d09acf",
-                "source_handle_id": "6aed9e19-9d26-457b-966d-0a9112f84070",
-                "target_node_id": "c6e3aced-1fc9-48d2-ae55-d2a880e359cb",
-                "target_handle_id": "1e126004-9de7-42c0-b1e1-87f9eb0642e2",
-                "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
@@ -283,22 +168,12 @@ def test_serialize_workflow_with_code():
 
     # AND its raw data should be what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    assert len(workflow_raw_data["edges"]) == 3
-    assert len(workflow_raw_data["nodes"]) == 4
 
     # AND each node should be serialized correctly
-    entrypoint_node = workflow_raw_data["nodes"][0]
-    assert entrypoint_node == {
-        "id": "22555158-d8ba-41b4-a6fc-87c3b25bd073",
-        "type": "ENTRYPOINT",
-        "inputs": [],
-        "data": {"label": "Entrypoint Node", "source_handle_id": "e82390bb-c68c-48c1-9f87-7fbfff494c45"},
-        "display_data": {"position": {"x": 0.0, "y": -50.0}},
-        "base": None,
-        "definition": None,
-    }
 
-    code_execution_node = workflow_raw_data["nodes"][1]
+    code_execution_node = next(
+        n for n in workflow_raw_data["nodes"] if (n.get("base") or {}).get("name") == "CodeExecutionNode"
+    )
     assert code_execution_node == {
         "id": "ae9c5da6-242e-4e0d-abe6-344e2ada3ce3",
         "type": "CODE_EXECUTION",
@@ -351,127 +226,22 @@ def test_serialize_workflow_with_code():
             "merge_behavior": "AWAIT_ANY",
         },
         "ports": [{"id": "d2560e23-cbb6-4b73-b082-294982da72aa", "name": "default", "type": "DEFAULT"}],
-    }
-    assert not DeepDiff(
-        [
+        "outputs": [
             {
-                "id": "52f285fe-1f52-4920-b01b-499762b95220",
-                "type": "TERMINAL",
-                "data": {
-                    "label": "Final Output",
-                    "name": "result",
-                    "target_handle_id": "de8f2cc2-8c32-4782-87d5-4eb5afcd42e3",
-                    "output_id": "283d6849-f3ed-4beb-b261-cf70f90e8d10",
-                    "output_type": "NUMBER",
-                    "node_input_id": "101d5222-8578-413a-be80-b1acf7559a0d",
-                },
-                "inputs": [
-                    {
-                        "id": "101d5222-8578-413a-be80-b1acf7559a0d",
-                        "key": "node_input",
-                        "value": {
-                            "rules": [
-                                {
-                                    "type": "NODE_OUTPUT",
-                                    "data": {
-                                        "node_id": "ae9c5da6-242e-4e0d-abe6-344e2ada3ce3",
-                                        "output_id": "9b2b939b-c840-4e01-99c1-43a3418d7fb7",
-                                    },
-                                }
-                            ],
-                            "combinator": "OR",
-                        },
-                    }
-                ],
-                "display_data": {"position": {"x": 400.0, "y": -175.0}},
-                "base": {
-                    "name": "FinalOutputNode",
-                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
-                },
-                "definition": None,
+                "id": "9b2b939b-c840-4e01-99c1-43a3418d7fb7",
+                "name": "result",
+                "schema": {"type": "integer"},
+                "type": "NUMBER",
+                "value": None,
             },
             {
-                "id": "eccf97c7-e766-471f-9703-4d2595800e66",
-                "type": "TERMINAL",
-                "data": {
-                    "label": "Final Output",
-                    "name": "log",
-                    "target_handle_id": "6b7d7f2c-5cc8-4005-9e66-cdb2c97b1998",
-                    "output_id": "4c136180-050b-4422-a7a4-2a1c6729042c",
-                    "output_type": "STRING",
-                    "node_input_id": "de7ea5c4-da2f-4cfd-9192-0257d0f9ba53",
-                },
-                "inputs": [
-                    {
-                        "id": "de7ea5c4-da2f-4cfd-9192-0257d0f9ba53",
-                        "key": "node_input",
-                        "value": {
-                            "rules": [
-                                {
-                                    "type": "NODE_OUTPUT",
-                                    "data": {
-                                        "node_id": "ae9c5da6-242e-4e0d-abe6-344e2ada3ce3",
-                                        "output_id": "e3536bf9-cf37-4d16-b835-770a9ee6cc53",
-                                    },
-                                }
-                            ],
-                            "combinator": "OR",
-                        },
-                    }
-                ],
-                "display_data": {"position": {"x": 400.0, "y": 75.0}},
-                "base": {
-                    "name": "FinalOutputNode",
-                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
-                },
-                "definition": None,
+                "id": "e3536bf9-cf37-4d16-b835-770a9ee6cc53",
+                "name": "log",
+                "schema": {"type": "string"},
+                "type": "STRING",
+                "value": None,
             },
         ],
-        workflow_raw_data["nodes"][2:],
-        ignore_order=True,
-    )
-
-    # AND each edge should be serialized correctly
-    serialized_edges = workflow_raw_data["edges"]
-    assert not DeepDiff(
-        [
-            {
-                "id": "22dbd9cd-446a-4053-a815-a9a991f69282",
-                "source_node_id": "22555158-d8ba-41b4-a6fc-87c3b25bd073",
-                "source_handle_id": "e82390bb-c68c-48c1-9f87-7fbfff494c45",
-                "target_node_id": "ae9c5da6-242e-4e0d-abe6-344e2ada3ce3",
-                "target_handle_id": "e4a738e8-bbcf-47d2-b7c9-6a034dd412f0",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "2ac757e4-87c3-402c-928f-a3845df10c9f",
-                "source_node_id": "ae9c5da6-242e-4e0d-abe6-344e2ada3ce3",
-                "source_handle_id": "d2560e23-cbb6-4b73-b082-294982da72aa",
-                "target_node_id": "eccf97c7-e766-471f-9703-4d2595800e66",
-                "target_handle_id": "6b7d7f2c-5cc8-4005-9e66-cdb2c97b1998",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "fcc6353a-265c-4a65-9e70-4eb92a04e4e1",
-                "source_node_id": "ae9c5da6-242e-4e0d-abe6-344e2ada3ce3",
-                "source_handle_id": "d2560e23-cbb6-4b73-b082-294982da72aa",
-                "target_node_id": "52f285fe-1f52-4920-b01b-499762b95220",
-                "target_handle_id": "de8f2cc2-8c32-4782-87d5-4eb5afcd42e3",
-                "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
@@ -514,27 +284,12 @@ def test_serialize_workflow__try_wrapped():
 
     # AND its raw data should be what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    assert len(workflow_raw_data["edges"]) == 3
-    assert len(workflow_raw_data["nodes"]) == 4
 
     # AND each node should be serialized correctly
-    entrypoint_node = workflow_raw_data["nodes"][0]
-    assert entrypoint_node == {
-        "id": "1b300de0-cf41-493f-ab41-6fdadf406f6a",
-        "type": "ENTRYPOINT",
-        "inputs": [],
-        "data": {
-            "label": "Entrypoint Node",
-            "source_handle_id": "8cd1e612-39aa-4471-88cf-f7999b713fa6",
-        },
-        "base": None,
-        "definition": None,
-        "display_data": {
-            "position": {"x": 0.0, "y": -50.0},
-        },
-    }
 
-    code_execution_node = workflow_raw_data["nodes"][1]
+    code_execution_node = next(
+        n for n in workflow_raw_data["nodes"] if (n.get("base") or {}).get("name") == "CodeExecutionNode"
+    )
     assert code_execution_node == {
         "id": "1c910367-dff1-4466-85bc-6a8ec4ca039d",
         "type": "CODE_EXECUTION",
@@ -615,128 +370,22 @@ def test_serialize_workflow__try_wrapped():
             "merge_behavior": "AWAIT_ANY",
         },
         "ports": [{"id": "dede45ee-e17a-447c-b1d3-ed0d29ff1057", "name": "default", "type": "DEFAULT"}],
-    }
-
-    final_output_nodes = workflow_raw_data["nodes"][2:]
-    assert not DeepDiff(
-        [
-            {
-                "id": "af4fc1ef-7701-43df-b5e7-4f354f707db2",
-                "type": "TERMINAL",
-                "data": {
-                    "label": "Final Output",
-                    "name": "log",
-                    "target_handle_id": "d243df8d-46f6-4928-ac31-7c775c5d73a9",
-                    "output_id": "5fbd27a0-9831-49c7-93c8-9c2a28c78696",
-                    "output_type": "STRING",
-                    "node_input_id": "a6ffaacf-1284-4e70-907e-042bb281e6fa",
-                },
-                "inputs": [
-                    {
-                        "id": "a6ffaacf-1284-4e70-907e-042bb281e6fa",
-                        "key": "node_input",
-                        "value": {
-                            "rules": [
-                                {
-                                    "type": "NODE_OUTPUT",
-                                    "data": {
-                                        "node_id": "1c910367-dff1-4466-85bc-6a8ec4ca039d",
-                                        "output_id": "0022691d-4707-476f-a9f5-4947b5e36f07",
-                                    },
-                                }
-                            ],
-                            "combinator": "OR",
-                        },
-                    }
-                ],
-                "display_data": {"position": {"x": 400.0, "y": 75.0}},
-                "base": {
-                    "name": "FinalOutputNode",
-                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
-                },
-                "definition": None,
-            },
+        "outputs": [
             {
-                "id": "4cbfa5f7-fc12-4ab2-81cb-168c5caef4f0",
-                "type": "TERMINAL",
-                "data": {
-                    "label": "Final Output",
-                    "name": "result",
-                    "target_handle_id": "9c43709e-25cb-4548-b840-3fcf6a1c9f3e",
-                    "output_id": "400f9ffe-e700-4204-a810-e06123565947",
-                    "output_type": "NUMBER",
-                    "node_input_id": "2cbb7351-a3e8-4a6b-b33e-204b04da7ad6",
-                },
-                "inputs": [
-                    {
-                        "id": "2cbb7351-a3e8-4a6b-b33e-204b04da7ad6",
-                        "key": "node_input",
-                        "value": {
-                            "rules": [
-                                {
-                                    "type": "NODE_OUTPUT",
-                                    "data": {
-                                        "node_id": "1c910367-dff1-4466-85bc-6a8ec4ca039d",
-                                        "output_id": "a483823a-e856-4ad3-ab3b-cac3c2961536",
-                                    },
-                                }
-                            ],
-                            "combinator": "OR",
-                        },
-                    }
-                ],
-                "display_data": {"position": {"x": 400.0, "y": -175.0}},
-                "base": {
-                    "name": "FinalOutputNode",
-                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
-                },
-                "definition": None,
-            },
-        ],
-        sorted(final_output_nodes, key=lambda x: x["id"], reverse=True),
-    )
-
-    # AND each edge should be serialized correctly
-    serialized_edges = workflow_raw_data["edges"]
-    assert not DeepDiff(
-        [
-            {
-                "id": "71c4cd73-1ce9-4261-96a1-dc40b70e62e4",
-                "source_node_id": "1b300de0-cf41-493f-ab41-6fdadf406f6a",
-                "source_handle_id": "8cd1e612-39aa-4471-88cf-f7999b713fa6",
-                "target_node_id": "1c910367-dff1-4466-85bc-6a8ec4ca039d",
-                "target_handle_id": "f4a58613-628e-4a6d-aeae-4f81cf96bbf4",
-                "type": "DEFAULT",
+                "id": "a483823a-e856-4ad3-ab3b-cac3c2961536",
+                "name": "result",
+                "schema": {"type": "integer"},
+                "type": "NUMBER",
+                "value": None,
             },
             {
-                "id": "ac96ad63-d91f-465c-9c52-629877e56492",
-                "source_node_id": "1c910367-dff1-4466-85bc-6a8ec4ca039d",
-                "source_handle_id": "dede45ee-e17a-447c-b1d3-ed0d29ff1057",
-                "target_node_id": "af4fc1ef-7701-43df-b5e7-4f354f707db2",
-                "target_handle_id": "d243df8d-46f6-4928-ac31-7c775c5d73a9",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "636f4540-e0e3-4740-af72-45f78b700cf9",
-                "source_node_id": "1c910367-dff1-4466-85bc-6a8ec4ca039d",
-                "source_handle_id": "dede45ee-e17a-447c-b1d3-ed0d29ff1057",
-                "target_node_id": "4cbfa5f7-fc12-4ab2-81cb-168c5caef4f0",
-                "target_handle_id": "9c43709e-25cb-4548-b840-3fcf6a1c9f3e",
-                "type": "DEFAULT",
+                "id": "0022691d-4707-476f-a9f5-4947b5e36f07",
+                "name": "log",
+                "schema": {"type": "string"},
+                "type": "STRING",
+                "value": None,
             },
         ],
-        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