vellum-ai 0.12.17__py3-none-any.whl → 0.13.0__py3-none-any.whl

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

@@ -0,0 +1,355 @@
+from uuid import uuid4
+
+from deepdiff import DeepDiff
+
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.ports.port import Port
+from vellum_ee.workflows.display.base import WorkflowInputsDisplay
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
+from vellum_ee.workflows.display.nodes.vellum.base_node import BaseNodeDisplay
+
+
+class Inputs(BaseInputs):
+    input: str
+
+
+class BasicGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+
+def test_serialize_node__basic(serialize_node):
+    serialized_node = serialize_node(BasicGenericNode)
+    assert not DeepDiff(
+        {
+            "id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
+            "label": "BasicGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "BasicGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_ports_serialization",
+                ],
+            },
+            "trigger": {"id": "9d3a1b3d-4a38-4f2e-bbf1-dd8be152bce8", "merge_behavior": "AWAIT_ANY"},
+            "ports": [
+                {
+                    "id": "4fbf0fff-a42e-4410-852a-238b5059198e",
+                    "type": "DEFAULT",
+                }
+            ],
+            "adornments": None,
+            "attributes": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class IfGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+    class Ports(BaseNode.Ports):
+        if_branch = Port.on_if(Inputs.input.equals("hello"))
+
+
+def test_serialize_node__if(serialize_node):
+    input_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=IfGenericNode, global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)}
+    )
+
+    assert not DeepDiff(
+        {
+            "id": "31da54ae-1abb-4e9e-8a7d-6f4f30a78c72",
+            "label": "IfGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "IfGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_ports_serialization",
+                ],
+            },
+            "trigger": {"id": "a8afaebc-7333-4e3f-b221-24452b4a1d47", "merge_behavior": "AWAIT_ANY"},
+            "ports": [
+                {
+                    "id": "7605b4c0-a432-4517-b759-5858045a5146",
+                    "type": "IF",
+                    "expression": {
+                        "type": "BINARY_EXPRESSION",
+                        "lhs": {
+                            "type": "WORKFLOW_INPUT",
+                            "input_variable_id": str(input_id),
+                        },
+                        "operator": "=",
+                        "rhs": {
+                            "type": "CONSTANT_VALUE",
+                            "value": {
+                                "type": "STRING",
+                                "value": "hello",
+                            },
+                        },
+                    },
+                }
+            ],
+            "adornments": None,
+            "attributes": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class IfElseGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+    class Ports(BaseNode.Ports):
+        if_branch = Port.on_if(Inputs.input.equals("hello"))
+        else_branch = Port.on_else()
+
+
+def test_serialize_node__if_else(serialize_node):
+    input_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=IfElseGenericNode, global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)}
+    )
+
+    assert not DeepDiff(
+        {
+            "id": "1f499f82-8cc0-4060-bf4d-d20ac409d4aa",
+            "label": "IfElseGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "IfElseGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_ports_serialization",
+                ],
+            },
+            "trigger": {"id": "5b4f6553-69ca-4844-bbe4-9e5594bc8cae", "merge_behavior": "AWAIT_ANY"},
+            "ports": [
+                {
+                    "id": "3eeb7f03-7d65-45aa-b0e5-c7a453f5cbdf",
+                    "type": "IF",
+                    "expression": {
+                        "type": "BINARY_EXPRESSION",
+                        "lhs": {
+                            "type": "WORKFLOW_INPUT",
+                            "input_variable_id": str(input_id),
+                        },
+                        "operator": "=",
+                        "rhs": {
+                            "type": "CONSTANT_VALUE",
+                            "value": {
+                                "type": "STRING",
+                                "value": "hello",
+                            },
+                        },
+                    },
+                },
+                {
+                    "id": "b8472c77-74d5-4432-bf8b-6cd65d3dde06",
+                    "type": "ELSE",
+                    "expression": None,
+                },
+            ],
+            "adornments": None,
+            "attributes": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class IfElifElseGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+    class Ports(BaseNode.Ports):
+        if_branch = Port.on_if(Inputs.input.equals("hello"))
+        elif_branch = Port.on_elif(Inputs.input.equals("world"))
+        else_branch = Port.on_else()
+
+
+def test_serialize_node__if_elif_else(serialize_node):
+    input_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=IfElifElseGenericNode,
+        global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)},
+    )
+    assert not DeepDiff(
+        {
+            "id": "21c49bfb-a90c-4565-a4e6-8eb5187e81ca",
+            "label": "IfElifElseGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "IfElifElseGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_ports_serialization",
+                ],
+            },
+            "trigger": {"id": "22d55b5b-3545-4498-8658-9d0464202e78", "merge_behavior": "AWAIT_ANY"},
+            "ports": [
+                {
+                    "id": "f6e0a2c0-192d-452f-bde4-32fb938e91bc",
+                    "type": "IF",
+                    "expression": {
+                        "type": "BINARY_EXPRESSION",
+                        "lhs": {
+                            "type": "WORKFLOW_INPUT",
+                            "input_variable_id": str(input_id),
+                        },
+                        "operator": "=",
+                        "rhs": {
+                            "type": "CONSTANT_VALUE",
+                            "value": {
+                                "type": "STRING",
+                                "value": "hello",
+                            },
+                        },
+                    },
+                },
+                {
+                    "id": "7e44de04-e816-4da8-9251-cf389442a5d6",
+                    "type": "ELIF",
+                    "expression": {
+                        "type": "BINARY_EXPRESSION",
+                        "lhs": {
+                            "type": "WORKFLOW_INPUT",
+                            "input_variable_id": str(input_id),
+                        },
+                        "operator": "=",
+                        "rhs": {
+                            "type": "CONSTANT_VALUE",
+                            "value": {
+                                "type": "STRING",
+                                "value": "world",
+                            },
+                        },
+                    },
+                },
+                {
+                    "id": "00db3698-ddf5-413b-8408-fff664c212d7",
+                    "type": "ELSE",
+                    "expression": None,
+                },
+            ],
+            "adornments": None,
+            "attributes": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class NodeWithOutput(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+
+class NodeWithOutputDisplay(BaseNodeDisplay[NodeWithOutput]):
+    pass
+
+
+class GenericNodeReferencingOutput(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = NodeWithOutput.Outputs.output
+
+    class Ports(BaseNode.Ports):
+        if_branch = Port.on_if(NodeWithOutput.Outputs.output.equals("hello"))
+
+
+def test_serialize_node__node_output_reference(serialize_node):
+    workflow_input_id = uuid4()
+    node_output_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=GenericNodeReferencingOutput,
+        global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=workflow_input_id)},
+        global_node_displays={NodeWithOutput: NodeWithOutputDisplay()},
+        global_node_output_displays={
+            NodeWithOutput.Outputs.output: (NodeWithOutput, NodeOutputDisplay(id=node_output_id, name="output"))
+        },
+    )
+
+    assert not DeepDiff(
+        {
+            "id": "c1e2ce60-ac3a-4b17-915e-abe861734e03",
+            "label": "GenericNodeReferencingOutput",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "definition": {
+                "name": "GenericNodeReferencingOutput",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_ports_serialization",
+                ],
+            },
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "trigger": {"id": "449072ba-f7b6-4314-ac96-682123f225e5", "merge_behavior": "AWAIT_ANY"},
+            "ports": [
+                {
+                    "id": "ec9a79b8-65c3-4de8-bd29-42c914d72d4f",
+                    "type": "IF",
+                    "expression": {
+                        "type": "BINARY_EXPRESSION",
+                        "lhs": {
+                            "type": "NODE_OUTPUT",
+                            "node_id": "cd954d76-0b0a-4d9b-9bdf-347179c38cb6",
+                            "node_output_id": str(node_output_id),
+                        },
+                        "operator": "=",
+                        "rhs": {
+                            "type": "CONSTANT_VALUE",
+                            "value": {
+                                "type": "STRING",
+                                "value": "hello",
+                            },
+                        },
+                    },
+                }
+            ],
+            "adornments": None,
+            "attributes": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )

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

@@ -14,22 +14,6 @@ class BasicGenericNode(BaseNode):
         output = Inputs.input
 
 
-class AwaitAnyGenericNode(BaseNode):
-    class Outputs(BaseNode.Outputs):
-        output = Inputs.input
-
-    class Trigger(BaseNode.Trigger):
-        merge_behavior = MergeBehavior.AWAIT_ANY
-
-
-class AwaitAllGenericNode(BaseNode):
-    class Outputs(BaseNode.Outputs):
-        output = Inputs.input
-
-    class Trigger(BaseNode.Trigger):
-        merge_behavior = MergeBehavior.AWAIT_ALL
-
-
 def test_serialize_node__basic(serialize_node):
     serialized_node = serialize_node(BasicGenericNode)
     assert not DeepDiff(
@@ -38,6 +22,7 @@ def test_serialize_node__basic(serialize_node):
             "label": "BasicGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
             "definition": {
                 "name": "BasicGenericNode",
                 "module": [
@@ -49,10 +34,14 @@ def test_serialize_node__basic(serialize_node):
                     "generic_nodes",
                     "test_trigger_serialization",
                 ],
-                "bases": [{"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}],
             },
             "trigger": {"id": "9d3a1b3d-4a38-4f2e-bbf1-dd8be152bce8", "merge_behavior": "AWAIT_ANY"},
-            "ports": [],
+            "ports": [
+                {
+                    "id": "4fbf0fff-a42e-4410-852a-238b5059198e",
+                    "type": "DEFAULT",
+                }
+            ],
             "adornments": None,
             "attributes": [],
         },
@@ -61,6 +50,14 @@ def test_serialize_node__basic(serialize_node):
     )
 
 
+class AwaitAnyGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+    class Trigger(BaseNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ANY
+
+
 def test_serialize_node__await_any(serialize_node):
     serialized_node = serialize_node(AwaitAnyGenericNode)
     assert not DeepDiff(
@@ -69,6 +66,7 @@ def test_serialize_node__await_any(serialize_node):
             "label": "AwaitAnyGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
             "definition": {
                 "name": "AwaitAnyGenericNode",
                 "module": [
@@ -80,10 +78,14 @@ def test_serialize_node__await_any(serialize_node):
                     "generic_nodes",
                     "test_trigger_serialization",
                 ],
-                "bases": [{"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}],
             },
             "trigger": {"id": "ffa72187-9a18-453f-ae55-b77aad332630", "merge_behavior": "AWAIT_ANY"},
-            "ports": [],
+            "ports": [
+                {
+                    "id": "3e219c0a-e5f8-443a-ac78-1a458b189009",
+                    "type": "DEFAULT",
+                }
+            ],
             "adornments": None,
             "attributes": [],
         },
@@ -92,6 +94,14 @@ def test_serialize_node__await_any(serialize_node):
     )
 
 
+class AwaitAllGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+    class Trigger(BaseNode.Trigger):
+        merge_behavior = MergeBehavior.AWAIT_ALL
+
+
 def test_serialize_node__await_all(serialize_node):
     serialized_node = serialize_node(AwaitAllGenericNode)
     assert not DeepDiff(
@@ -100,6 +110,7 @@ def test_serialize_node__await_all(serialize_node):
             "label": "AwaitAllGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
             "definition": {
                 "name": "AwaitAllGenericNode",
                 "module": [
@@ -111,10 +122,14 @@ def test_serialize_node__await_all(serialize_node):
                     "generic_nodes",
                     "test_trigger_serialization",
                 ],
-                "bases": [{"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}],
             },
             "trigger": {"id": "62074276-c817-476d-b59d-da523ae3f218", "merge_behavior": "AWAIT_ALL"},
-            "ports": [],
+            "ports": [
+                {
+                    "id": "9797e93f-9fe6-48b2-a48e-476abbd20e32",
+                    "type": "DEFAULT",
+                }
+            ],
             "adornments": None,
             "attributes": [],
         },

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

@@ -66,17 +66,8 @@ def test_serialize_workflow(vellum_client):
         "inputs": [],
         "data": {"label": "Entrypoint Node", "source_handle_id": "8eaa7f02-25ff-4a00-9b0a-5185718d89b3"},
         "display_data": {"position": {"x": 0.0, "y": 0.0}},
-        "definition": {
-            "name": "BaseNode",
-            "module": [
-                "vellum",
-                "workflows",
-                "nodes",
-                "bases",
-                "base",
-            ],
-            "bases": [],
-        },
+        "base": None,
+        "definition": None,
     }
 
     api_node = workflow_raw_data["nodes"][1]
@@ -195,12 +186,13 @@ def test_serialize_workflow(vellum_client):
                 "status_code_output_id": "fecc16c3-400e-4fd3-8223-08366070e3b1",
             },
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {
+                "name": "APINode",
+                "module": ["vellum", "workflows", "nodes", "displayable", "api_node", "node"],
+            },
             "definition": {
                 "name": "SimpleAPINode",
                 "module": ["tests", "workflows", "basic_vellum_api_node", "workflow"],
-                "bases": [
-                    {"name": "APINode", "module": ["vellum", "workflows", "nodes", "displayable", "api_node", "node"]}
-                ],
             },
         },
         api_node,
@@ -239,7 +231,7 @@ def test_serialize_workflow(vellum_client):
                     }
                 ],
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
-                "definition": {
+                "base": {
                     "name": "FinalOutputNode",
                     "module": [
                         "vellum",
@@ -249,20 +241,8 @@ def test_serialize_workflow(vellum_client):
                         "final_output_node",
                         "node",
                     ],
-                    "bases": [
-                        {
-                            "name": "BaseNode",
-                            "module": [
-                                "vellum",
-                                "workflows",
-                                "nodes",
-                                "bases",
-                                "base",
-                            ],
-                            "bases": [],
-                        }
-                    ],
                 },
+                "definition": None,
             },
             {
                 "id": "736f9bd0-f487-42af-bdb3-780b4941c61c",
@@ -294,7 +274,7 @@ def test_serialize_workflow(vellum_client):
                     }
                 ],
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
-                "definition": {
+                "base": {
                     "name": "FinalOutputNode",
                     "module": [
                         "vellum",
@@ -304,20 +284,8 @@ def test_serialize_workflow(vellum_client):
                         "final_output_node",
                         "node",
                     ],
-                    "bases": [
-                        {
-                            "name": "BaseNode",
-                            "module": [
-                                "vellum",
-                                "workflows",
-                                "nodes",
-                                "bases",
-                                "base",
-                            ],
-                            "bases": [],
-                        }
-                    ],
                 },
+                "definition": None,
             },
             {
                 "id": "3f3ffc50-b156-48ac-b5f3-f68cb05c2b90",
@@ -349,7 +317,7 @@ def test_serialize_workflow(vellum_client):
                     }
                 ],
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
-                "definition": {
+                "base": {
                     "name": "FinalOutputNode",
                     "module": [
                         "vellum",
@@ -359,20 +327,8 @@ def test_serialize_workflow(vellum_client):
                         "final_output_node",
                         "node",
                     ],
-                    "bases": [
-                        {
-                            "name": "BaseNode",
-                            "module": [
-                                "vellum",
-                                "workflows",
-                                "nodes",
-                                "bases",
-                                "base",
-                            ],
-                            "bases": [],
-                        }
-                    ],
                 },
+                "definition": None,
             },
         ],
         final_output_nodes,