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

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

@@ -0,0 +1,289 @@
+from datetime import datetime
+from uuid import uuid4
+
+from deepdiff import DeepDiff
+
+from vellum import DeploymentRead
+
+from tests.workflows.basic_text_prompt_deployment.workflow import BasicTextPromptDeployment
+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(vellum_client):
+    # GIVEN a Workflow with stubbed out API calls
+    deployment = DeploymentRead(
+        id=str(uuid4()),
+        created=datetime.now(),
+        label="Example Prompt Deployment",
+        name="example_prompt_deployment",
+        last_deployed_on=datetime.now(),
+        input_variables=[],
+        active_model_version_ids=[],
+        last_deployed_history_item_id=str(uuid4()),
+    )
+    vellum_client.deployments.retrieve.return_value = deployment
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=BasicTextPromptDeployment
+    )
+    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) == 2
+    assert not DeepDiff(
+        [
+            {
+                "id": "52995b50-84c9-465f-8a4b-a4ee2a92e388",
+                "key": "city",
+                "type": "STRING",
+                "required": None,
+                "default": None,
+                "extensions": None,
+            },
+            {
+                "id": "aa3ca842-250c-4a3f-853f-23928c28d0f8",
+                "key": "date",
+                "type": "STRING",
+                "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 output_variables == [
+        {
+            "id": "a609ab19-db1b-4cd0-bdb0-aee5ed31dc28",
+            "key": "text",
+            "type": "STRING",
+        }
+    ]
+
+    # 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": "d680afbd-de64-4cf6-aa50-912686c48c64",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "definition": {
+            "bases": [],
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "name": "BaseNode",
+        },
+        "data": {
+            "label": "Entrypoint Node",
+            "source_handle_id": "7065a943-1cab-4afd-9690-e678c5b74a2f",
+        },
+        "display_data": {
+            "position": {"x": 0.0, "y": 0.0},
+        },
+    }
+
+    prompt_node = workflow_raw_data["nodes"][1]
+    assert prompt_node == {
+        "id": "56c74024-19a3-4c0d-a5f5-23e1e9f11b21",
+        "type": "PROMPT",
+        "inputs": [
+            {
+                "id": "509cf3d4-db72-483e-897a-0f7f20e70d03",
+                "key": "city",
+                "value": {
+                    "combinator": "OR",
+                    "rules": [
+                        {
+                            "type": "INPUT_VARIABLE",
+                            "data": {
+                                "input_variable_id": "52995b50-84c9-465f-8a4b-a4ee2a92e388",
+                            },
+                        },
+                    ],
+                },
+            },
+            {
+                "id": "7d6ec5ac-c582-4153-bd8a-b8794c367420",
+                "key": "date",
+                "value": {
+                    "combinator": "OR",
+                    "rules": [
+                        {
+                            "type": "INPUT_VARIABLE",
+                            "data": {
+                                "input_variable_id": "aa3ca842-250c-4a3f-853f-23928c28d0f8",
+                            },
+                        },
+                    ],
+                },
+            },
+        ],
+        "data": {
+            "label": "Example Prompt Deployment Node",
+            "output_id": "4d38b850-79e3-4b85-9158-a41d0c535410",
+            "error_output_id": None,
+            "array_output_id": "0cf47d33-6d5f-466f-b826-e814f1d0348b",
+            "source_handle_id": "2f26c7e0-283d-4f04-b639-adebb56bc679",
+            "target_handle_id": "b7605c48-0937-4ecc-914e-0d1058130e65",
+            "variant": "DEPLOYMENT",
+            "prompt_deployment_id": deployment.id,
+            "release_tag": "LATEST",
+        },
+        "display_data": {
+            "position": {
+                "x": 0.0,
+                "y": 0.0,
+            },
+        },
+        "definition": {
+            "bases": [
+                {
+                    "module": [
+                        "vellum",
+                        "workflows",
+                        "nodes",
+                        "displayable",
+                        "prompt_deployment_node",
+                        "node",
+                    ],
+                    "name": "PromptDeploymentNode",
+                }
+            ],
+            "module": [
+                "tests",
+                "workflows",
+                "basic_text_prompt_deployment",
+                "workflow",
+            ],
+            "name": "ExamplePromptDeploymentNode",
+        },
+    }
+
+    final_output_node = workflow_raw_data["nodes"][2]
+    assert final_output_node == {
+        "id": "64ff72c7-8ffc-4e1f-b7a7-e7cd0697f576",
+        "type": "TERMINAL",
+        "definition": {
+            "bases": [
+                {
+                    "bases": [],
+                    "module": [
+                        "vellum",
+                        "workflows",
+                        "nodes",
+                        "bases",
+                        "base",
+                    ],
+                    "name": "BaseNode",
+                },
+            ],
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "displayable",
+                "final_output_node",
+                "node",
+            ],
+            "name": "FinalOutputNode",
+        },
+        "inputs": [
+            {
+                "id": "5f0fdd0f-63fe-4437-bc1b-cd6a84bb84c5",
+                "key": "node_input",
+                "value": {
+                    "combinator": "OR",
+                    "rules": [
+                        {
+                            "data": {
+                                "node_id": "56c74024-19a3-4c0d-a5f5-23e1e9f11b21",
+                                "output_id": "4d38b850-79e3-4b85-9158-a41d0c535410",
+                            },
+                            "type": "NODE_OUTPUT",
+                        }
+                    ],
+                },
+            }
+        ],
+        "data": {
+            "label": "Final Output",
+            "name": "text",
+            "node_input_id": "5f0fdd0f-63fe-4437-bc1b-cd6a84bb84c5",
+            "output_id": "a609ab19-db1b-4cd0-bdb0-aee5ed31dc28",
+            "output_type": "STRING",
+            "target_handle_id": "dced939a-9122-4290-8482-7daa9525dad6",
+        },
+        "display_data": {
+            "position": {
+                "x": 0.0,
+                "y": 0.0,
+            },
+        },
+    }
+
+    # AND each edge should be serialized correctly
+    serialized_edges = workflow_raw_data["edges"]
+    assert serialized_edges == [
+        {
+            "id": "8961d02b-074e-45ab-9f77-4e94606a4344",
+            "source_handle_id": "7065a943-1cab-4afd-9690-e678c5b74a2f",
+            "source_node_id": "d680afbd-de64-4cf6-aa50-912686c48c64",
+            "target_handle_id": "b7605c48-0937-4ecc-914e-0d1058130e65",
+            "target_node_id": "56c74024-19a3-4c0d-a5f5-23e1e9f11b21",
+            "type": "DEFAULT",
+        },
+        {
+            "id": "c2cbf6ef-8582-45c8-a643-fc6ae8fe482f",
+            "source_handle_id": "2f26c7e0-283d-4f04-b639-adebb56bc679",
+            "source_node_id": "56c74024-19a3-4c0d-a5f5-23e1e9f11b21",
+            "target_handle_id": "dced939a-9122-4290-8482-7daa9525dad6",
+            "target_node_id": "64ff72c7-8ffc-4e1f-b7a7-e7cd0697f576",
+            "type": "DEFAULT",
+        },
+    ]
+
+    # 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": "BasicTextPromptDeployment",
+        "module": [
+            "tests",
+            "workflows",
+            "basic_text_prompt_deployment",
+            "workflow",
+        ],
+    }

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

@@ -0,0 +1,354 @@
+from datetime import datetime
+from uuid import uuid4
+
+from deepdiff import DeepDiff
+
+from vellum import WorkflowDeploymentRead
+
+from tests.workflows.basic_subworkflow_deployment.workflow import BasicSubworkflowDeploymentWorkflow
+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(vellum_client):
+    # GIVEN a Workflow with stubbed out API calls
+    deployment = WorkflowDeploymentRead(
+        id=str(uuid4()),
+        created=datetime.now(),
+        label="Example Subworkflow Deployment",
+        name="example_subworkflow_deployment",
+        input_variables=[],
+        output_variables=[],
+        last_deployed_on=datetime.now(),
+        last_deployed_history_item_id=str(uuid4()),
+    )
+    vellum_client.workflow_deployments.retrieve.return_value = deployment
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=BasicSubworkflowDeploymentWorkflow
+    )
+    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) == 2
+    assert not DeepDiff(
+        [
+            {
+                "id": "693cc9a5-8d74-4a58-bdcf-2b4989cdf250",
+                "key": "city",
+                "type": "STRING",
+                "required": None,
+                "default": None,
+                "extensions": None,
+            },
+            {
+                "id": "19a78824-9a98-4ae8-a1fc-61f81a422a17",
+                "key": "date",
+                "type": "STRING",
+                "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) == 2
+    assert not DeepDiff(
+        [
+            {
+                "id": "3f487916-126f-4d6c-95b4-fa72d875b793",
+                "key": "temperature",
+                "type": "NUMBER",
+            },
+            {
+                "id": "45d53a1e-26e8-4c43-a010-80d141acc249",
+                "key": "reasoning",
+                "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"]) == 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": "f0eea82b-39cc-44e3-9c0d-12205ed5652c",
+        "type": "ENTRYPOINT",
+        "definition": {
+            "bases": [],
+            "module": [
+                "vellum",
+                "workflows",
+                "nodes",
+                "bases",
+                "base",
+            ],
+            "name": "BaseNode",
+        },
+        "inputs": [],
+        "data": {
+            "label": "Entrypoint Node",
+            "source_handle_id": "13d9eb34-aecb-496d-9e57-d5e786b0bc7c",
+        },
+        "display_data": {
+            "position": {"x": 0.0, "y": 0.0},
+        },
+    }
+
+    subworkflow_node = workflow_raw_data["nodes"][1]
+    assert subworkflow_node == {
+        "id": "d71f674e-8a6b-44ab-b552-7f4637a4e7a6",
+        "type": "SUBWORKFLOW",
+        "inputs": [
+            {
+                "id": "dade23b9-dab6-4760-9247-da189f1019d2",
+                "key": "city",
+                "value": {
+                    "rules": [
+                        {
+                            "type": "INPUT_VARIABLE",
+                            "data": {"input_variable_id": "693cc9a5-8d74-4a58-bdcf-2b4989cdf250"},
+                        }
+                    ],
+                    "combinator": "OR",
+                },
+            },
+            {
+                "id": "8d73270e-2cf9-4146-b053-4780b99857a6",
+                "key": "date",
+                "value": {
+                    "rules": [
+                        {
+                            "type": "INPUT_VARIABLE",
+                            "data": {"input_variable_id": "19a78824-9a98-4ae8-a1fc-61f81a422a17"},
+                        }
+                    ],
+                    "combinator": "OR",
+                },
+            },
+        ],
+        "data": {
+            "label": "Example Subworkflow Deployment Node",
+            "error_output_id": None,
+            "source_handle_id": "ab0db8a9-7b53-4d88-8667-273b31303273",
+            "target_handle_id": "e4d80502-9281-42c8-91e3-10817bcd7d9e",
+            "variant": "DEPLOYMENT",
+            "workflow_deployment_id": deployment.id,
+            "release_tag": "LATEST",
+        },
+        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "definition": {
+            "bases": [
+                {
+                    "module": [
+                        "vellum",
+                        "workflows",
+                        "nodes",
+                        "displayable",
+                        "subworkflow_deployment_node",
+                        "node",
+                    ],
+                    "name": "SubworkflowDeploymentNode",
+                }
+            ],
+            "module": [
+                "tests",
+                "workflows",
+                "basic_subworkflow_deployment",
+                "workflow",
+            ],
+            "name": "ExampleSubworkflowDeploymentNode",
+        },
+    }
+
+    assert not DeepDiff(
+        [
+            {
+                "id": "18170041-1a70-4836-9fa0-adceba2a1f4f",
+                "type": "TERMINAL",
+                "definition": {
+                    "bases": [
+                        {
+                            "bases": [],
+                            "module": [
+                                "vellum",
+                                "workflows",
+                                "nodes",
+                                "bases",
+                                "base",
+                            ],
+                            "name": "BaseNode",
+                        },
+                    ],
+                    "module": [
+                        "vellum",
+                        "workflows",
+                        "nodes",
+                        "displayable",
+                        "final_output_node",
+                        "node",
+                    ],
+                    "name": "FinalOutputNode",
+                },
+                "data": {
+                    "label": "Final Output",
+                    "name": "temperature",
+                    "target_handle_id": "23117248-df28-4519-bebc-abcb24f966b3",
+                    "output_id": "3f487916-126f-4d6c-95b4-fa72d875b793",
+                    "output_type": "NUMBER",
+                    "node_input_id": "29e10bd7-eb39-4e0f-a406-c3a55b834d6f",
+                },
+                "inputs": [
+                    {
+                        "id": "29e10bd7-eb39-4e0f-a406-c3a55b834d6f",
+                        "key": "node_input",
+                        "value": {
+                            "rules": [
+                                {
+                                    "type": "NODE_OUTPUT",
+                                    "data": {
+                                        "node_id": "d71f674e-8a6b-44ab-b552-7f4637a4e7a6",
+                                        "output_id": "5170e2fb-2b7c-4cb1-8958-6c0f2a359e1e",
+                                    },
+                                }
+                            ],
+                            "combinator": "OR",
+                        },
+                    }
+                ],
+                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            },
+            {
+                "id": "94afd0ac-1ec4-486b-a6fb-fa1ec7029d19",
+                "type": "TERMINAL",
+                "definition": {
+                    "bases": [
+                        {
+                            "bases": [],
+                            "module": [
+                                "vellum",
+                                "workflows",
+                                "nodes",
+                                "bases",
+                                "base",
+                            ],
+                            "name": "BaseNode",
+                        },
+                    ],
+                    "module": [
+                        "vellum",
+                        "workflows",
+                        "nodes",
+                        "displayable",
+                        "final_output_node",
+                        "node",
+                    ],
+                    "name": "FinalOutputNode",
+                },
+                "data": {
+                    "label": "Final Output",
+                    "name": "reasoning",
+                    "target_handle_id": "c3aeba92-4faf-4814-9842-eec7436ee555",
+                    "output_id": "45d53a1e-26e8-4c43-a010-80d141acc249",
+                    "output_type": "STRING",
+                    "node_input_id": "daaf0664-0aca-4619-af30-07caa1486b90",
+                },
+                "inputs": [
+                    {
+                        "id": "daaf0664-0aca-4619-af30-07caa1486b90",
+                        "key": "node_input",
+                        "value": {
+                            "rules": [
+                                {
+                                    "type": "NODE_OUTPUT",
+                                    "data": {
+                                        "node_id": "d71f674e-8a6b-44ab-b552-7f4637a4e7a6",
+                                        "output_id": "d9ab77e4-226f-436f-ad70-585b57510001",
+                                    },
+                                }
+                            ],
+                            "combinator": "OR",
+                        },
+                    }
+                ],
+                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            },
+        ],
+        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": "77f5b6e0-16e3-4bb5-9344-4d9557289802",
+                "source_handle_id": "13d9eb34-aecb-496d-9e57-d5e786b0bc7c",
+                "source_node_id": "f0eea82b-39cc-44e3-9c0d-12205ed5652c",
+                "target_handle_id": "e4d80502-9281-42c8-91e3-10817bcd7d9e",
+                "target_node_id": "d71f674e-8a6b-44ab-b552-7f4637a4e7a6",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "69933897-e91e-4c6c-9ba3-ed3e3c265c73",
+                "source_handle_id": "ab0db8a9-7b53-4d88-8667-273b31303273",
+                "source_node_id": "d71f674e-8a6b-44ab-b552-7f4637a4e7a6",
+                "target_handle_id": "c3aeba92-4faf-4814-9842-eec7436ee555",
+                "target_node_id": "94afd0ac-1ec4-486b-a6fb-fa1ec7029d19",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "86a9af31-f78e-45ac-b170-f66bbba98f9d",
+                "source_handle_id": "ab0db8a9-7b53-4d88-8667-273b31303273",
+                "source_node_id": "d71f674e-8a6b-44ab-b552-7f4637a4e7a6",
+                "target_handle_id": "23117248-df28-4519-bebc-abcb24f966b3",
+                "target_node_id": "18170041-1a70-4836-9fa0-adceba2a1f4f",
+                "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": "BasicSubworkflowDeploymentWorkflow",
+        "module": [
+            "tests",
+            "workflows",
+            "basic_subworkflow_deployment",
+            "workflow",
+        ],
+    }