vellum-ai 0.13.28__py3-none-any.whl → 0.14.1__py3-none-any.whl

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_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_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_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_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_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_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_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_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_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_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_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_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_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_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(

vellum_ee/workflows/display/utils/expressions.py

@@ -0,0 +1,26 @@
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.references.lazy import LazyReference
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+
+
+def get_child_descriptor(value: LazyReference, display_context: WorkflowDisplayContext) -> BaseDescriptor:
+    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 getattr(node.Outputs, output_name)
+
+        raise Exception(f"Failed to parse lazy reference: {value._get}")
+
+    return value._get()

vellum_ee/workflows/display/utils/vellum.py

@@ -32,10 +32,12 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
+from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.references.node import NodeReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 from vellum.workflows.vellum_client import create_vellum_client
+from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 from vellum_ee.workflows.display.vellum import (
     ConstantValuePointer,
     ExecutionCounterData,
@@ -88,6 +90,9 @@ def create_node_input_value_pointer_rule(
         return NodeOutputPointer(
             data=NodeOutputData(node_id=str(upstream_node_display.node_id), output_id=str(output_display.id)),
         )
+    if isinstance(value, LazyReference):
+        child_descriptor = get_child_descriptor(value, display_context)
+        return create_node_input_value_pointer_rule(child_descriptor, display_context)
     if isinstance(value, WorkflowInputReference):
         workflow_input_display = display_context.global_workflow_input_displays[value]
         return InputVariablePointer(data=InputVariableData(input_variable_id=str(workflow_input_display.id)))

vellum_ee/workflows/display/vellum.py

@@ -14,6 +14,8 @@ from vellum_ee.workflows.display.base import (
     EdgeDisplayOverrides,
     EntrypointDisplay,
     EntrypointDisplayOverrides,
+    StateValueDisplay,
+    StateValueDisplayOverrides,
     WorkflowInputsDisplay,
     WorkflowInputsDisplayOverrides,
     WorkflowMetaDisplay,
@@ -84,6 +86,18 @@ class WorkflowInputsVellumDisplay(WorkflowInputsVellumDisplayOverrides):
     pass
 
 
+@dataclass
+class StateValueVellumDisplayOverrides(StateValueDisplay, StateValueDisplayOverrides):
+    name: Optional[str] = None
+    required: Optional[bool] = None
+    color: Optional[str] = None
+
+
+@dataclass
+class StateValueVellumDisplay(StateValueVellumDisplayOverrides):
+    pass
+
+
 @dataclass
 class EdgeVellumDisplayOverrides(EdgeDisplay, EdgeDisplayOverrides):
     pass

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -13,7 +13,7 @@ from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.utils import get_wrapped_node
 from vellum.workflows.ports import Port
-from vellum.workflows.references import OutputReference, WorkflowInputReference
+from vellum.workflows.references import OutputReference, StateValueReference, WorkflowInputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.utils.uuids import uuid4_from_hash
@@ -22,6 +22,8 @@ from vellum_ee.workflows.display.base import (
     EdgeDisplayType,
     EntrypointDisplayOverridesType,
     EntrypointDisplayType,
+    StateValueDisplayOverridesType,
+    StateValueDisplayType,
     WorkflowInputsDisplayOverridesType,
     WorkflowInputsDisplayType,
     WorkflowMetaDisplayOverridesType,
@@ -49,6 +51,8 @@ class BaseWorkflowDisplay(
         WorkflowMetaDisplayOverridesType,
         WorkflowInputsDisplayType,
         WorkflowInputsDisplayOverridesType,
+        StateValueDisplayType,
+        StateValueDisplayOverridesType,
         NodeDisplayType,
         EntrypointDisplayType,
         EntrypointDisplayOverridesType,
@@ -64,6 +68,9 @@ class BaseWorkflowDisplay(
     # Used to explicitly specify display data for a workflow's inputs.
     inputs_display: Dict[WorkflowInputReference, WorkflowInputsDisplayOverridesType] = {}
 
+    # Used to explicitly specify display data for a workflow's state values.
+    state_value_displays: Dict[StateValueReference, StateValueDisplayOverridesType] = {}
+
     # Used to explicitly specify display data for a workflow's entrypoints.
     entrypoint_displays: Dict[Type[BaseNode], EntrypointDisplayOverridesType] = {}
 
@@ -91,6 +98,7 @@ class BaseWorkflowDisplay(
             WorkflowDisplayContext[
                 WorkflowMetaDisplayType,
                 WorkflowInputsDisplayType,
+                StateValueDisplayType,
                 NodeDisplayType,
                 EntrypointDisplayType,
                 WorkflowOutputDisplayType,
@@ -191,6 +199,7 @@ class BaseWorkflowDisplay(
     ) -> WorkflowDisplayContext[
         WorkflowMetaDisplayType,
         WorkflowInputsDisplayType,
+        StateValueDisplayType,
         NodeDisplayType,
         EntrypointDisplayType,
         WorkflowOutputDisplayType,
@@ -244,6 +253,18 @@ class BaseWorkflowDisplay(
             workflow_input_displays[workflow_input] = input_display
             global_workflow_input_displays[workflow_input] = input_display
 
+        state_value_displays: Dict[StateValueReference, StateValueDisplayType] = {}
+        global_state_value_displays = (
+            copy(self._parent_display_context.global_state_value_displays) if self._parent_display_context else {}
+        )
+        for state_value in self._workflow.get_state_class():
+            state_value_display_overrides = self.state_value_displays.get(state_value)
+            state_value_display = self._generate_state_value_display(
+                state_value, overrides=state_value_display_overrides
+            )
+            state_value_displays[state_value] = state_value_display
+            global_state_value_displays[state_value] = state_value_display
+
         entrypoint_displays: Dict[Type[BaseNode], EntrypointDisplayType] = {}
         for entrypoint in self._workflow.get_entrypoints():
             if entrypoint in entrypoint_displays:
@@ -286,6 +307,8 @@ class BaseWorkflowDisplay(
             workflow_display=workflow_display,
             workflow_input_displays=workflow_input_displays,
             global_workflow_input_displays=global_workflow_input_displays,
+            state_value_displays=state_value_displays,
+            global_state_value_displays=global_state_value_displays,
             node_displays=node_displays,
             global_node_output_displays=global_node_output_displays,
             global_node_displays=global_node_displays,
@@ -306,6 +329,12 @@ class BaseWorkflowDisplay(
     ) -> WorkflowInputsDisplayType:
         pass
 
+    @abstractmethod
+    def _generate_state_value_display(
+        self, state_value: StateValueReference, overrides: Optional[StateValueDisplayOverridesType] = None
+    ) -> StateValueDisplayType:
+        pass
+
     @abstractmethod
     def _generate_entrypoint_display(
         self,

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -25,6 +25,8 @@ from vellum_ee.workflows.display.vellum import (
     EntrypointVellumDisplay,
     EntrypointVellumDisplayOverrides,
     NodeDisplayData,
+    StateValueVellumDisplay,
+    StateValueVellumDisplayOverrides,
     WorkflowInputsVellumDisplay,
     WorkflowInputsVellumDisplayOverrides,
     WorkflowMetaVellumDisplay,
@@ -44,6 +46,8 @@ class VellumWorkflowDisplay(
         WorkflowMetaVellumDisplayOverrides,
         WorkflowInputsVellumDisplay,
         WorkflowInputsVellumDisplayOverrides,
+        StateValueVellumDisplay,
+        StateValueVellumDisplayOverrides,
         BaseNodeDisplay,
         EntrypointVellumDisplay,
         EntrypointVellumDisplayOverrides,
@@ -76,6 +80,25 @@ class VellumWorkflowDisplay(
                 }
             )
 
+        state_variables: JsonArray = []
+        for state_value, state_value_display in self.display_context.state_value_displays.items():
+            default = primitive_to_vellum_value(state_value.instance) if state_value.instance else None
+            required = (
+                state_value_display.required
+                if state_value_display.required is not None
+                else type(None) not in state_value.types
+            )
+            state_variables.append(
+                {
+                    "id": str(state_value_display.id),
+                    "key": state_value_display.name or state_value.name,
+                    "type": infer_vellum_variable_type(state_value),
+                    "default": default.dict() if default else None,
+                    "required": required,
+                    "extensions": {"color": state_value_display.color},
+                }
+            )
+
         nodes: JsonArray = []
         edges: JsonArray = []
 
@@ -244,6 +267,7 @@ class VellumWorkflowDisplay(
                 },
             },
             "input_variables": input_variables,
+            "state_variables": state_variables,
             "output_variables": output_variables,
         }
 
@@ -283,6 +307,23 @@ class VellumWorkflowDisplay(
 
         return WorkflowInputsVellumDisplay(id=workflow_input_id, name=name, required=required, color=color)
 
+    def _generate_state_value_display(
+        self, state_value: BaseDescriptor, overrides: Optional[StateValueVellumDisplayOverrides] = None
+    ) -> StateValueVellumDisplay:
+        state_value_id: UUID
+        name = None
+        required = None
+        color = None
+        if overrides:
+            state_value_id = overrides.id
+            name = overrides.name
+            required = overrides.required
+            color = overrides.color
+        else:
+            state_value_id = uuid4_from_hash(f"{self.workflow_id}|state_values|id|{state_value.name}")
+
+        return StateValueVellumDisplay(id=state_value_id, name=name, required=required, color=color)
+
     def _generate_entrypoint_display(
         self,
         entrypoint: Type[BaseNode],