vellum-ai 0.14.72__py3-none-any.whl → 0.14.74__py3-none-any.whl

vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py

@@ -2,6 +2,7 @@ import pytest
 from uuid import UUID
 from typing import Type
 
+from vellum.client.types.prompt_parameters import PromptParameters
 from vellum.client.types.variable_prompt_block import VariablePromptBlock
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs import BaseInputs
@@ -280,3 +281,104 @@ def test_serialize_node__port_groups():
 
     # AND the legacy source_handle_id should be the default port
     assert my_prompt_node["data"]["source_handle_id"] == "149d97a4-3da3-44a9-95f7-ea7b8d38b877"
+
+
+def test_serialize_node__prompt_parameters__dynamic_references():
+    # GIVEN input definition
+    class MyInputs(BaseInputs):
+        input_value: str
+
+    # AND a prompt node with PromptParameters containing dynamic references
+    class DynamicPromptNode(InlinePromptNode):
+        blocks = []
+        ml_model = "gpt-4o"
+        parameters = PromptParameters(custom_parameters={"json_schema": MyInputs.input_value})
+
+    # AND a workflow with the prompt node
+    class Workflow(BaseWorkflow[MyInputs, BaseState]):
+        graph = DynamicPromptNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the PromptParameters
+    dynamic_prompt_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(DynamicPromptNode.__id__)
+    )
+
+    # AND the parameters should be properly serialized in exec_config
+    exec_config = dynamic_prompt_node["data"]["exec_config"]
+    assert "parameters" in exec_config
+
+    parameters = exec_config["parameters"]
+    assert parameters == {}
+
+    # AND the parameters should also be serialized in the attributes array
+    parameters_attribute = next(
+        (attr for attr in dynamic_prompt_node.get("attributes", []) if attr["name"] == "parameters"), None
+    )
+    assert parameters_attribute is not None
+    assert parameters_attribute["name"] == "parameters"
+    assert parameters_attribute["value"]["type"] == "DICTIONARY_REFERENCE"
+    assert parameters_attribute["value"]["entries"] == [
+        {
+            "id": "6b63ff96-a2eb-4c6e-bad1-bde01605fa86",
+            "key": "stop",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+        },
+        {
+            "id": "265a1c17-2089-4ac1-b2ce-361b6b9a3335",
+            "key": "temperature",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+        },
+        {
+            "id": "699976ec-8ec2-476a-a011-7cf810a8a307",
+            "key": "max_tokens",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+        },
+        {
+            "id": "a87e23da-9794-41ff-ba80-c3a77e976e75",
+            "key": "top_p",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+        },
+        {
+            "id": "18eb53c2-ec1a-4115-9f21-083af430df67",
+            "key": "top_k",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+        },
+        {
+            "id": "295509a2-5837-452c-893d-f47b67c63c8a",
+            "key": "frequency_penalty",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+        },
+        {
+            "id": "5fc64379-5566-426a-a909-dd56c3305aa5",
+            "key": "presence_penalty",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+        },
+        {
+            "id": "5d326da0-c096-4425-8bf1-3a18764e96e3",
+            "key": "logit_bias",
+            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+        },
+        {
+            "id": "cd1a0e1b-6667-48a0-9964-257e1ec8851d",
+            "key": "custom_parameters",
+            "value": {
+                "entries": [
+                    {
+                        "id": "a9a3092e-dd18-4533-b6b5-24588ebd8f7f",
+                        "key": "json_schema",
+                        "value": {
+                            "input_variable_id": "c02d1201-86d1-4364-b3b3-4fc6824db8a4",
+                            "type": "WORKFLOW_INPUT",
+                        },
+                    }
+                ],
+                "type": "DICTIONARY_REFERENCE",
+            },
+        },
+    ]

vellum_ee/workflows/display/tests/test_base_workflow_display.py

@@ -39,7 +39,7 @@ def test_base_workflow_display__serialize_empty_workflow():
                     "data": {"label": "Entrypoint Node", "source_handle_id": "0af025a4-3b25-457d-a7ae-e3a7ba15c86c"},
                     "base": None,
                     "definition": None,
-                    "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                    "display_data": {"position": {"x": 0.0, "y": -50.0}},
                     "id": "3c41cdd9-999a-48b8-9088-f6dfa1369bfd",
                     "inputs": [],
                     "type": "ENTRYPOINT",

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

@@ -249,7 +249,7 @@ def test_serialize_node__stacked():
             "id": "074833b0-e142-4bbc-8dec-209a35e178a3",
             "label": "test_serialize_node__stacked.<locals>.InnerStackedGenericNode",
             "type": "GENERIC",
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
             "definition": {
                 "name": "InnerStackedGenericNode",

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

@@ -1,7 +1,10 @@
+import pytest
+from dataclasses import dataclass
 from uuid import uuid4
 from typing import List
 
 from deepdiff import DeepDiff
+from pydantic import BaseModel
 
 from vellum.client.types.chat_message import ChatMessage
 from vellum.workflows.inputs.base import BaseInputs
@@ -70,6 +73,28 @@ def test_serialize_node__constant_value(serialize_node):
     )
 
 
+@pytest.mark.parametrize(
+    "boolean_value, expected_value",
+    [
+        (True, True),
+        (False, False),
+    ],
+)
+def test_serialize_node__constant_boolean_value(serialize_node, boolean_value, expected_value):
+    class BooleanValueGenericNode(BaseNode):
+        attr: bool = boolean_value
+
+    serialized_node = serialize_node(BooleanValueGenericNode)
+
+    assert serialized_node["attributes"][0]["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {
+            "type": "JSON",
+            "value": expected_value,
+        },
+    }
+
+
 def test_serialize_node__constant_value_reference(serialize_node):
     class ConstantValueReferenceGenericNode(BaseNode):
         attr: str = ConstantValueReference("hello")
@@ -269,7 +294,7 @@ def test_serialize_node__workflow_input_as_nested_chat_history():
             "id": "11be9d37-0069-4695-a317-14a3b6519d4e",
             "label": "test_serialize_node__workflow_input_as_nested_chat_history.<locals>.GenericNode",
             "type": "GENERIC",
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
             "definition": {
                 "name": "GenericNode",
@@ -580,3 +605,68 @@ def test_serialize_node__coalesce(serialize_node):
         serialized_node,
         ignore_order=True,
     )
+
+
+def test_serialize_node__dataclass_with_node_output_reference(serialize_node):
+    @dataclass
+    class MyDataClass:
+        name: str
+        node_ref: str
+
+    class NodeWithOutput(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class NodeWithOutputDisplay(BaseNodeDisplay[NodeWithOutput]):
+        pass
+
+    class GenericNodeWithDataclass(BaseNode):
+        attr = MyDataClass(name="test", node_ref=NodeWithOutput.Outputs.result)
+
+    node_output_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=GenericNodeWithDataclass,
+        global_node_displays={NodeWithOutput: NodeWithOutputDisplay()},
+        global_node_output_displays={
+            NodeWithOutput.Outputs.result: (NodeWithOutput, NodeOutputDisplay(id=node_output_id, name="result"))
+        },
+    )
+
+    attr_value = serialized_node["attributes"][0]["value"]
+    assert attr_value["type"] == "DICTIONARY_REFERENCE"
+
+    assert any(
+        entry["key"] == "node_ref" and entry["value"]["type"] == "NODE_OUTPUT" for entry in attr_value["entries"]
+    )
+
+
+def test_serialize_node__pydantic_with_node_output_reference(serialize_node):
+    class MyPydanticModel(BaseModel):
+        name: str
+        node_ref: str
+
+    class NodeWithOutput(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class NodeWithOutputDisplay(BaseNodeDisplay[NodeWithOutput]):
+        pass
+
+    class GenericNodeWithPydantic(BaseNode):
+        attr = MyPydanticModel(name="test", node_ref=NodeWithOutput.Outputs.result)
+
+    node_output_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=GenericNodeWithPydantic,
+        global_node_displays={NodeWithOutput: NodeWithOutputDisplay()},
+        global_node_output_displays={
+            NodeWithOutput.Outputs.result: (NodeWithOutput, NodeOutputDisplay(id=node_output_id, name="result"))
+        },
+    )
+
+    attr_value = serialized_node["attributes"][0]["value"]
+    assert attr_value["type"] == "DICTIONARY_REFERENCE"
+
+    assert any(
+        entry["key"] == "node_ref" and entry["value"]["type"] == "NODE_OUTPUT" for entry in attr_value["entries"]
+    )

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

@@ -64,7 +64,7 @@ def test_serialize_workflow(vellum_client):
         "type": "ENTRYPOINT",
         "inputs": [],
         "data": {"label": "Entrypoint Node", "source_handle_id": "8eaa7f02-25ff-4a00-9b0a-5185718d89b3"},
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 0.0, "y": -50.0}},
         "base": None,
         "definition": None,
     }
@@ -184,7 +184,7 @@ def test_serialize_workflow(vellum_client):
                 "json_output_id": "12e4a99d-883d-4da5-aa51-35817d94013e",
                 "status_code_output_id": "fecc16c3-400e-4fd3-8223-08366070e3b1",
             },
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {
                 "name": "APINode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "api_node", "node"],
@@ -230,7 +230,7 @@ def test_serialize_workflow(vellum_client):
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": 200.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -266,7 +266,7 @@ def test_serialize_workflow(vellum_client):
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": -50.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -302,7 +302,7 @@ def test_serialize_workflow(vellum_client):
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": -300.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],

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

@@ -50,7 +50,7 @@ def test_serialize_workflow_with_filepath():
         "type": "ENTRYPOINT",
         "inputs": [],
         "data": {"label": "Entrypoint Node", "source_handle_id": "118e4298-aa79-467c-b8b4-2df540905e86"},
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 0.0, "y": -50.0}},
         "base": None,
         "definition": None,
     }
@@ -94,7 +94,7 @@ def test_serialize_workflow_with_filepath():
             "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
             "log_output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
         },
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 200.0, "y": -50.0}},
         "base": {
             "module": [
                 "vellum",
@@ -143,7 +143,7 @@ def test_serialize_workflow_with_filepath():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": -175.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -179,7 +179,7 @@ def test_serialize_workflow_with_filepath():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": 75.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -289,7 +289,7 @@ def test_serialize_workflow_with_code():
         "type": "ENTRYPOINT",
         "inputs": [],
         "data": {"label": "Entrypoint Node", "source_handle_id": "e82390bb-c68c-48c1-9f87-7fbfff494c45"},
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 0.0, "y": -50.0}},
         "base": None,
         "definition": None,
     }
@@ -333,7 +333,7 @@ def test_serialize_workflow_with_code():
             "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
             "log_output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
         },
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 200.0, "y": -50.0}},
         "base": {
             "name": "CodeExecutionNode",
             "module": ["vellum", "workflows", "nodes", "displayable", "code_execution_node", "node"],
@@ -375,7 +375,7 @@ def test_serialize_workflow_with_code():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": -175.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -411,7 +411,7 @@ def test_serialize_workflow_with_code():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": 75.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -522,7 +522,7 @@ def test_serialize_workflow__try_wrapped():
         "base": None,
         "definition": None,
         "display_data": {
-            "position": {"x": 0.0, "y": 0.0},
+            "position": {"x": 0.0, "y": -50.0},
         },
     }
 
@@ -565,7 +565,7 @@ def test_serialize_workflow__try_wrapped():
             "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
             "log_output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
         },
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 200.0, "y": -50.0}},
         "base": {
             "module": [
                 "vellum",
@@ -637,7 +637,7 @@ def test_serialize_workflow__try_wrapped():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": 75.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -673,7 +673,7 @@ def test_serialize_workflow__try_wrapped():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 400.0, "y": -175.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],

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

@@ -111,7 +111,7 @@ def test_serialize_workflow():
             "source_handle_id": "c2f0871d-0d9d-417f-8b0e-c813ccf880ac",
         },
         "display_data": {
-            "position": {"x": 0.0, "y": 0.0},
+            "position": {"x": 0.0, "y": -50.0},
         },
     }
 
@@ -415,7 +415,7 @@ def test_serialize_workflow():
                 ],
                 "version": "2",
             },
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {
                 "name": "ConditionalNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "conditional_node", "node"],
@@ -544,7 +544,7 @@ def test_serialize_workflow():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 600.0, "y": -50.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -580,7 +580,7 @@ def test_serialize_workflow():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 600.0, "y": -550.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -616,7 +616,7 @@ def test_serialize_workflow():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 600.0, "y": 200.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -652,7 +652,7 @@ def test_serialize_workflow():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 600.0, "y": -300.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -688,7 +688,7 @@ def test_serialize_workflow():
                         },
                     }
                 ],
-                "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                "display_data": {"position": {"x": 600.0, "y": 450.0}},
                 "base": {
                     "name": "FinalOutputNode",
                     "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -943,7 +943,7 @@ def test_conditional_node_serialize_all_operators_with_lhs_and_rhs(descriptor, o
                 ],
                 "version": "2",
             },
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {
                 "name": "ConditionalNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "conditional_node", "node"],
@@ -1054,7 +1054,7 @@ def test_conditional_node_serialize_all_operators_with_expression(descriptor, op
                 ],
                 "version": "2",
             },
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {
                 "name": "ConditionalNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "conditional_node", "node"],
@@ -1177,7 +1177,7 @@ def test_conditional_node_serialize_all_operators_with_value_and_start_and_end(d
                 ],
                 "version": "2",
             },
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 200.0, "y": -50.0}},
             "base": {
                 "name": "ConditionalNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "conditional_node", "node"],

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

@@ -93,7 +93,7 @@ def test_serialize_workflow():
         "base": None,
         "definition": None,
         "display_data": {
-            "position": {"x": 0.0, "y": 0.0},
+            "position": {"x": 0.0, "y": -50.0},
         },
     }
 
@@ -133,7 +133,7 @@ def test_serialize_workflow():
                     },
                 }
             ],
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 400.0, "y": -175.0}},
             "base": {
                 "name": "FinalOutputNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -173,7 +173,7 @@ def test_serialize_workflow():
                     },
                 }
             ],
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 400.0, "y": 75.0}},
             "base": {
                 "name": "FinalOutputNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],

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

@@ -68,7 +68,7 @@ def test_serialize_workflow():
             "label": "Entrypoint Node",
             "source_handle_id": "7d86498b-84ed-4feb-8e62-2188058c2c4e",
         },
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 0.0, "y": -50.0}},
         "base": None,
         "definition": None,
     }
@@ -105,7 +105,7 @@ def test_serialize_workflow():
                 "target_handle_id": "70c19f1c-309c-4a5d-ba65-664c0bb2fedf",
                 "error_source_input_id": "8e4c8d76-2e02-4d7e-a7bf-d71af392dc49",
             },
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 400.0, "y": 75.0}},
             "base": {
                 "name": "ErrorNode",
                 "module": ["vellum", "workflows", "nodes", "core", "error_node", "node"],
@@ -153,7 +153,7 @@ def test_serialize_workflow():
                     },
                 }
             ],
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 600.0, "y": -50.0}},
             "base": {
                 "name": "FinalOutputNode",
                 "module": [

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

@@ -56,15 +56,15 @@ def test_serialize_workflow():
 
     # AND each node should be serialized correctly
     entrypoint_node = workflow_raw_data["nodes"][0]
-    assert entrypoint_node == {
-        "id": "f1e4678f-c470-400b-a40e-c8922cc99a86",
-        "type": "ENTRYPOINT",
-        "inputs": [],
-        "data": {"label": "Entrypoint Node", "source_handle_id": "40201804-8beb-43ad-8873-a027759512f1"},
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
-        "base": None,
-        "definition": None,
+    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"] == "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f"
@@ -100,7 +100,7 @@ def test_serialize_workflow():
                     },
                 }
             ],
-            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "display_data": {"position": {"x": 400.0, "y": -50.0}},
             "base": {
                 "name": "FinalOutputNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -146,6 +146,17 @@ def test_serialize_workflow():
         }
     }
 
+    # AND the nodes should have been auto-positioned since they all started at (0,0)
+    nodes = workflow_raw_data["nodes"]
+    positions = [node["display_data"]["position"] for node in nodes]
+
+    nodes_at_zero = sum(1 for pos in positions if pos["x"] == 0.0 and pos["y"] == 0.0)
+    assert nodes_at_zero < len(nodes), "Auto layout should have positioned nodes away from (0,0)"
+
+    x_positions = [pos["x"] for pos in positions]
+    assert all(x >= 0 for x in x_positions), "All x positions should be non-negative"
+    assert len(set(x_positions)) > 1 or len(nodes) == 1, "Nodes should be spread across different x positions"
+
     # AND the definition should be what we expect
     definition = workflow_raw_data["definition"]
     assert definition == {

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

@@ -65,7 +65,7 @@ def test_serialize_workflow():
         "definition": None,
         "inputs": [],
         "data": {"label": "Entrypoint Node", "source_handle_id": "41840690-8d85-486e-a864-b0661ccf0f2e"},
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 0.0, "y": -50.0}},
     }
 
     guardrail_node = workflow_raw_data["nodes"][1]
@@ -108,7 +108,7 @@ def test_serialize_workflow():
             "metric_definition_id": "example_metric_definition",
             "release_tag": "LATEST",
         },
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 200.0, "y": -50.0}},
         "base": {
             "module": ["vellum", "workflows", "nodes", "displayable", "guardrail_node", "node"],
             "name": "GuardrailNode",
@@ -162,7 +162,7 @@ def test_serialize_workflow():
                 },
             }
         ],
-        "display_data": {"position": {"x": 0.0, "y": 0.0}},
+        "display_data": {"position": {"x": 400.0, "y": -50.0}},
     }
 
     # AND each edge should be serialized correctly