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

@@ -1,5 +1,11 @@
 from deepdiff import DeepDiff
 
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes import MapNode
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
+from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.state.base import BaseState
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 from tests.workflows.basic_map_node.workflow import SimpleMapExample
@@ -31,6 +37,7 @@ def test_serialize_workflow():
                 "required": True,
                 "default": None,
                 "extensions": {"color": None},
+                "schema": {"type": "array", "items": {"type": "string"}},
             },
         ],
         input_variables,
@@ -54,333 +61,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": "c0aa464d-1685-4f15-a051-31b426fec92e",
-        "type": "ENTRYPOINT",
-        "inputs": [],
-        "data": {
-            "label": "Entrypoint Node",
-            "source_handle_id": "844d992e-60ab-4af2-a8ff-52cd858386f7",
-        },
-        "base": None,
-        "definition": None,
-        "display_data": {
-            "position": {"x": 0.0, "y": -50.0},
-        },
-    }
 
-    map_node = workflow_raw_data["nodes"][1]
-    assert not DeepDiff(
-        {
-            "id": "f2f94af1-fcbe-497c-80ce-80952c8903c8",
-            "type": "MAP",
-            "inputs": [
-                {
-                    "id": "0acef1e7-caa1-4d9f-bb69-6981729af18d",
-                    "key": "items",
-                    "value": {
-                        "rules": [
-                            {
-                                "type": "INPUT_VARIABLE",
-                                "data": {"input_variable_id": "db2eb237-38e4-417a-8bfc-5bda0f3165ca"},
-                            }
-                        ],
-                        "combinator": "OR",
-                    },
-                }
-            ],
-            "data": {
-                "label": "Map Fruits Node",
-                "error_output_id": None,
-                "source_handle_id": "aff8a80e-7ce7-43d2-9c9e-9d137efd3b33",
-                "target_handle_id": "e16e9d55-5f26-4d89-8c7a-939f1f463d80",
-                "variant": "INLINE",
-                "workflow_raw_data": {
-                    "nodes": [
-                        {
-                            "id": "ff9bfe6e-839d-4d40-b8fc-313b3bbd0ab0",
-                            "type": "ENTRYPOINT",
-                            "inputs": [],
-                            "data": {
-                                "label": "Entrypoint Node",
-                                "source_handle_id": "520d3616-8369-4e79-9da5-3febae299c2a",
-                            },
-                            "display_data": {"position": {"x": 0.0, "y": -50.0}},
-                            "base": None,
-                            "definition": None,
-                        },
-                        {
-                            "id": "9dffa9e5-00e3-4d20-8ab8-ad0412f35bf2",
-                            "label": "Iteration",
-                            "type": "GENERIC",
-                            "should_file_merge": True,
-                            "display_data": {"position": {"x": 200.0, "y": -50.0}},
-                            "base": {
-                                "name": "BaseNode",
-                                "module": ["vellum", "workflows", "nodes", "bases", "base"],
-                            },
-                            "definition": {
-                                "name": "Iteration",
-                                "module": ["tests", "workflows", "basic_map_node", "workflow"],
-                            },
-                            "trigger": {
-                                "id": "40631b74-8a35-4cb1-897d-3c93f8d5b625",
-                                "merge_behavior": "AWAIT_ATTRIBUTES",
-                            },
-                            "ports": [
-                                {"id": "8d5d36aa-f7c2-49bf-b4d5-cda6f6cd18f9", "type": "DEFAULT", "name": "default"}
-                            ],
-                            "adornments": None,
-                            "attributes": [
-                                {
-                                    "id": "52a43441-9162-40b1-96c6-445250af1c34",
-                                    "name": "item",
-                                    "value": {
-                                        "type": "WORKFLOW_INPUT",
-                                        "input_variable_id": "b29bb546-9bc8-4136-857d-8c7a464ba9d4",
-                                    },
-                                },
-                                {
-                                    "id": "dadea752-95db-4e5f-a0d4-f3920376ee8f",
-                                    "name": "index",
-                                    "value": {
-                                        "type": "WORKFLOW_INPUT",
-                                        "input_variable_id": "17e7ca49-668f-450d-a792-e1f97d13db67",
-                                    },
-                                },
-                            ],
-                            "outputs": [
-                                {
-                                    "id": "3d97725d-a423-4072-9a73-7b6517edc76e",
-                                    "name": "count",
-                                    "type": "NUMBER",
-                                    "value": None,
-                                }
-                            ],
-                        },
-                        {
-                            "id": "6f4883b2-70b1-4e1c-ae15-7d0f5aec810b",
-                            "type": "TERMINAL",
-                            "data": {
-                                "label": "Final Output",
-                                "name": "count",
-                                "target_handle_id": "9d74571f-b7f5-4c1d-8b7c-b9c648738a4d",
-                                "output_id": "2a957315-fae0-4366-8a35-f0b315c5eade",
-                                "output_type": "NUMBER",
-                                "node_input_id": "05f2e4cc-0e36-4c6b-a535-be38281e5280",
-                            },
-                            "inputs": [
-                                {
-                                    "id": "05f2e4cc-0e36-4c6b-a535-be38281e5280",
-                                    "key": "node_input",
-                                    "value": {
-                                        "rules": [
-                                            {
-                                                "type": "NODE_OUTPUT",
-                                                "data": {
-                                                    "node_id": "9dffa9e5-00e3-4d20-8ab8-ad0412f35bf2",
-                                                    "output_id": "3d97725d-a423-4072-9a73-7b6517edc76e",
-                                                },
-                                            }
-                                        ],
-                                        "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,
-                        },
-                    ],
-                    "edges": [
-                        {
-                            "id": "8b52ffcb-1ccf-41eb-ad89-705c6abc2369",
-                            "source_node_id": "ff9bfe6e-839d-4d40-b8fc-313b3bbd0ab0",
-                            "source_handle_id": "520d3616-8369-4e79-9da5-3febae299c2a",
-                            "target_node_id": "9dffa9e5-00e3-4d20-8ab8-ad0412f35bf2",
-                            "target_handle_id": "40631b74-8a35-4cb1-897d-3c93f8d5b625",
-                            "type": "DEFAULT",
-                        },
-                        {
-                            "id": "14152688-6996-4d64-9231-a6e66a4827eb",
-                            "source_node_id": "9dffa9e5-00e3-4d20-8ab8-ad0412f35bf2",
-                            "source_handle_id": "8d5d36aa-f7c2-49bf-b4d5-cda6f6cd18f9",
-                            "target_node_id": "6f4883b2-70b1-4e1c-ae15-7d0f5aec810b",
-                            "target_handle_id": "9d74571f-b7f5-4c1d-8b7c-b9c648738a4d",
-                            "type": "DEFAULT",
-                        },
-                    ],
-                    "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
-                    "definition": {
-                        "name": "IterationSubworkflow",
-                        "module": [
-                            "tests",
-                            "workflows",
-                            "basic_map_node",
-                            "workflow",
-                        ],
-                    },
-                    "output_values": [
-                        {
-                            "output_variable_id": "2a957315-fae0-4366-8a35-f0b315c5eade",
-                            "value": {
-                                "type": "NODE_OUTPUT",
-                                "node_id": "9dffa9e5-00e3-4d20-8ab8-ad0412f35bf2",
-                                "node_output_id": "3d97725d-a423-4072-9a73-7b6517edc76e",
-                            },
-                        }
-                    ],
-                },
-                "input_variables": [
-                    {
-                        "id": "b29bb546-9bc8-4136-857d-8c7a464ba9d4",
-                        "key": "item",
-                        "type": "JSON",
-                        "required": True,
-                        "default": None,
-                        "extensions": {"color": None},
-                    },
-                    {
-                        "id": "17e7ca49-668f-450d-a792-e1f97d13db67",
-                        "key": "index",
-                        "type": "NUMBER",
-                        "required": True,
-                        "default": None,
-                        "extensions": {"color": None},
-                    },
-                    {
-                        "id": "0acef1e7-caa1-4d9f-bb69-6981729af18d",
-                        "key": "items",
-                        "type": "JSON",
-                        "required": True,
-                        "default": None,
-                        "extensions": {"color": None},
-                    },
-                ],
-                "output_variables": [{"id": "2a957315-fae0-4366-8a35-f0b315c5eade", "key": "count", "type": "NUMBER"}],
-                "concurrency": None,
-                "items_input_id": "0acef1e7-caa1-4d9f-bb69-6981729af18d",
-                "item_input_id": "b29bb546-9bc8-4136-857d-8c7a464ba9d4",
-                "index_input_id": "17e7ca49-668f-450d-a792-e1f97d13db67",
-            },
-            "display_data": {"position": {"x": 200.0, "y": -50.0}},
-            "base": {
-                "name": "MapNode",
-                "module": ["vellum", "workflows", "nodes", "core", "map_node", "node"],
-            },
-            "definition": {
-                "name": "MapFruitsNode",
-                "module": ["tests", "workflows", "basic_map_node", "workflow"],
-            },
-            "trigger": {
-                "id": "e16e9d55-5f26-4d89-8c7a-939f1f463d80",
-                "merge_behavior": "AWAIT_ATTRIBUTES",
-            },
-            "ports": [{"id": "aff8a80e-7ce7-43d2-9c9e-9d137efd3b33", "name": "default", "type": "DEFAULT"}],
-        },
-        map_node,
-        ignore_order=True,
-    )
-
-    assert not DeepDiff(
-        {
-            "id": "bacc5d55-07d4-4a0a-a69e-831524480de5",
-            "type": "TERMINAL",
-            "data": {
-                "label": "Final Output",
-                "name": "final_value",
-                "target_handle_id": "720dd872-2f3d-47b9-8245-89387f04f300",
-                "output_id": "145b0b68-224b-4f83-90e6-eea3457e6c3e",
-                "output_type": "JSON",
-                "node_input_id": "57fe72e9-cc9f-48f2-b5ea-25ac52637a33",
-            },
-            "inputs": [
-                {
-                    "id": "57fe72e9-cc9f-48f2-b5ea-25ac52637a33",
-                    "key": "node_input",
-                    "value": {
-                        "rules": [
-                            {
-                                "type": "NODE_OUTPUT",
-                                "data": {
-                                    "node_id": "f2f94af1-fcbe-497c-80ce-80952c8903c8",
-                                    "output_id": "2a957315-fae0-4366-8a35-f0b315c5eade",
-                                },
-                            }
-                        ],
-                        "combinator": "OR",
-                    },
-                }
-            ],
-            "display_data": {"position": {"x": 400.0, "y": -50.0}},
-            "base": {
-                "module": [
-                    "vellum",
-                    "workflows",
-                    "nodes",
-                    "displayable",
-                    "final_output_node",
-                    "node",
-                ],
-                "name": "FinalOutputNode",
-            },
-            "definition": None,
-        },
-        workflow_raw_data["nodes"][2],
-        # TODO: Fix output ID not referencing map node workflow output
-        # https://app.shortcut.com/vellum/story/5667/fix-output-display-reference-on-map-nodes
-        exclude_regex_paths=r"root\['inputs'\]\[0\]\['value'\]\['rules'\]\[0\]\['data'\]\['output_id'\]",
-    )
-
-    # AND each edge should be serialized correctly
-    serialized_edges = workflow_raw_data["edges"]
-    assert not DeepDiff(
-        [
-            {
-                "id": "528eb20a-9db4-4c01-87c0-39b9f5f09753",
-                "source_node_id": "c0aa464d-1685-4f15-a051-31b426fec92e",
-                "source_handle_id": "844d992e-60ab-4af2-a8ff-52cd858386f7",
-                "target_node_id": "f2f94af1-fcbe-497c-80ce-80952c8903c8",
-                "target_handle_id": "e16e9d55-5f26-4d89-8c7a-939f1f463d80",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "47a34f6e-d139-4702-aa46-6212bb8a150f",
-                "source_node_id": "f2f94af1-fcbe-497c-80ce-80952c8903c8",
-                "source_handle_id": "aff8a80e-7ce7-43d2-9c9e-9d137efd3b33",
-                "target_node_id": "bacc5d55-07d4-4a0a-a69e-831524480de5",
-                "target_handle_id": "720dd872-2f3d-47b9-8245-89387f04f300",
-                "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,
-        }
-    }
+    map_node = next(n for n in workflow_raw_data["nodes"] if (n.get("base") or {}).get("name") == "MapNode")
 
     # AND the definition should be what we expect
     definition = workflow_raw_data["definition"]
@@ -397,3 +81,57 @@ def test_serialize_workflow():
     # AND the map node's items input ID should match the subworkflow's items input ID
     items_input_id = map_node["data"]["items_input_id"]
     assert map_node["inputs"][0]["id"] == items_input_id
+
+
+def test_serialize_workflow__dynamic_field_reference_items():
+    """Tests that a map node with items set to a dynamic field reference serializes correctly."""
+
+    # GIVEN a prompt node with json output
+    class MyPromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = []
+
+    # AND a simple subworkflow for the map node
+    class IterationNode(BaseNode):
+        item = MapNode.SubworkflowInputs.item
+
+        class Outputs(BaseOutputs):
+            value: str
+
+        def run(self) -> Outputs:
+            return self.Outputs(value=str(self.item))
+
+    class IterationSubworkflow(BaseWorkflow[MapNode.SubworkflowInputs, BaseState]):
+        graph = IterationNode
+
+        class Outputs(BaseOutputs):
+            value = IterationNode.Outputs.value
+
+    # AND a map node with items set to a dynamic field reference
+    class DynamicItemsMapNode(MapNode):
+        items = MyPromptNode.Outputs.json["items"]
+        subworkflow = IterationSubworkflow
+
+    # AND a workflow using these nodes
+    class DynamicFieldReferenceWorkflow(BaseWorkflow):
+        graph = MyPromptNode >> DynamicItemsMapNode
+
+        class Outputs(BaseOutputs):
+            result = DynamicItemsMapNode.Outputs.value
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=DynamicFieldReferenceWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the map node should be serialized
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    map_node = next(n for n in workflow_raw_data["nodes"] if n.get("type") == "MAP")
+
+    # AND the attributes should contain the items with a BINARY_EXPRESSION
+    attributes = map_node["attributes"]
+    items_attribute = next(attr for attr in attributes if attr["name"] == "items")
+    items_value = items_attribute["value"]
+    assert items_value["type"] == "BINARY_EXPRESSION"
+    assert items_value["operator"] == "accessField"
+    assert items_value["rhs"]["type"] == "CONSTANT_VALUE"
+    assert items_value["rhs"]["value"]["value"] == "items"

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

@@ -36,8 +36,6 @@ def test_serialize_workflow__await_all():
 
     # AND its raw data should be what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    assert len(workflow_raw_data["edges"]) == 6
-    assert len(workflow_raw_data["nodes"]) == 6
 
     # AND each node should be serialized correctly
     entrypoint_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "ENTRYPOINT")
@@ -93,71 +91,10 @@ def test_serialize_workflow__await_all():
         ignore_order_func=lambda x: x.path() == "root['data']['target_handles']",
     )
 
-    final_output_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "TERMINAL")
-    assert final_output_node == {
-        "id": "8187ce10-62b7-4a2c-8c0f-297387915467",
-        "type": "TERMINAL",
-        "data": {
-            "label": "Final Output",
-            "name": "value",
-            "target_handle_id": "ff55701c-16d3-4348-a633-6a298e71377d",
-            "output_id": "959ba00d-d30b-402e-8676-76efc4c3d2ae",
-            "output_type": "STRING",
-            "node_input_id": "fed9d343-6504-460c-968b-d3f9658193d0",
-        },
-        "base": {
-            "module": [
-                "vellum",
-                "workflows",
-                "nodes",
-                "displayable",
-                "final_output_node",
-                "node",
-            ],
-            "name": "FinalOutputNode",
-        },
-        "definition": None,
-        "inputs": [
-            {
-                "id": "fed9d343-6504-460c-968b-d3f9658193d0",
-                "key": "node_input",
-                "value": {
-                    "rules": [
-                        {
-                            "type": "NODE_OUTPUT",
-                            "data": {
-                                "node_id": "6c34a839-552b-436d-ba6a-501f663883c8",
-                                "output_id": "f6148c39-9ec1-4b10-9d24-d19bec9eeea8",
-                            },
-                        }
-                    ],
-                    "combinator": "OR",
-                },
-            }
-        ],
-        "display_data": {"position": {"x": 800.0, "y": -50.0}},
-    }
-
-    # AND each edge should be serialized correctly
-    serialized_edges = workflow_raw_data["edges"]
+    # AND each edge feeding into the merge node should be serialized correctly
+    merge_target_edges = [edge for edge in workflow_raw_data["edges"] if edge["target_node_id"] == merge_node["id"]]
     assert not DeepDiff(
         [
-            {
-                "id": "ea747507-68d5-4c72-8d00-a464ff7a55f1",
-                "source_node_id": "dc8aecd0-49ba-4464-a45f-29d3bfd686e4",
-                "source_handle_id": "017d40f5-8326-4e42-a409-b08995defaa8",
-                "target_node_id": "0871708d-8f05-4bc8-b3fb-a8624dae51de",
-                "target_handle_id": "01543a95-6bcb-46e2-a16c-570008972b88",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "2c28f06b-4f4a-4531-8cf9-860089ce1cc2",
-                "source_node_id": "dc8aecd0-49ba-4464-a45f-29d3bfd686e4",
-                "source_handle_id": "017d40f5-8326-4e42-a409-b08995defaa8",
-                "target_node_id": "0306d2a2-8e2a-49d1-bc4d-4026fbd98c4c",
-                "target_handle_id": "5b10f22e-ae81-4ca5-a682-12d3c3ec73c1",
-                "type": "DEFAULT",
-            },
             {
                 "id": "3870f290-8da7-43a8-b875-60510c060380",
                 "source_node_id": "0306d2a2-8e2a-49d1-bc4d-4026fbd98c4c",
@@ -174,24 +111,8 @@ def test_serialize_workflow__await_all():
                 "target_handle_id": "6da3e50a-8c6d-4de1-8ee9-da26f7c9552f",
                 "type": "DEFAULT",
             },
-            {
-                "id": "5344fbf3-6a71-4890-93a6-c6a2c5a4e50c",
-                "source_node_id": "f07c263c-65a3-4b58-83c1-f4a29123f167",
-                "source_handle_id": "da1bdfe9-8e99-4d06-842f-a76af95a713a",
-                "target_node_id": "6c34a839-552b-436d-ba6a-501f663883c8",
-                "target_handle_id": "413c3699-a37b-49cd-93fd-fe6d041a0119",
-                "type": "DEFAULT",
-            },
-            {
-                "id": "3d031c93-09b1-4937-9f98-c30a7ba6823d",
-                "source_node_id": "6c34a839-552b-436d-ba6a-501f663883c8",
-                "source_handle_id": "48beca40-00f3-4273-b2ea-5e1ac6494839",
-                "target_node_id": "8187ce10-62b7-4a2c-8c0f-297387915467",
-                "target_handle_id": "ff55701c-16d3-4348-a633-6a298e71377d",
-                "type": "DEFAULT",
-            },
         ],
-        serialized_edges,
+        merge_target_edges,
         ignore_order=True,
     )