vellum-ai 0.14.0__tar.gz → 0.14.1__tar.gz

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.0
+Version: 0.14.1
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/base.py

@@ -31,6 +31,20 @@ WorkflowInputsDisplayType = TypeVar("WorkflowInputsDisplayType", bound=WorkflowI
 WorkflowInputsDisplayOverridesType = TypeVar("WorkflowInputsDisplayOverridesType", bound=WorkflowInputsDisplayOverrides)
 
 
+@dataclass
+class StateValueDisplayOverrides:
+    id: UUID
+
+
+@dataclass
+class StateValueDisplay(StateValueDisplayOverrides):
+    pass
+
+
+StateValueDisplayType = TypeVar("StateValueDisplayType", bound=StateValueDisplay)
+StateValueDisplayOverridesType = TypeVar("StateValueDisplayOverridesType", bound=StateValueDisplayOverrides)
+
+
 @dataclass
 class EdgeDisplayOverrides:
     id: UUID

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -34,6 +34,7 @@ from vellum.workflows.references import OutputReference
 from vellum.workflows.references.constant import ConstantValueReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.lazy import LazyReference
+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 JsonArray, JsonObject
@@ -44,6 +45,7 @@ from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator, primitive_to_vellum_value
 from vellum_ee.workflows.display.vellum import CodeResourceDefinition, GenericNodeDisplayData
 
@@ -343,28 +345,8 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             return self.serialize_value(display_context, value._value)
 
         if isinstance(value, LazyReference):
-            if isinstance(value._get, str):
-                reference_parts = value._get.split(".")
-                if len(reference_parts) < 3:
-                    raise Exception(
-                        f"Failed to parse lazy reference: {value._get}. Only Node Output references are supported."
-                    )
-
-                output_name = reference_parts[-1]
-                nested_class_name = reference_parts[-2]
-                if nested_class_name != "Outputs":
-                    raise Exception(
-                        f"Failed to parse lazy reference: {value._get}. Outputs are the only node reference supported."
-                    )
-
-                node_class_name = ".".join(reference_parts[:-2])
-                for node in display_context.node_displays.keys():
-                    if node.__name__ == node_class_name:
-                        return self.serialize_value(display_context, getattr(node.Outputs, output_name))
-
-                raise NotImplementedError(f"Failed to find a LazyReference for: {value._get}")
-
-            return self.serialize_value(display_context, value._get())
+            child_descriptor = get_child_descriptor(value, display_context)
+            return self.serialize_value(display_context, child_descriptor)
 
         if isinstance(value, WorkflowInputReference):
             workflow_input_display = display_context.global_workflow_input_displays[value]
@@ -373,6 +355,13 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 "input_variable_id": str(workflow_input_display.id),
             }
 
+        if isinstance(value, StateValueReference):
+            state_value_display = display_context.global_state_value_displays[value]
+            return {
+                "type": "STATE_VALUE",
+                "state_variable_id": str(state_value_display.id),
+            }
+
         if isinstance(value, OutputReference):
             upstream_node, output_display = display_context.global_node_output_displays[value]
             upstream_node_display = display_context.global_node_displays[upstream_node]

vellum_ai-0.14.1/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py

@@ -0,0 +1,52 @@
+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.references.lazy import LazyReference
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+
+
+def test_serialize_node__lazy_reference_in_prompt_inputs():
+    # GIVEN a prompt node with a lazy reference in the prompt inputs
+    class LazyReferencePromptNode(InlinePromptNode):
+        prompt_inputs = {"attr": LazyReference[str]("OtherNode.Outputs.result")}
+        blocks = []
+        ml_model = "gpt-4o"
+
+    class OtherNode(BaseNode):
+        class Outputs:
+            result: str
+
+    # AND a workflow with both nodes
+    class Workflow(BaseWorkflow):
+        graph = LazyReferencePromptNode >> OtherNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the attribute reference
+    lazy_reference_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(LazyReferencePromptNode.__id__)
+    )
+
+    assert lazy_reference_node["inputs"] == [
+        {
+            "id": "fba6a4d5-835a-4e99-afb7-f6a4aed15110",
+            "key": "attr",
+            "value": {
+                "combinator": "OR",
+                "rules": [
+                    {
+                        "type": "NODE_OUTPUT",
+                        "data": {
+                            "node_id": str(OtherNode.__id__),
+                            "output_id": "7f377cb8-4eca-4f1c-9239-9925f9495d84",
+                        },
+                    }
+                ],
+            },
+        }
+    ]

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/test_vellum_workflow_display.py

@@ -25,6 +25,7 @@ def test_vellum_workflow_display__serialize_empty_workflow():
     # THEN it should return the expected config
     assert exec_config == {
         "input_variables": [],
+        "state_variables": [],
         "output_variables": [],
         "workflow_raw_data": {
             "definition": {

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py

@@ -4,10 +4,11 @@ from typing import Any, Dict, Tuple, Type
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.references.output import OutputReference
+from vellum.workflows.references.state_value import StateValueReference
 from vellum.workflows.references.workflow_input import WorkflowInputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
-from vellum_ee.workflows.display.base import WorkflowInputsDisplayType
+from vellum_ee.workflows.display.base import StateValueDisplayType, WorkflowInputsDisplayType
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
@@ -22,6 +23,7 @@ def serialize_node():
         node_class: Type[NodeType],
         base_class: type[BaseNodeDisplay[Any]] = BaseNodeDisplay,
         global_workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {},
+        global_state_value_displays: Dict[StateValueReference, StateValueDisplayType] = {},
         global_node_displays: Dict[Type[BaseNode], NodeDisplayType] = {},
         global_node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]] = {},
     ) -> JsonObject:
@@ -37,6 +39,7 @@ def serialize_node():
             ),
             node_displays={node_class: node_display},
             global_workflow_input_displays=global_workflow_input_displays,
+            global_state_value_displays=global_state_value_displays,
             global_node_displays=global_node_displays,
             global_node_output_displays=global_node_output_displays,
         )

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py

@@ -32,6 +32,7 @@ def test_serialize_workflow(vellum_client):
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py

@@ -21,6 +21,7 @@ def test_serialize_workflow_with_filepath():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 
@@ -261,6 +262,7 @@ def test_serialize_workflow_with_code():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 
@@ -489,6 +491,7 @@ def test_serialize_workflow__try_wrapped():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py

@@ -37,6 +37,7 @@ def test_serialize_workflow():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 
@@ -785,6 +786,7 @@ def test_conditional_node_serialize_all_operators_with_lhs_and_rhs(descriptor, o
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 
@@ -890,6 +892,7 @@ def test_conditional_node_serialize_all_operators_with_expression(descriptor, op
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 
@@ -982,6 +985,7 @@ def test_conditional_node_serialize_all_operators_with_value_and_start_and_end(d
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

vellum_ai-0.14.1/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_default_state_serialization.py

@@ -0,0 +1,243 @@
+from deepdiff import DeepDiff
+
+from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_default_state.workflow import BasicDefaultStateWorkflow
+
+
+def test_serialize_workflow():
+    # GIVEN a Workflow that has a simple state definition
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=BasicDefaultStateWorkflow
+    )
+
+    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",
+        "state_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "0bbe085c-31d3-4b48-b74d-d501f592da90",
+                "key": "example",
+                "type": "STRING",
+                "default": {"type": "STRING", "value": "hello"},
+                "required": True,
+                "extensions": {"color": None},
+            },
+        ],
+        input_variables,
+        ignore_order=True,
+    )
+
+    # AND its state variables should be what we expect
+    state_variables = serialized_workflow["state_variables"]
+    assert len(state_variables) == 1
+    assert not DeepDiff(
+        [
+            {
+                "id": "812ec99b-1859-4361-b795-228628657bac",
+                "key": "example",
+                "type": "NUMBER",
+                "default": {"type": "NUMBER", "value": 5.0},
+                "required": True,
+                "extensions": {"color": None},
+            },
+        ],
+        state_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": "6e7eeaa5-9559-4ae3-8606-e52ead5805a5",
+                "key": "example_input",
+                "type": "STRING",
+            },
+            {
+                "id": "e3ae0fe3-7590-4eac-b808-45901d82f2ba",
+                "key": "example_state",
+                "type": "NUMBER",
+            },
+        ],
+        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": "32684932-7c7c-4b1c-aed2-553de29bf3f7",
+        "type": "ENTRYPOINT",
+        "inputs": [],
+        "data": {
+            "label": "Entrypoint Node",
+            "source_handle_id": "e4136ee4-a51a-4ca3-9a3a-aa96f5de2347",
+        },
+        "base": None,
+        "definition": None,
+        "display_data": {
+            "position": {"x": 0.0, "y": 0.0},
+        },
+    }
+
+    final_output_node_state = next(
+        n for n in workflow_raw_data["nodes"] if n["type"] == "TERMINAL" and n["data"]["name"] == "example_state"
+    )
+    final_output_node_input = next(
+        n for n in workflow_raw_data["nodes"] if n["type"] == "TERMINAL" and n["data"]["name"] == "example_input"
+    )
+    assert not DeepDiff(
+        {
+            "id": "27fdaa45-b8ce-464d-be50-cf71cc56bc10",
+            "type": "TERMINAL",
+            "data": {
+                "label": "Final Output",
+                "name": "example_state",
+                "target_handle_id": "e7a09eb2-c9fb-4d57-b436-9cd9384c8960",
+                "output_id": "e3ae0fe3-7590-4eac-b808-45901d82f2ba",
+                "output_type": "NUMBER",
+                "node_input_id": "8de6a408-cf76-4d04-9845-f75211b611be",
+            },
+            "inputs": [
+                {
+                    "id": "8de6a408-cf76-4d04-9845-f75211b611be",
+                    "key": "node_input",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "NODE_OUTPUT",
+                                "data": {
+                                    "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
+                                    "output_id": "84b59a1a-82bf-46bb-9826-9b393402d0fe",
+                                },
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {
+                "name": "FinalOutputNode",
+                "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
+            },
+            "definition": None,
+        },
+        final_output_node_state,
+    )
+
+    assert not DeepDiff(
+        {
+            "id": "ca8bb585-c9a8-4bf7-bf9d-534b600fe23b",
+            "type": "TERMINAL",
+            "data": {
+                "label": "Final Output",
+                "name": "example_input",
+                "target_handle_id": "8a4a7efd-0e18-43ed-ba32-803e22e3ba0a",
+                "output_id": "6e7eeaa5-9559-4ae3-8606-e52ead5805a5",
+                "output_type": "STRING",
+                "node_input_id": "796b4a0b-da10-403a-acc3-8ebd3ebd3667",
+            },
+            "inputs": [
+                {
+                    "id": "796b4a0b-da10-403a-acc3-8ebd3ebd3667",
+                    "key": "node_input",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "NODE_OUTPUT",
+                                "data": {
+                                    "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
+                                    "output_id": "f305c61e-6e8f-4cea-a53e-92656136b545",
+                                },
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {
+                "name": "FinalOutputNode",
+                "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
+            },
+            "definition": None,
+        },
+        final_output_node_input,
+    )
+
+    # AND each edge should be serialized correctly
+    serialized_edges = workflow_raw_data["edges"]
+    assert not DeepDiff(
+        [
+            {
+                "id": "26003bf1-b5a3-41f4-a419-c7ef5a7595df",
+                "source_node_id": "32684932-7c7c-4b1c-aed2-553de29bf3f7",
+                "source_handle_id": "e4136ee4-a51a-4ca3-9a3a-aa96f5de2347",
+                "target_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
+                "target_handle_id": "a95a34f2-e894-4fb6-a2c9-15d12c1e3135",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "b0a57a5f-a1e4-4dc9-85dd-946f08304738",
+                "source_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
+                "source_handle_id": "1e739e86-a285-4438-9725-a152c15a63e3",
+                "target_node_id": "ca8bb585-c9a8-4bf7-bf9d-534b600fe23b",
+                "target_handle_id": "8a4a7efd-0e18-43ed-ba32-803e22e3ba0a",
+                "type": "DEFAULT",
+            },
+            {
+                "id": "e4366583-94a5-40b0-9b6f-1e965695b1fe",
+                "source_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
+                "source_handle_id": "1e739e86-a285-4438-9725-a152c15a63e3",
+                "target_node_id": "27fdaa45-b8ce-464d-be50-cf71cc56bc10",
+                "target_handle_id": "e7a09eb2-c9fb-4d57-b436-9cd9384c8960",
+                "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": "BasicDefaultStateWorkflow",
+        "module": [
+            "tests",
+            "workflows",
+            "basic_default_state",
+            "workflow",
+        ],
+    }

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py

@@ -19,6 +19,7 @@ def test_serialize_workflow():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py

@@ -19,6 +19,7 @@ def test_serialize_workflow(vellum_client):
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py

@@ -16,7 +16,7 @@ def test_serialize_workflow():
     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"}
+    assert serialized_workflow.keys() == {"workflow_raw_data", "input_variables", "state_variables", "output_variables"}
 
     # AND its input variables should be what we expect
     input_variables = serialized_workflow["input_variables"]

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py

@@ -18,6 +18,7 @@ def test_serialize_workflow():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py

@@ -16,6 +16,7 @@ def test_serialize_workflow():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py

@@ -18,6 +18,7 @@ def test_serialize_workflow__await_all():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py

@@ -34,6 +34,7 @@ def test_serialize_workflow(vellum_client):
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_search_node_serialization.py

@@ -15,7 +15,7 @@ def test_serialize_workflow():
     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"}
+    assert serialized_workflow.keys() == {"workflow_raw_data", "input_variables", "state_variables", "output_variables"}
 
     # AND its input variables should be what we expect
     input_variables = serialized_workflow["input_variables"]

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_deployment_serialization.py

@@ -34,6 +34,7 @@ def test_serialize_workflow(vellum_client):
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py

@@ -20,6 +20,7 @@ def test_serialize_workflow():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py

@@ -16,6 +16,7 @@ def test_serialize_workflow():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py

@@ -16,6 +16,7 @@ def test_serialize_workflow():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py

@@ -22,6 +22,7 @@ def test_serialize_workflow__missing_final_output_node():
     assert serialized_workflow.keys() == {
         "workflow_raw_data",
         "input_variables",
+        "state_variables",
         "output_variables",
     }
 

{vellum_ai-0.14.0 → vellum_ai-0.14.1}/ee/vellum_ee/workflows/display/types.py

@@ -6,10 +6,11 @@ from vellum.client.core import UniversalBaseModel
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.ports import Port
-from vellum.workflows.references import OutputReference, WorkflowInputReference
+from vellum.workflows.references import OutputReference, StateValueReference, WorkflowInputReference
 from vellum_ee.workflows.display.base import (
     EdgeDisplayType,
     EntrypointDisplayType,
+    StateValueDisplayType,
     WorkflowInputsDisplayType,
     WorkflowMetaDisplayType,
     WorkflowOutputDisplayType,
@@ -41,6 +42,7 @@ class WorkflowDisplayContext(
     Generic[
         WorkflowMetaDisplayType,
         WorkflowInputsDisplayType,
+        StateValueDisplayType,
         NodeDisplayType,
         EntrypointDisplayType,
         WorkflowOutputDisplayType,
@@ -53,6 +55,8 @@ class WorkflowDisplayContext(
     global_workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = field(
         default_factory=dict
     )
+    state_value_displays: Dict[StateValueReference, StateValueDisplayType] = field(default_factory=dict)
+    global_state_value_displays: Dict[StateValueReference, StateValueDisplayType] = field(default_factory=dict)
     node_displays: Dict[Type[BaseNode], "NodeDisplayType"] = field(default_factory=dict)
     global_node_displays: Dict[Type[BaseNode], NodeDisplayType] = field(default_factory=dict)
     global_node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], "NodeOutputDisplay"]] = field(