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_error_node_serialization.py

@@ -32,6 +32,7 @@ def test_serialize_workflow():
                 "default": None,
                 "required": True,
                 "extensions": {"color": None},
+                "schema": {"type": "integer"},
             }
         ],
         input_variables,
@@ -55,122 +56,30 @@ def test_serialize_workflow():
 
     # AND its raw data should be what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    assert len(workflow_raw_data["edges"]) == 4
-    assert len(workflow_raw_data["nodes"]) == 5
-
-    # AND each node should be serialized correctly
-    entrypoint_node = workflow_raw_data["nodes"][0]
-    assert entrypoint_node == {
-        "id": "10e90662-e998-421d-a5c9-ec16e37a8de1",
-        "type": "ENTRYPOINT",
-        "inputs": [],
-        "data": {
-            "label": "Entrypoint Node",
-            "source_handle_id": "7d86498b-84ed-4feb-8e62-2188058c2c4e",
-        },
-        "display_data": {"position": {"x": 0.0, "y": -50.0}},
-        "base": None,
-        "definition": None,
-    }
 
-    error_node, error_index = next(
-        (
-            (node, index)
-            for index, node in enumerate(workflow_raw_data["nodes"])
-            if node.get("data", {}).get("label") == "Fail Node"
-        ),
-        (None, None),
-    )
-    assert not DeepDiff(
-        {
-            "id": "98597c0e-8951-4320-98ed-efd7f79c00cf",
-            "type": "ERROR",
-            "inputs": [
-                {
-                    "id": "f7fc097f-db5b-48c1-8c33-391834678521",
-                    "key": "error_source_input_id",
-                    "value": {
-                        "rules": [
-                            {
-                                "type": "CONSTANT_VALUE",
-                                "data": {"type": "STRING", "value": "Input threshold was too low"},
-                            }
-                        ],
-                        "combinator": "OR",
-                    },
-                }
-            ],
-            "data": {
-                "label": "Fail Node",
-                "target_handle_id": "865dcf9c-8fe8-44c0-bf36-2c92eb47927e",
-                "error_source_input_id": "f7fc097f-db5b-48c1-8c33-391834678521",
-            },
-            "display_data": {"position": {"x": 400.0, "y": -175.0}},
-            "base": {
-                "name": "ErrorNode",
-                "module": ["vellum", "workflows", "nodes", "core", "error_node", "node"],
-            },
-            "definition": {
-                "name": "FailNode",
-                "module": ["tests", "workflows", "basic_error_node", "workflow"],
-            },
-            "trigger": {
-                "id": "865dcf9c-8fe8-44c0-bf36-2c92eb47927e",
-                "merge_behavior": "AWAIT_ATTRIBUTES",
-            },
-            "ports": [],
-        },
-        error_node,
-        ignore_order=True,
-    )
+    # AND the error node should be serialized correctly
+    error_node = next(n for n in workflow_raw_data["nodes"] if (n.get("base") or {}).get("name") == "ErrorNode")
+
+    assert error_node["id"] == "98597c0e-8951-4320-98ed-efd7f79c00cf"
+    assert error_node["type"] == "ERROR"
+
+    assert error_node["base"]["name"] == "ErrorNode"
+    assert error_node["base"]["module"] == ["vellum", "workflows", "nodes", "core", "error_node", "node"]
+
+    assert error_node["definition"]["name"] == "FailNode"
+    assert error_node["definition"]["module"] == ["tests", "workflows", "basic_error_node", "workflow"]
+
+    assert error_node["trigger"]["id"] == "865dcf9c-8fe8-44c0-bf36-2c92eb47927e"
+    assert error_node["trigger"]["merge_behavior"] == "AWAIT_ATTRIBUTES"
+
+    assert error_node["ports"] == []
+
+    assert error_node["data"]["label"] == "Fail Node"
+    assert error_node["data"]["target_handle_id"] == "865dcf9c-8fe8-44c0-bf36-2c92eb47927e"
+    assert error_node["data"]["error_source_input_id"] == "f7fc097f-db5b-48c1-8c33-391834678521"
+
+    assert len(error_node["inputs"]) == 1
+    assert error_node["inputs"][0]["key"] == "error_source_input_id"
 
     passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "GENERIC"]
     assert len(passthrough_nodes) == 2
-
-    terminal_node = workflow_raw_data["nodes"][-1]
-    assert not DeepDiff(
-        {
-            "id": "e5fff999-80c7-4cbc-9d99-06c653f3ec77",
-            "type": "TERMINAL",
-            "data": {
-                "label": "Final Output",
-                "name": "final_value",
-                "target_handle_id": "b070e9bc-e9b7-46d3-8f5b-0b646bd25cf0",
-                "output_id": "04c5c6be-f5e1-41b8-b668-39e179790d9e",
-                "output_type": "NUMBER",
-                "node_input_id": "c191a5c1-8912-49ad-bf4b-c6d4d58ed482",
-            },
-            "inputs": [
-                {
-                    "id": "c191a5c1-8912-49ad-bf4b-c6d4d58ed482",
-                    "key": "node_input",
-                    "value": {
-                        "rules": [
-                            {
-                                "type": "NODE_OUTPUT",
-                                "data": {
-                                    "node_id": "f490ecde-1679-414b-8db5-2b706be25669",
-                                    "output_id": "29fbd4af-bb7f-40c3-a5c7-2bb47e12c534",
-                                },
-                            }
-                        ],
-                        "combinator": "OR",
-                    },
-                }
-            ],
-            "display_data": {"position": {"x": 600.0, "y": -50.0}},
-            "base": {
-                "name": "FinalOutputNode",
-                "module": [
-                    "vellum",
-                    "workflows",
-                    "nodes",
-                    "displayable",
-                    "final_output_node",
-                    "node",
-                ],
-            },
-            "definition": None,
-        },
-        terminal_node,
-    )

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

@@ -32,6 +32,7 @@ def test_serialize_workflow():
                 "default": None,
                 "required": True,
                 "extensions": {"color": None},
+                "schema": {"type": "string"},
             },
         ],
         input_variables,
@@ -51,100 +52,10 @@ def test_serialize_workflow():
 
     # AND its raw data should be what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    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"] == "f1e4678f-c470-400b-a40e-c8922cc99a86"
-    assert entrypoint_node["type"] == "ENTRYPOINT"
-    assert entrypoint_node["inputs"] == []
-    assert entrypoint_node["data"] == {
-        "label": "Entrypoint Node",
-        "source_handle_id": "40201804-8beb-43ad-8873-a027759512f1",
-    }
-    assert entrypoint_node["base"] is None
-    assert entrypoint_node["definition"] is None
-
-    api_node = workflow_raw_data["nodes"][1]
-    assert api_node["id"] == "bf98371c-65d3-43c1-99a2-0f5369397847"
-
-    final_output_node = workflow_raw_data["nodes"][2]
-    assert not DeepDiff(
-        {
-            "id": "50e3b446-afcd-4a5d-8c6f-5f05eaf2200e",
-            "type": "TERMINAL",
-            "data": {
-                "label": "Final Output",
-                "name": "output",
-                "target_handle_id": "8bd9f4f3-9f66-4d95-8e84-529b0002c531",
-                "output_id": "2b6389d0-266a-4be4-843e-4e543dd3d727",
-                "output_type": "STRING",
-                "node_input_id": "545d6001-cfb5-4ccc-bcdf-3b03ccd67d90",
-            },
-            "inputs": [
-                {
-                    "id": "545d6001-cfb5-4ccc-bcdf-3b03ccd67d90",
-                    "key": "node_input",
-                    "value": {
-                        "rules": [
-                            {
-                                "type": "NODE_OUTPUT",
-                                "data": {
-                                    "node_id": "bf98371c-65d3-43c1-99a2-0f5369397847",
-                                    "output_id": "83a720ef-d397-429b-8124-1db8286b3b00",
-                                },
-                            }
-                        ],
-                        "combinator": "OR",
-                    },
-                }
-            ],
-            "display_data": {"position": {"x": 400.0, "y": -50.0}},
-            "base": {
-                "name": "FinalOutputNode",
-                "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
-            },
-            "definition": None,
-        },
-        final_output_node,
-        ignore_order=True,
-    )
-
-    # AND each edge should be serialized correctly
-    serialized_edges = workflow_raw_data["edges"]
-    assert not DeepDiff(
-        [
-            {
-                "id": "cd4d1a87-1fa2-46df-89fe-7fed9bb4600c",
-                "source_node_id": "f1e4678f-c470-400b-a40e-c8922cc99a86",
-                "source_handle_id": "40201804-8beb-43ad-8873-a027759512f1",
-                "target_node_id": "bf98371c-65d3-43c1-99a2-0f5369397847",
-                "target_handle_id": "43d9db30-98f6-4d75-a487-f74f3c660d8a",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "b741c861-cf67-4649-b9ef-b43a4add72b1",
-                "source_node_id": "bf98371c-65d3-43c1-99a2-0f5369397847",
-                "source_handle_id": "f710ef29-a056-420d-9342-8f0ac79ce4ca",
-                "target_node_id": "50e3b446-afcd-4a5d-8c6f-5f05eaf2200e",
-                "target_handle_id": "8bd9f4f3-9f66-4d95-8e84-529b0002c531",
-                "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 generic node should be serialized correctly
+    generic_node = next(n for n in workflow_raw_data["nodes"] if (n.get("base") or {}).get("name") == "BaseNode")
+    assert generic_node["id"] == "bf98371c-65d3-43c1-99a2-0f5369397847"
 
     # AND the nodes should have been auto-positioned since they all started at (0,0)
     nodes = workflow_raw_data["nodes"]

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

@@ -32,6 +32,7 @@ def test_serialize_workflow():
                 "required": True,
                 "default": None,
                 "extensions": {"color": None},
+                "schema": {"type": "string"},
             },
             {
                 "id": "545ff95e-e86f-4d06-a991-602781e72605",
@@ -40,6 +41,7 @@ def test_serialize_workflow():
                 "required": True,
                 "default": None,
                 "extensions": {"color": None},
+                "schema": {"type": "string"},
             },
         ],
         input_variables,
@@ -53,22 +55,13 @@ def test_serialize_workflow():
 
     # AND its raw data is what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    assert len(workflow_raw_data["nodes"]) == 3
-    assert len(workflow_raw_data["edges"]) == 2
 
     # AND each node should be serialized correctly
-    entrypoint_node = workflow_raw_data["nodes"][0]
-    assert entrypoint_node == {
-        "id": "54c5c7d0-ab86-4ae9-b0b8-ea9ca7b87c14",
-        "type": "ENTRYPOINT",
-        "base": None,
-        "definition": None,
-        "inputs": [],
-        "data": {"label": "Entrypoint Node", "source_handle_id": "41840690-8d85-486e-a864-b0661ccf0f2e"},
-        "display_data": {"position": {"x": 0.0, "y": -50.0}},
-    }
-
-    guardrail_node = workflow_raw_data["nodes"][1]
+    guardrail_node = next(
+        node
+        for node in workflow_raw_data["nodes"]
+        if node["type"] == "METRIC" and node["data"]["label"] == "Example Guardrail Node"
+    )
     assert guardrail_node == {
         "id": "7fef2bbc-cdfc-4f66-80eb-2a52ee52da5f",
         "type": "METRIC",
@@ -124,72 +117,6 @@ def test_serialize_workflow():
         "ports": [{"id": "baa8baa7-8849-4b96-a90d-c0545a60d3a8", "name": "default", "type": "DEFAULT"}],
     }
 
-    final_output_node = workflow_raw_data["nodes"][2]
-    assert final_output_node == {
-        "id": "cbc7197e-67c9-4af5-b781-879c8fd3e4c9",
-        "type": "TERMINAL",
-        "base": {
-            "module": [
-                "vellum",
-                "workflows",
-                "nodes",
-                "displayable",
-                "final_output_node",
-                "node",
-            ],
-            "name": "FinalOutputNode",
-        },
-        "definition": None,
-        "data": {
-            "label": "Final Output",
-            "name": "score",
-            "target_handle_id": "001b97f6-2bc8-4d1e-9572-028dcf17df4e",
-            "output_id": "2abd2b3b-c301-4834-a43f-5db3604f8422",
-            "output_type": "NUMBER",
-            "node_input_id": "6321442a-0d0d-4e25-965d-c24ff24712c5",
-        },
-        "inputs": [
-            {
-                "id": "6321442a-0d0d-4e25-965d-c24ff24712c5",
-                "key": "node_input",
-                "value": {
-                    "rules": [
-                        {
-                            "type": "NODE_OUTPUT",
-                            "data": {
-                                "node_id": "7fef2bbc-cdfc-4f66-80eb-2a52ee52da5f",
-                                "output_id": "eb9c8043-4c04-467b-944e-633c54de6876",
-                            },
-                        }
-                    ],
-                    "combinator": "OR",
-                },
-            }
-        ],
-        "display_data": {"position": {"x": 400.0, "y": -50.0}},
-    }
-
-    # AND each edge should be serialized correctly
-    serialized_edges = workflow_raw_data["edges"]
-    assert serialized_edges == [
-        {
-            "id": "4737867f-1967-45a1-966b-a0bda81a583d",
-            "source_node_id": "54c5c7d0-ab86-4ae9-b0b8-ea9ca7b87c14",
-            "source_handle_id": "41840690-8d85-486e-a864-b0661ccf0f2e",
-            "target_node_id": "7fef2bbc-cdfc-4f66-80eb-2a52ee52da5f",
-            "target_handle_id": "53c299c7-1df2-4d54-bb0d-559a4947c16d",
-            "type": "DEFAULT",
-        },
-        {
-            "id": "5c456a17-a92b-4dad-9569-306043707c9f",
-            "source_node_id": "7fef2bbc-cdfc-4f66-80eb-2a52ee52da5f",
-            "source_handle_id": "baa8baa7-8849-4b96-a90d-c0545a60d3a8",
-            "target_node_id": "cbc7197e-67c9-4af5-b781-879c8fd3e4c9",
-            "target_handle_id": "001b97f6-2bc8-4d1e-9572-028dcf17df4e",
-            "type": "DEFAULT",
-        },
-    ]
-
     # AND the display data is what we expect
     display_data = workflow_raw_data["display_data"]
     assert display_data == {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}}

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

@@ -39,6 +39,7 @@ def test_serialize_workflow():
                 "default": None,
                 "required": True,
                 "extensions": {"color": None},
+                "schema": {"type": "string"},
             }
         ],
         input_variables,
@@ -50,7 +51,7 @@ def test_serialize_workflow():
     assert len(output_variables) == 2
     assert not DeepDiff(
         [
-            {"id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e", "key": "results", "type": "JSON"},
+            {"id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e", "key": "results", "type": "ARRAY"},
             {"id": "0ef1608e-1737-41cc-9b90-a8e124138f70", "key": "json", "type": "JSON"},
         ],
         output_variables,
@@ -59,22 +60,10 @@ def test_serialize_workflow():
 
     # 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": "382842a3-0490-4dee-b87b-eef86766f07c",
-        "type": "ENTRYPOINT",
-        "inputs": [],
-        "data": {"label": "Entrypoint Node", "source_handle_id": "8294baa6-8bf4-4b54-a56b-407b64851b77"},
-        "display_data": {"position": {"x": 0.0, "y": -50.0}},
-        "base": None,
-        "definition": None,
-    }
 
-    prompt_node = workflow_raw_data["nodes"][1]
+    prompt_node = next(n for n in workflow_raw_data["nodes"] if (n.get("base") or {}).get("name") == "InlinePromptNode")
     assert not DeepDiff(
         {
             "id": "f800ecab-fe14-498f-88cf-8f67b3f04338",
@@ -115,7 +104,11 @@ def test_serialize_workflow():
                         "custom_parameters": None,
                     },
                     "input_variables": [
-                        {"id": "15381676-75eb-4688-8ae1-7f9f937d6bb0", "key": "noun", "type": "STRING"}
+                        {
+                            "id": "15381676-75eb-4688-8ae1-7f9f937d6bb0",
+                            "key": "noun",
+                            "type": "STRING",
+                        }
                     ],
                     "prompt_template_block_data": {
                         "version": 1,
@@ -278,91 +271,6 @@ def test_serialize_workflow():
         ignore_order=True,
     )
 
-    final_output_node = workflow_raw_data["nodes"][2]
-    assert not DeepDiff(
-        {
-            "id": "42318326-3ae8-417f-9609-f6d8ae47eafb",
-            "type": "TERMINAL",
-            "data": {
-                "label": "Final Output",
-                "name": "results",
-                "target_handle_id": "46c99277-2b4b-477d-851c-ea497aef6b16",
-                "output_id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e",
-                "output_type": "JSON",
-                "node_input_id": "d7c89dce-765b-494d-a256-aba4bcf87b42",
-            },
-            "inputs": [
-                {
-                    "id": "d7c89dce-765b-494d-a256-aba4bcf87b42",
-                    "key": "node_input",
-                    "value": {
-                        "rules": [
-                            {
-                                "type": "NODE_OUTPUT",
-                                "data": {
-                                    "node_id": "f800ecab-fe14-498f-88cf-8f67b3f04338",
-                                    "output_id": "f5180d8d-89e4-479d-8baf-f6db8f9defa6",
-                                },
-                            }
-                        ],
-                        "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,
-        },
-        final_output_node,
-        ignore_order=True,
-    )
-
-    # AND each edge should be serialized correctly
-    serialized_edges = workflow_raw_data["edges"]
-    assert not DeepDiff(
-        [
-            {
-                "id": "bd2af66e-7724-45e1-8ab6-d37d5b2769f5",
-                "source_node_id": "382842a3-0490-4dee-b87b-eef86766f07c",
-                "source_handle_id": "8294baa6-8bf4-4b54-a56b-407b64851b77",
-                "target_node_id": "f800ecab-fe14-498f-88cf-8f67b3f04338",
-                "target_handle_id": "7040f290-6b61-4519-86f5-d004c38a6905",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "05ca58fb-e02d-48d4-9207-2dad0833a25b",
-                "source_node_id": "f800ecab-fe14-498f-88cf-8f67b3f04338",
-                "source_handle_id": "6fad8947-ecce-498f-8160-46af26b75a81",
-                "target_node_id": "42318326-3ae8-417f-9609-f6d8ae47eafb",
-                "target_handle_id": "46c99277-2b4b-477d-851c-ea497aef6b16",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "0b1a2960-4cd5-4045-844f-42b6c87487aa",
-                "source_node_id": "f800ecab-fe14-498f-88cf-8f67b3f04338",
-                "source_handle_id": "6fad8947-ecce-498f-8160-46af26b75a81",
-                "target_node_id": "1f4e3b7b-6af1-42c8-ab33-05b0f01e2b62",
-                "target_handle_id": "7d94907f-c840-4ced-b813-ee3b17f2a8a9",
-                "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 == {
@@ -620,7 +528,7 @@ def test_inline_prompt_node__coalesce_expression_serialization():
     assert chat_history_entry["value"]["operator"] == "coalesce"
     assert chat_history_entry["value"]["lhs"] == {
         "type": "WORKFLOW_STATE",
-        "state_variable_id": "6012a4f7-a8ff-464d-bd62-7c41fde06fa4",
+        "state_variable_id": "34dc3ea7-e44e-45ac-af42-8d765c4d3c00",
     }
     assert chat_history_entry["value"]["rhs"] == {
         "type": "CONSTANT_VALUE",