vellum-ai 0.14.46__py3-none-any.whl → 0.14.48__py3-none-any.whl

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

@@ -2,9 +2,11 @@ import pytest
 from uuid import UUID
 from typing import Type
 
+from vellum.client.types.variable_prompt_block import VariablePromptBlock
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
+from vellum.workflows.ports.port import Port
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state.base import BaseState
 from vellum_ee.workflows.display.nodes.vellum.inline_prompt_node import BaseInlinePromptNodeDisplay
@@ -78,7 +80,7 @@ def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[Inl
 @pytest.mark.parametrize(
     ["GetDisplayClass", "expected_input_id"],
     [
-        (_no_display_class, "8aa4ce7f-5eb8-41b7-abd0-ea2b40c8fb88"),
+        (_no_display_class, "9b036991-67ff-4cd0-a4d7-b4ed581e8b6d"),
         (_display_class_with_node_input_ids_by_name, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
         (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
     ],
@@ -165,7 +167,7 @@ def test_serialize_node__prompt_inputs__state_reference():
             },
         },
         {
-            "id": "3750feb9-5d5c-4150-b62d-a9924f466888",
+            "id": "b83c40f7-0159-442f-af03-e80870363c52",
             "key": "bar",
             "value": {
                 "rules": [
@@ -181,3 +183,66 @@ def test_serialize_node__prompt_inputs__state_reference():
             },
         },
     ]
+
+
+def test_serialize_node__unreferenced_variable_block__still_serializes():
+    # GIVEN a prompt node with an unreferenced variable block
+    class MyPromptNode(InlinePromptNode):
+        blocks = [VariablePromptBlock(input_variable="foo")]
+
+    # AND a workflow with the prompt node
+    class MyWorkflow(BaseWorkflow):
+        graph = MyPromptNode
+
+    # WHEN the prompt node is serialized
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should skip the state reference input rule
+    assert serialized_workflow["workflow_raw_data"]["nodes"][1]["data"]["exec_config"]["prompt_template_block_data"][
+        "blocks"
+    ] == [
+        {
+            "id": "fecbb5f3-e0a3-42ed-9774-6c68fd5db50c",
+            "block_type": "VARIABLE",
+            "input_variable_id": "ea3f6348-8553-4375-bd27-527df4e4f3c2",
+            "state": "ENABLED",
+            "cache_config": None,
+        }
+    ]
+
+    # AND we should have a warning of the invalid reference
+    # TODO: Come up with a proposal for how nodes should propagate warnings
+    # warnings = list(workflow_display.errors)
+    # assert len(warnings) == 1
+    # assert "Missing input variable 'foo' for prompt block 0" in str(warnings[0])
+
+
+def test_serialize_node__port_groups():
+    # GIVEN a prompt node with ports
+    class MyPromptNode(InlinePromptNode):
+        class Ports(InlinePromptNode.Ports):
+            apple = Port.on_if(LazyReference(lambda: MyPromptNode.Outputs.text).equals("apple"))
+            banana = Port.on_if(LazyReference(lambda: MyPromptNode.Outputs.text).equals("banana"))
+
+    # AND a workflow with the prompt node
+    class MyWorkflow(BaseWorkflow):
+        graph = MyPromptNode
+
+    # WHEN the prompt node is serialized
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should have the ports serialized
+    my_prompt_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MyPromptNode.__id__)
+    )
+    ports = my_prompt_node["ports"]
+    assert len(ports) == 2
+    assert ports[0]["id"] == "149d97a4-3da3-44a9-95f7-ea7b8d38b877"
+    assert ports[1]["id"] == "71f2d2b3-194f-4492-bc1c-a5ca1f60fb0a"
+    assert ports[0]["name"] == "apple"
+    assert ports[1]["name"] == "banana"
+
+    # AND the legacy source_handle_id should be the default port
+    assert my_prompt_node["data"]["source_handle_id"] == "149d97a4-3da3-44a9-95f7-ea7b8d38b877"

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

@@ -4,7 +4,8 @@ from uuid import UUID, uuid4
 from typing import Type
 
 from vellum.workflows import BaseWorkflow
-from vellum.workflows.nodes import SubworkflowDeploymentNode
+from vellum.workflows.nodes.displayable.subworkflow_deployment_node.node import SubworkflowDeploymentNode
+from vellum_ee.workflows.display.nodes.vellum.subworkflow_deployment_node import BaseSubworkflowDeploymentNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
 
@@ -13,14 +14,14 @@ def _no_display_class(Node: Type[SubworkflowDeploymentNode]):  # type: ignore
 
 
 def _display_class_with_node_input_ids_by_name(Node: Type[SubworkflowDeploymentNode]):
-    class SubworkflowNodeDisplay(SubworkflowDeploymentNode[Node]):  # type: ignore[valid-type]
+    class SubworkflowNodeDisplay(BaseSubworkflowDeploymentNodeDisplay[Node]):  # type: ignore[valid-type]
         node_input_ids_by_name = {"foo": UUID("aff4f838-577e-44b9-ac5c-6d8213abbb9c")}
 
     return SubworkflowNodeDisplay
 
 
 def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[SubworkflowDeploymentNode]):
-    class SubworkflowNodeDisplay(SubworkflowDeploymentNode[Node]):  # type: ignore[valid-type]
+    class SubworkflowNodeDisplay(BaseSubworkflowDeploymentNodeDisplay[Node]):  # type: ignore[valid-type]
         node_input_ids_by_name = {"subworkflow_inputs.foo": UUID("aff4f838-577e-44b9-ac5c-6d8213abbb9c")}
 
     return SubworkflowNodeDisplay
@@ -55,7 +56,7 @@ def mock_fetch_deployment(mocker):
 @pytest.mark.parametrize(
     ["GetDisplayClass", "expected_input_id"],
     [
-        (_no_display_class, "aff4f838-577e-44b9-ac5c-6d8213abbb9c"),
+        (_no_display_class, "394132c2-1817-455e-9f3f-b7073eb63a2b"),
         (_display_class_with_node_input_ids_by_name, "aff4f838-577e-44b9-ac5c-6d8213abbb9c"),
         (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "aff4f838-577e-44b9-ac5c-6d8213abbb9c"),
     ],

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

@@ -29,7 +29,7 @@ def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[Tem
 @pytest.mark.parametrize(
     ["GetDisplayClass", "expected_input_id"],
     [
-        (_no_display_class, "d3519cec-590c-416d-8eb1-96051aed5ddd"),
+        (_no_display_class, "91d982a9-6c41-42ac-aff9-7b623c450a55"),
         (_display_class_with_node_input_ids_by_name, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
         (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "fba6a4d5-835a-4e99-afb7-f6a4aed15110"),
     ],

vellum_ee/workflows/display/tests/test_base_workflow_display.py

@@ -3,6 +3,8 @@ from typing import Dict
 
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import BaseNode
+from vellum.workflows.ports.port import Port
+from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state import BaseState
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.vellum import WorkflowInputsVellumDisplayOverrides
@@ -285,3 +287,45 @@ def test_vellum_workflow_display__serialize_with_parse_json_expression():
             "operator": "parseJson",
         },
     }
+
+
+def test_serialize__port_with_lazy_reference():
+    # GIVEN a node with a lazy reference in a Port
+    class MyNode(BaseNode):
+        class Ports(BaseNode.Ports):
+            foo = Port.on_if(LazyReference(lambda: MyNode.Outputs.bar))
+
+        class Outputs(BaseNode.Outputs):
+            bar: bool
+
+    # AND a workflow that uses the node
+    class Workflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    exec_config = workflow_display.serialize()
+
+    # THEN the lazy reference should be serialized correctly
+    raw_data = exec_config["workflow_raw_data"]
+    assert isinstance(raw_data, dict)
+
+    nodes = raw_data["nodes"]
+    assert isinstance(nodes, list)
+
+    my_node = nodes[1]
+    assert isinstance(my_node, dict)
+    ports = my_node.get("ports")
+    assert isinstance(ports, list)
+    assert ports == [
+        {
+            "id": "6c26bc2b-6469-47c1-b858-d63f0d311ea6",
+            "name": "foo",
+            "type": "IF",
+            "expression": {
+                "type": "NODE_OUTPUT",
+                "node_id": str(MyNode.__id__),
+                "node_output_id": str(MyNode.__output_ids__["bar"]),
+            },
+        }
+    ]

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

@@ -87,7 +87,7 @@ def test_serialize_workflow():
             "type": "ERROR",
             "inputs": [
                 {
-                    "id": "690d825f-6ffd-493e-8141-c86d384e6150",
+                    "id": "8e4c8d76-2e02-4d7e-a7bf-d71af392dc49",
                     "key": "error_source_input_id",
                     "value": {
                         "rules": [
@@ -101,11 +101,9 @@ def test_serialize_workflow():
                 }
             ],
             "data": {
-                "name": "error-node",
                 "label": "Fail Node",
                 "target_handle_id": "70c19f1c-309c-4a5d-ba65-664c0bb2fedf",
-                "error_source_input_id": "None",
-                "error_output_id": "None",
+                "error_source_input_id": "8e4c8d76-2e02-4d7e-a7bf-d71af392dc49",
             },
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {

vellum_ee/workflows/display/utils/expressions.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Dict, cast
 
 from vellum.client.types.logical_operator import LogicalOperator
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -37,7 +37,7 @@ from vellum.workflows.references.output import OutputReference
 from vellum.workflows.references.state_value import StateValueReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.references.workflow_input import WorkflowInputReference
-from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
 
 if TYPE_CHECKING:
@@ -116,7 +116,7 @@ def get_child_descriptor(value: LazyReference, display_context: "WorkflowDisplay
     return value._get()
 
 
-def serialize_condition(display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
+def _serialize_condition(display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
     if isinstance(
         condition,
         (
@@ -235,6 +235,33 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
             "node_id": str(node_class_display.node_id),
         }
 
+    if isinstance(value, list):
+        serialized_items = [serialize_value(display_context, item) for item in value]
+        if all(isinstance(item, dict) and item["type"] == "CONSTANT_VALUE" for item in serialized_items):
+            constant_values = []
+            for item in serialized_items:
+                item_dict = cast(Dict[str, Any], item)
+                value_inner = item_dict["value"]
+
+                if value_inner["type"] == "JSON" and "items" in value_inner:
+                    # Nested JSON list
+                    constant_values.append(value_inner["items"])
+                else:
+                    constant_values.append(value_inner["value"])
+
+            return {
+                "type": "CONSTANT_VALUE",
+                "value": {
+                    "type": "JSON",
+                    "items": constant_values,
+                },
+            }
+        else:
+            return {
+                "type": "ARRAY_REFERENCE",
+                "items": cast(JsonArray, serialized_items),  # list[JsonObject] -> JsonArray
+            }
+
     if isinstance(value, dict) and any(isinstance(v, BaseDescriptor) for v in value.values()):
         raise ValueError("Nested references are not supported.")
 
@@ -247,4 +274,4 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
 
     # If it's not any of the references we know about,
     # then try to serialize it as a nested value
-    return serialize_condition(display_context, value)
+    return _serialize_condition(display_context, value)

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -425,3 +425,165 @@ def test_serialize_workflow__nested_lazy_reference():
         "node_id": str(OuterNode.__id__),
         "node_output_id": str(OuterNode.__output_ids__["bar"]),
     }
+
+
+def test_serialize_workflow__array_values():
+    # GIVEN a node with array and nested array values
+    class MyNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            array_value = ["item1", "item2", "item3"]
+            nested_array_value = [["item1", "item2", "item3"], ["item4", "item5", "item6"]]
+            mixed_array_value = [["item1"], "item2", "item3"]
+
+    # AND a workflow that uses these outputs
+    class MyWorkflow(BaseWorkflow):
+        graph = MyNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            array_output = MyNode.Outputs.array_value
+            nested_array_output = MyNode.Outputs.nested_array_value
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN it should properly serialize the array and dictionary values
+    assert isinstance(data["workflow_raw_data"], dict)
+    assert isinstance(data["workflow_raw_data"]["nodes"], list)
+    raw_nodes = data["workflow_raw_data"]["nodes"]
+    generic_nodes = [node for node in raw_nodes if isinstance(node, dict) and node["type"] == "GENERIC"]
+    assert len(generic_nodes) > 0
+    my_node = generic_nodes[0]
+
+    outputs = my_node["outputs"]
+    assert isinstance(outputs, list)
+
+    array_outputs = [val for val in outputs if isinstance(val, dict) and val["name"] == "array_value"]
+    assert len(array_outputs) > 0
+    array_output = array_outputs[0]
+
+    assert isinstance(array_output, dict)
+    assert "value" in array_output
+    assert array_output["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {"type": "JSON", "items": ["item1", "item2", "item3"]},
+    }
+
+    nested_array_outputs = [val for val in outputs if isinstance(val, dict) and val["name"] == "nested_array_value"]
+    assert len(nested_array_outputs) > 0
+    nested_array_output = nested_array_outputs[0]
+
+    assert isinstance(nested_array_output, dict)
+    assert "value" in nested_array_output
+    assert nested_array_output["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {"type": "JSON", "items": [["item1", "item2", "item3"], ["item4", "item5", "item6"]]},
+    }
+
+    mixed_array_outputs = [val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_array_value"]
+    assert len(mixed_array_outputs) > 0
+    mixed_array_output = mixed_array_outputs[0]
+
+    assert isinstance(mixed_array_output, dict)
+    assert "value" in mixed_array_output
+    assert mixed_array_output["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {"type": "JSON", "items": [["item1"], "item2", "item3"]},
+    }
+
+
+def test_serialize_workflow__array_reference():
+    # GIVEN a node with array containing non-constant values (node references)
+    class FirstNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            value1: str
+            value2: str
+
+    class SecondNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            # Array containing a mix of constants and node references
+            mixed_array = ["constant1", FirstNode.Outputs.value1, "constant2", FirstNode.Outputs.value2]
+            mixed_nested_array = [["constant1", FirstNode.Outputs.value1], ["constant2", FirstNode.Outputs.value2]]
+
+    # AND a workflow that uses these outputs
+    class MyWorkflow(BaseWorkflow):
+        graph = FirstNode >> SecondNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            mixed_array_output = SecondNode.Outputs.mixed_array
+            mixed_nested_array_output = SecondNode.Outputs.mixed_nested_array
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN it should serialize as an ARRAY_REFERENCE
+    assert isinstance(data["workflow_raw_data"], dict)
+    assert isinstance(data["workflow_raw_data"]["nodes"], list)
+    assert len(data["workflow_raw_data"]["nodes"]) == 5
+    second_node = data["workflow_raw_data"]["nodes"][2]
+    assert isinstance(second_node, dict)
+
+    assert "outputs" in second_node
+    assert isinstance(second_node["outputs"], list)
+    outputs = second_node["outputs"]
+
+    mixed_array_outputs = [val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_array"]
+    assert len(mixed_array_outputs) > 0
+    mixed_array_output = mixed_array_outputs[0]
+
+    assert isinstance(mixed_array_output, dict)
+    assert "value" in mixed_array_output
+    assert mixed_array_output["value"] == {
+        "type": "ARRAY_REFERENCE",
+        "items": [
+            {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant1"}},
+            {
+                "type": "NODE_OUTPUT",
+                "node_id": "702a08b5-61e8-4a7a-a83d-77f49e39c5be",
+                "node_output_id": "419b6afa-fab5-493a-ba1e-4606f4641616",
+            },
+            {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant2"}},
+            {
+                "type": "NODE_OUTPUT",
+                "node_id": "702a08b5-61e8-4a7a-a83d-77f49e39c5be",
+                "node_output_id": "d1cacc41-478d-49a3-a6b3-1ba2d51291e2",
+            },
+        ],
+    }
+
+    mixed_nested_array_outputs = [
+        val for val in outputs if isinstance(val, dict) and val["name"] == "mixed_nested_array"
+    ]
+    assert len(mixed_nested_array_outputs) > 0
+    mixed_nested_array_output = mixed_nested_array_outputs[0]
+
+    assert isinstance(mixed_nested_array_output, dict)
+    assert "value" in mixed_nested_array_output
+    assert mixed_nested_array_output["value"] == {
+        "type": "ARRAY_REFERENCE",
+        "items": [
+            {
+                "type": "ARRAY_REFERENCE",
+                "items": [
+                    {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant1"}},
+                    {
+                        "type": "NODE_OUTPUT",
+                        "node_id": "702a08b5-61e8-4a7a-a83d-77f49e39c5be",
+                        "node_output_id": "419b6afa-fab5-493a-ba1e-4606f4641616",
+                    },
+                ],
+            },
+            {
+                "type": "ARRAY_REFERENCE",
+                "items": [
+                    {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "constant2"}},
+                    {
+                        "type": "NODE_OUTPUT",
+                        "node_id": "702a08b5-61e8-4a7a-a83d-77f49e39c5be",
+                        "node_output_id": "d1cacc41-478d-49a3-a6b3-1ba2d51291e2",
+                    },
+                ],
+            },
+        ],
+    }