vellum-ai 0.13.1__py3-none-any.whl → 0.13.2__py3-none-any.whl

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

@@ -1,34 +1,39 @@
-import pytest
+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.references.vellum_secret import VellumSecretReference
+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 ConstantValueGenericNode(BaseNode):
     attr: str = "hello"
 
-    class Outputs(BaseNode.Outputs):
-        output = Inputs.input
 
+def test_serialize_node__constant_value(serialize_node):
+    input_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=ConstantValueGenericNode,
+        global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)},
+    )
 
-@pytest.mark.skip(reason="not implemented")
-def test_serialize_node__try(serialize_node):
-    serialized_node = serialize_node(BasicGenericNode)
     assert not DeepDiff(
         {
-            "id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
-            "label": "BasicGenericNode",
+            "id": "be892bc8-e4de-47ef-ab89-dc9d869af1fe",
+            "label": "ConstantValueGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
             "definition": {
-                "name": "BasicGenericNode",
+                "name": "ConstantValueGenericNode",
                 "module": [
                     "vellum_ee",
                     "workflows",
@@ -39,17 +44,12 @@ def test_serialize_node__try(serialize_node):
                     "test_attributes_serialization",
                 ],
             },
-            "trigger": {"id": "9d3a1b3d-4a38-4f2e-bbf1-dd8be152bce8", "merge_behavior": "AWAIT_ANY"},
-            "ports": [
-                {
-                    "id": "4fbf0fff-a42e-4410-852a-238b5059198e",
-                    "type": "DEFAULT",
-                }
-            ],
+            "trigger": {"id": "279e8228-9b82-43a3-8c31-affc036e3a0b", "merge_behavior": "AWAIT_ANY"},
+            "ports": [{"id": "e6cf13b0-5590-421c-84f2-5a0b169677e1", "type": "DEFAULT", "name": "default"}],
             "adornments": None,
             "attributes": [
                 {
-                    "id": "e3c8b6e4-4b7d-4e3c-8a5c-6e5f4e3c8b6e",
+                    "id": "4cbbfd98-9ab6-41a8-bf4e-ae65f0eafe47",
                     "name": "attr",
                     "value": {
                         "type": "CONSTANT_VALUE",
@@ -60,6 +60,226 @@ def test_serialize_node__try(serialize_node):
                     },
                 }
             ],
+            "outputs": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class WorkflowInputGenericNode(BaseNode):
+    attr: str = Inputs.input
+
+
+def test_serialize_node__workflow_input(serialize_node):
+    input_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=WorkflowInputGenericNode,
+        global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)},
+    )
+    assert not DeepDiff(
+        {
+            "id": "ddfa947f-0830-476b-b07e-ac573968f9a7",
+            "label": "WorkflowInputGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "WorkflowInputGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "b1a5d749-bac0-4f11-8427-191febb2198e", "merge_behavior": "AWAIT_ANY"},
+            "ports": [{"id": "f013bf3f-49f6-41cd-ac13-7439b71a304a", "type": "DEFAULT", "name": "default"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "56d44313-cfdd-4d75-9b19-0beb94e59c4e",
+                    "name": "attr",
+                    "value": {
+                        "type": "WORKFLOW_INPUT",
+                        "input_variable_id": str(input_id),
+                    },
+                }
+            ],
+            "outputs": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class NodeWithOutput(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+
+class NodeWithOutputDisplay(BaseNodeDisplay[NodeWithOutput]):
+    pass
+
+
+class GenericNodeReferencingOutput(BaseNode):
+    attr = NodeWithOutput.Outputs.output
+
+
+def test_serialize_node__node_output(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}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "GenericNodeReferencingOutput",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "449072ba-f7b6-4314-ac96-682123f225e5", "merge_behavior": "AWAIT_ANY"},
+            "ports": [{"id": "ec9a79b8-65c3-4de8-bd29-42c914d72d4f", "type": "DEFAULT", "name": "default"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "73e6a103-1339-41ec-a245-42d43b0637c1",
+                    "name": "attr",
+                    "value": {
+                        "type": "NODE_OUTPUT",
+                        "node_id": "cd954d76-0b0a-4d9b-9bdf-347179c38cb6",
+                        "node_output_id": str(node_output_id),
+                    },
+                }
+            ],
+            "outputs": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class VellumSecretGenericNode(BaseNode):
+    attr = VellumSecretReference(name="hello")
+
+
+def test_serialize_node__vellum_secret(serialize_node):
+    input_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=VellumSecretGenericNode,
+        global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)},
+    )
+    assert not DeepDiff(
+        {
+            "id": "89aa6faa-b533-4179-8912-70a048bf0712",
+            "label": "VellumSecretGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "VellumSecretGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "3ea0305d-d8ea-45fe-8cf1-f6c1c85e6979", "merge_behavior": "AWAIT_ANY"},
+            "ports": [{"id": "69e49322-8f8a-42d2-b36f-3fbe53dad616", "type": "DEFAULT", "name": "default"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "8edd27da-eec1-4539-8bec-629b5ef7a9f9",
+                    "name": "attr",
+                    "value": {"type": "VELLUM_SECRET", "vellum_secret_name": "hello"},
+                }
+            ],
+            "outputs": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class NodeWithExecutions(BaseNode):
+    pass
+
+
+class NodeWithExecutionsDisplay(BaseNodeDisplay[NodeWithExecutions]):
+    pass
+
+
+class GenericNodeReferencingExecutions(BaseNode):
+    attr: int = NodeWithExecutions.Execution.count
+
+
+def test_serialize_node__node_execution(serialize_node):
+    workflow_input_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=GenericNodeReferencingExecutions,
+        global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=workflow_input_id)},
+        global_node_displays={NodeWithExecutions: NodeWithExecutionsDisplay()},
+    )
+
+    assert not DeepDiff(
+        {
+            "id": "6e4d2fb7-891e-492e-97a1-adf44693f518",
+            "label": "GenericNodeReferencingExecutions",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "GenericNodeReferencingExecutions",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "68a91426-4c30-4194-a4c0-cff224d3c0f3", "merge_behavior": "AWAIT_ANY"},
+            "ports": [{"id": "1794c2eb-5cab-49fe-9354-dfc29f11b374", "type": "DEFAULT", "name": "default"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "090e18d7-d5b9-4d5f-9aec-0f562e4b33a8",
+                    "name": "attr",
+                    "value": {
+                        "type": "EXECUTION_COUNTER",
+                        "node_id": "c09bd5a6-dc04-4036-90d4-580acd43c71f",
+                    },
+                }
+            ],
+            "outputs": [],
         },
         serialized_node,
         ignore_order=True,

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

@@ -0,0 +1,177 @@
+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_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 AnnotatedOutputGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output: int
+
+
+def test_serialize_node__annotated_output(serialize_node):
+    serialized_node = serialize_node(AnnotatedOutputGenericNode)
+
+    assert not DeepDiff(
+        {
+            "id": "c0b71cfa-0d9e-4329-bde0-967c44be5c3c",
+            "label": "AnnotatedOutputGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "AnnotatedOutputGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_outputs_serialization",
+                ],
+            },
+            "trigger": {"id": "256ef76c-39a6-4a8f-8bda-922f5972a1d4", "merge_behavior": "AWAIT_ANY"},
+            "ports": [{"id": "378d8939-7a0f-4c87-a849-7fd00067cd2d", "type": "DEFAULT", "name": "default"}],
+            "adornments": None,
+            "attributes": [],
+            "outputs": [
+                {
+                    "id": "8c3c9aff-e1d5-49f4-af75-3ec2fcbb4af2",
+                    "name": "output",
+                    "type": "NUMBER",
+                    "value": None,
+                }
+            ],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class WorkflowInputGenericNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output = Inputs.input
+
+
+def test_serialize_node__workflow_input(serialize_node):
+    input_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=WorkflowInputGenericNode,
+        global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)},
+    )
+
+    assert not DeepDiff(
+        {
+            "id": "ddfa947f-0830-476b-b07e-ac573968f9a7",
+            "label": "WorkflowInputGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "WorkflowInputGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_outputs_serialization",
+                ],
+            },
+            "trigger": {"id": "b1a5d749-bac0-4f11-8427-191febb2198e", "merge_behavior": "AWAIT_ANY"},
+            "ports": [{"id": "f013bf3f-49f6-41cd-ac13-7439b71a304a", "type": "DEFAULT", "name": "default"}],
+            "adornments": None,
+            "attributes": [],
+            "outputs": [
+                {
+                    "id": "2c4a85c0-b017-4cea-a261-e8e8498570c9",
+                    "name": "output",
+                    "type": "STRING",
+                    "value": {
+                        "type": "WORKFLOW_INPUT",
+                        "input_variable_id": str(input_id),
+                    },
+                }
+            ],
+        },
+        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
+
+
+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}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "GenericNodeReferencingOutput",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_outputs_serialization",
+                ],
+            },
+            "trigger": {"id": "449072ba-f7b6-4314-ac96-682123f225e5", "merge_behavior": "AWAIT_ANY"},
+            "ports": [{"id": "ec9a79b8-65c3-4de8-bd29-42c914d72d4f", "type": "DEFAULT", "name": "default"}],
+            "adornments": None,
+            "attributes": [],
+            "outputs": [
+                {
+                    "id": "db010db3-7076-4df9-ae1b-069caa16fa20",
+                    "name": "output",
+                    "type": "STRING",
+                    "value": {
+                        "type": "NODE_OUTPUT",
+                        "node_id": "cd954d76-0b0a-4d9b-9bdf-347179c38cb6",
+                        "node_output_id": str(node_output_id),
+                    },
+                }
+            ],
+        },
+        serialized_node,
+        ignore_order=True,
+    )