vellum-ai 0.14.89__py3-none-any.whl → 1.0.1__py3-none-any.whl

vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -11,7 +11,6 @@ from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.expressions import serialize_value
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
 
@@ -173,7 +172,6 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
             }
 
         elif prompt_block.block_type == "CHAT_MESSAGE":
-
             chat_properties: JsonObject = {
                 "chat_role": prompt_block.chat_role,
                 "chat_source": prompt_block.chat_source,
@@ -248,30 +246,6 @@ class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generi
 
         return block
 
-    def _serialize_attributes(self, display_context: "WorkflowDisplayContext"):
-        attributes = []
-        for attribute in self._node:
-            if attribute in self.__unserializable_attributes__:
-                continue
-
-            id = (
-                str(self.attribute_ids_by_name[attribute.name])
-                if self.attribute_ids_by_name.get(attribute.name)
-                else str(uuid4_from_hash(f"{self.node_id}|{attribute.name}"))
-            )
-            try:
-                attributes.append(
-                    {
-                        "id": id,
-                        "name": attribute.name,
-                        "value": serialize_value(display_context, attribute.instance),
-                    }
-                )
-            except ValueError as e:
-                raise ValueError(f"Failed to serialize attribute '{attribute.name}': {e}")
-
-        return attributes
-
     def _serialize_parameters(self, parameters, display_context: "WorkflowDisplayContext") -> JsonObject:
         """Serialize parameters, returning empty object when nested descriptors are detected."""
         params = raise_if_descriptor(parameters)

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

@@ -0,0 +1,31 @@
+from vellum.workflows.constants import APIRequestMethod
+from vellum.workflows.nodes.displayable.api_node.node import APINode
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_serialize_node__api_node_with_timeout():
+    # GIVEN an API node with timeout specified
+    class MyAPINode(APINode):
+        url = "https://api.example.com"
+        method = APIRequestMethod.GET
+        timeout = 30  # This is the key attribute we're testing
+
+    # AND a workflow with the API node
+    class Workflow(BaseWorkflow):
+        graph = MyAPINode
+
+    # 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 timeout
+    my_api_node = next(node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] == "API")
+
+    # Assert that timeout is present in the serialized attributes
+    timeout_attribute = next((attr for attr in my_api_node.get("attributes", []) if attr["name"] == "timeout"), None)
+
+    assert timeout_attribute is not None, "timeout attribute should be present in serialized attributes"
+    assert timeout_attribute["value"]["type"] == "CONSTANT_VALUE"
+    assert timeout_attribute["value"]["value"]["type"] == "NUMBER"
+    assert timeout_attribute["value"]["value"]["value"] == 30.0

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

@@ -8,6 +8,7 @@ from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.references import LazyReference
+from vellum_ee.workflows.display.base import WorkflowInputsDisplay, WorkflowMetaDisplay
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
@@ -21,7 +22,6 @@ from vellum_ee.workflows.display.utils.vellum import (
     NodeOutputData,
     NodeOutputPointer,
 )
-from vellum_ee.workflows.display.vellum import WorkflowInputsVellumDisplayOverrides, WorkflowMetaVellumDisplay
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
@@ -110,13 +110,13 @@ def test_create_node_input_value_pointer_rules(
         descriptor,
         WorkflowDisplayContext(
             workflow_display_class=BaseWorkflowDisplay,
-            workflow_display=WorkflowMetaVellumDisplay(
+            workflow_display=WorkflowMetaDisplay(
                 entrypoint_node_id=uuid4(),
                 entrypoint_node_source_handle_id=uuid4(),
                 entrypoint_node_display=NodeDisplayData(),
             ),
             global_workflow_input_displays={
-                Inputs.example_workflow_input: WorkflowInputsVellumDisplayOverrides(
+                Inputs.example_workflow_input: WorkflowInputsDisplay(
                     id=UUID("a154c29d-fac0-4cd0-ba88-bc52034f5470"),
                 ),
             },

vellum_ee/workflows/display/tests/test_base_workflow_display.py

@@ -7,7 +7,7 @@ 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
+from vellum_ee.workflows.display.base import WorkflowInputsDisplay
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -64,9 +64,7 @@ def test_vellum_workflow_display__serialize_input_variables_with_capitalized_var
 
     class ExampleWorkflowDisplay(BaseWorkflowDisplay[ExampleWorkflow]):
         inputs_display = {
-            Inputs.foo: WorkflowInputsVellumDisplayOverrides(
-                id=UUID("97b63d71-5413-417f-9cf5-49e1b4fd56e4"), name="Foo", required=True
-            )
+            Inputs.foo: WorkflowInputsDisplay(id=UUID("97b63d71-5413-417f-9cf5-49e1b4fd56e4"), name="Foo")
         }
 
     display = get_workflow_display(

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

@@ -4,6 +4,7 @@ from typing import Type
 
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
+from vellum_ee.workflows.display.base import WorkflowMetaDisplay
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.types import (
@@ -13,7 +14,6 @@ from vellum_ee.workflows.display.types import (
     WorkflowDisplayContext,
     WorkflowInputsDisplays,
 )
-from vellum_ee.workflows.display.vellum import WorkflowMetaVellumDisplay
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
@@ -31,7 +31,7 @@ def serialize_node():
 
         context: WorkflowDisplayContext = WorkflowDisplayContext(
             workflow_display_class=BaseWorkflowDisplay,
-            workflow_display=WorkflowMetaVellumDisplay(
+            workflow_display=WorkflowMetaDisplay(
                 entrypoint_node_id=uuid4(),
                 entrypoint_node_source_handle_id=uuid4(),
                 entrypoint_node_display=NodeDisplayData(),

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

@@ -162,6 +162,16 @@ def test_serialize_workflow(vellum_client):
                     },
                 },
             ],
+            "attributes": [
+                {
+                    "id": "ad719e65-0032-4012-a0bd-9b5162194bce",
+                    "name": "timeout",
+                    "value": {
+                        "type": "CONSTANT_VALUE",
+                        "value": {"type": "JSON", "value": None},
+                    },
+                },
+            ],
             "data": {
                 "label": "Simple A P I Node",
                 "error_output_id": None,

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

@@ -285,7 +285,7 @@ def test_serialize_workflow():
                     {
                         "id": "de7b0b4e-7803-4d36-a275-2e7e3f60342b",
                         "type": "IF",
-                        "source_handle_id": "561b4e3a-8db3-448a-8933-1115937082ff",
+                        "source_handle_id": "3a45b81f-95e4-4cbd-8997-bfdbe30251e8",
                         "data": {
                             "id": "2ccd0730-26d1-4fb4-baa9-1a2a182dd9a0",
                             "rules": [
@@ -309,7 +309,7 @@ def test_serialize_workflow():
                     {
                         "id": "5e783d17-6808-441a-ac6c-33a4e184f4e0",
                         "type": "ELIF",
-                        "source_handle_id": "0644f22e-2680-441a-9554-eedf1d3d22a9",
+                        "source_handle_id": "7202f702-1ebc-4067-ab1e-ec67e49158ee",
                         "data": {
                             "id": "cc3f0d92-b603-42cc-b2e9-83e3b23b3bcb",
                             "rules": [
@@ -333,7 +333,7 @@ def test_serialize_workflow():
                     {
                         "id": "6bd2f643-9cf5-4e7f-9113-f90e5c8057be",
                         "type": "ELIF",
-                        "source_handle_id": "2a48d274-ecfc-4f40-95ac-bc697663f10c",
+                        "source_handle_id": "cf45705d-1a47-43a6-9d24-a7fdf78baae0",
                         "data": {
                             "id": "a5a0f391-7052-452f-9fe1-a5781a491591",
                             "rules": [
@@ -357,7 +357,7 @@ def test_serialize_workflow():
                     {
                         "id": "0a058485-18a4-4e20-8a30-6da8196ac46f",
                         "type": "ELIF",
-                        "source_handle_id": "90fc9dc8-0a74-4a98-b6ac-55ffce4a2881",
+                        "source_handle_id": "f04610dd-61cf-41b0-b337-2235e101cdb0",
                         "data": {
                             "id": "efe7a851-2a67-4189-99ec-bc193242b270",
                             "rules": [
@@ -409,7 +409,7 @@ def test_serialize_workflow():
                     {
                         "id": "c2fa8a44-923b-462a-b0d2-fa800a152e52",
                         "type": "ELSE",
-                        "source_handle_id": "493024f4-8010-4e1a-abae-b6adbc6fb208",
+                        "source_handle_id": "f9dde637-ea90-465f-a871-caf8380ae377",
                         "data": None,
                     },
                 ],
@@ -913,7 +913,7 @@ def test_conditional_node_serialize_all_operators_with_lhs_and_rhs(descriptor, o
                     {
                         "id": "a4c32611-fd58-4b98-9d08-313cfd1c214e",
                         "type": "IF",
-                        "source_handle_id": "8124a6cf-4a34-4149-adc0-68696c11bd4e",
+                        "source_handle_id": "2ff87aa6-37cf-43dd-af9d-13b9198ab70a",
                         "data": {
                             "id": "650e7105-3e76-43ca-858f-b290970b438b",
                             "rules": [
@@ -937,7 +937,7 @@ def test_conditional_node_serialize_all_operators_with_lhs_and_rhs(descriptor, o
                     {
                         "id": "342e5497-ea2b-4e5c-99cf-e6492f133a3c",
                         "type": "ELSE",
-                        "source_handle_id": "4df924c0-7bed-4f4a-9db4-2bfe51841755",
+                        "source_handle_id": "9a7e8c2e-0228-4321-8f74-61cb5778f3df",
                         "data": None,
                     },
                 ],
@@ -1024,7 +1024,7 @@ def test_conditional_node_serialize_all_operators_with_expression(descriptor, op
                     {
                         "id": "a4c32611-fd58-4b98-9d08-313cfd1c214e",
                         "type": "IF",
-                        "source_handle_id": "8124a6cf-4a34-4149-adc0-68696c11bd4e",
+                        "source_handle_id": "2ff87aa6-37cf-43dd-af9d-13b9198ab70a",
                         "data": {
                             "id": "650e7105-3e76-43ca-858f-b290970b438b",
                             "rules": [
@@ -1048,7 +1048,7 @@ def test_conditional_node_serialize_all_operators_with_expression(descriptor, op
                     {
                         "id": "342e5497-ea2b-4e5c-99cf-e6492f133a3c",
                         "type": "ELSE",
-                        "source_handle_id": "4df924c0-7bed-4f4a-9db4-2bfe51841755",
+                        "source_handle_id": "9a7e8c2e-0228-4321-8f74-61cb5778f3df",
                         "data": None,
                     },
                 ],
@@ -1147,7 +1147,7 @@ def test_conditional_node_serialize_all_operators_with_value_and_start_and_end(d
                     {
                         "id": "a4c32611-fd58-4b98-9d08-313cfd1c214e",
                         "type": "IF",
-                        "source_handle_id": "8124a6cf-4a34-4149-adc0-68696c11bd4e",
+                        "source_handle_id": "2ff87aa6-37cf-43dd-af9d-13b9198ab70a",
                         "data": {
                             "id": "650e7105-3e76-43ca-858f-b290970b438b",
                             "rules": [
@@ -1171,7 +1171,7 @@ def test_conditional_node_serialize_all_operators_with_value_and_start_and_end(d
                     {
                         "id": "342e5497-ea2b-4e5c-99cf-e6492f133a3c",
                         "type": "ELSE",
-                        "source_handle_id": "4df924c0-7bed-4f4a-9db4-2bfe51841755",
+                        "source_handle_id": "9a7e8c2e-0228-4321-8f74-61cb5778f3df",
                         "data": None,
                     },
                 ],

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

@@ -407,6 +407,27 @@ def test_serialize_workflow():
                     ],
                 },
             },
+            {
+                "id": "229cd1ca-dc2f-4586-b933-c4d4966f7bd1",
+                "name": "parameters",
+                "value": {
+                    "type": "CONSTANT_VALUE",
+                    "value": {
+                        "type": "JSON",
+                        "value": {
+                            "stop": [],
+                            "temperature": 0.0,
+                            "max_tokens": 4096.0,
+                            "top_p": 1.0,
+                            "top_k": 0.0,
+                            "frequency_penalty": 0.0,
+                            "presence_penalty": 0.0,
+                            "logit_bias": None,
+                            "custom_parameters": None,
+                        },
+                    },
+                },
+            },
             {
                 "id": "1668419e-a193-43a5-8a97-3394e89bf278",
                 "name": "max_prompt_iterations",

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

@@ -169,6 +169,27 @@ def test_serialize_workflow():
                     ],
                 },
             },
+            {
+                "id": "229cd1ca-dc2f-4586-b933-c4d4966f7bd1",
+                "name": "parameters",
+                "value": {
+                    "type": "CONSTANT_VALUE",
+                    "value": {
+                        "type": "JSON",
+                        "value": {
+                            "stop": [],
+                            "temperature": 0.0,
+                            "max_tokens": 4096.0,
+                            "top_p": 1.0,
+                            "top_k": 0.0,
+                            "frequency_penalty": 0.0,
+                            "presence_penalty": 0.0,
+                            "logit_bias": None,
+                            "custom_parameters": None,
+                        },
+                    },
+                },
+            },
             {
                 "id": "1668419e-a193-43a5-8a97-3394e89bf278",
                 "name": "max_prompt_iterations",

vellum_ee/workflows/display/vellum.py

@@ -1,135 +1,9 @@
-from dataclasses import dataclass
-from uuid import UUID
-from typing import List, Literal, Optional
+from typing import List, Literal
 
-from vellum.core import UniversalBaseModel
-from vellum_ee.workflows.display.base import (
-    EdgeDisplayOverrides,
-    EntrypointDisplayOverrides,
-    StateValueDisplayOverrides,
-    WorkflowInputsDisplay,
-    WorkflowMetaDisplayOverrides,
-    WorkflowOutputDisplayOverrides,
-)
-from vellum_ee.workflows.display.base import WorkflowDisplayData  # noqa: F401 - Remove in 0.15.0
-from vellum_ee.workflows.display.base import WorkflowDisplayDataViewport  # noqa: F401 - Remove in 0.15.0
-from vellum_ee.workflows.display.editor.types import NodeDisplayComment  # noqa: F401 - Remove in 0.15.0
-from vellum_ee.workflows.display.editor.types import NodeDisplayData
-from vellum_ee.workflows.display.editor.types import NodeDisplayPosition  # noqa: F401 - Remove in 0.15.0
+from vellum.client.core import UniversalBaseModel
 from vellum_ee.workflows.display.utils.vellum import NodeInputValuePointerRule
 
 
-@dataclass
-class WorkflowMetaVellumDisplayOverrides(WorkflowMetaDisplayOverrides):
-    """
-    DEPRECATED: Use WorkflowMetaDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class WorkflowMetaVellumDisplay(WorkflowMetaVellumDisplayOverrides):
-    """
-    DEPRECATED: Use WorkflowMetaDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class WorkflowInputsVellumDisplayOverrides(WorkflowInputsDisplay):
-    """
-    DEPRECATED: Use WorkflowInputsDisplay instead. Will be removed in 0.15.0
-    """
-
-    required: Optional[bool] = None
-
-
-@dataclass
-class WorkflowInputsVellumDisplay(WorkflowInputsVellumDisplayOverrides):
-    pass
-
-
-@dataclass
-class StateValueVellumDisplayOverrides(StateValueDisplayOverrides):
-    """
-    DEPRECATED: Use StateValueDisplay instead. Will be removed in 0.15.0
-    """
-
-    required: Optional[bool] = None
-
-
-@dataclass
-class StateValueVellumDisplay(StateValueVellumDisplayOverrides):
-    """
-    DEPRECATED: Use StateValueDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class EdgeVellumDisplayOverrides(EdgeDisplayOverrides):
-    """
-    DEPRECATED: Use EdgeDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class EdgeVellumDisplay(EdgeVellumDisplayOverrides):
-    """
-    DEPRECATED: Use EdgeDisplay instead. Will be removed in 0.15.0
-    """
-
-    source_node_id: UUID
-    source_handle_id: UUID
-    target_node_id: UUID
-    target_handle_id: UUID
-    type: Literal["DEFAULT"] = "DEFAULT"
-
-
-@dataclass
-class EntrypointVellumDisplayOverrides(EntrypointDisplayOverrides):
-    """
-    DEPRECATED: Use EntrypointDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class EntrypointVellumDisplay(EntrypointVellumDisplayOverrides):
-    """
-    DEPRECATED: Use EntrypointDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
-@dataclass
-class WorkflowOutputVellumDisplayOverrides(WorkflowOutputDisplayOverrides):
-    """
-    DEPRECATED: Use WorkflowOutputDisplay instead. Will be removed in 0.15.0
-    """
-
-    label: Optional[str] = None
-    node_id: Optional[UUID] = None
-    display_data: Optional[NodeDisplayData] = None
-    target_handle_id: Optional[UUID] = None
-
-
-@dataclass
-class WorkflowOutputVellumDisplay(WorkflowOutputVellumDisplayOverrides):
-    """
-    DEPRECATED: Use WorkflowOutputDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
 class NodeInputValuePointer(UniversalBaseModel):
     rules: List[NodeInputValuePointerRule]
     combinator: Literal["OR"] = "OR"

vellum_ee/workflows/display/workflows/__init__.py

@@ -2,4 +2,3 @@
 
 # These must be imported to register the classes in the registry
 from .base_workflow_display import BaseWorkflowDisplay
-from .vellum_workflow_display import VellumWorkflowDisplay

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -9,7 +9,7 @@ from uuid import UUID
 from typing import Any, Dict, ForwardRef, Generic, Iterator, List, Optional, Tuple, Type, TypeVar, Union, cast, get_args
 
 from vellum.client import Vellum as VellumClient
-from vellum.core.pydantic_utilities import UniversalBaseModel
+from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -136,13 +136,19 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 if workflow_input_reference.instance
                 else None
             )
+
+            # An input is required if it has no default value AND is not optional
+            has_default = workflow_input_reference.instance is not undefined
+            is_optional = type(None) in workflow_input_reference.types
+            is_required = not has_default and not is_optional
+
             input_variables.append(
                 {
                     "id": str(workflow_input_display.id),
                     "key": workflow_input_display.name or workflow_input_reference.name,
                     "type": infer_vellum_variable_type(workflow_input_reference),
                     "default": default.dict() if default else None,
-                    "required": workflow_input_reference.instance is undefined,
+                    "required": is_required,
                     "extensions": {"color": workflow_input_display.color},
                 }
             )
@@ -152,13 +158,19 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             default = (
                 primitive_to_vellum_value(state_value_reference.instance) if state_value_reference.instance else None
             )
+
+            # A state variable is required if it has no default value AND is not optional
+            has_default = state_value_reference.instance is not undefined
+            is_optional = type(None) in state_value_reference.types
+            is_required = not has_default and not is_optional
+
             state_variables.append(
                 {
                     "id": str(state_value_display.id),
                     "key": state_value_display.name or state_value_reference.name,
                     "type": infer_vellum_variable_type(state_value_reference),
                     "default": default.dict() if default else None,
-                    "required": state_value_reference.instance is undefined,
+                    "required": is_required,
                     "extensions": {"color": state_value_display.color},
                 }
             )

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

@@ -1,6 +1,8 @@
 import pytest
 from uuid import uuid4
+from typing import Optional
 
+from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.inline_subworkflow_node.node import InlineSubworkflowNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
@@ -8,6 +10,7 @@ from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.nodes.displayable.final_output_node.node import FinalOutputNode
 from vellum.workflows.references.lazy import LazyReference
+from vellum.workflows.state.base import BaseState
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.nodes import BaseNodeDisplay
@@ -828,3 +831,127 @@ def test_serialize_workflow__empty_rules_indexerror():
     terminal_nodes = [node for node in nodes if node.get("type") == "TERMINAL"]
     assert len(terminal_nodes) == 1
     assert terminal_nodes[0]["data"]["name"] == "problematic_output"
+
+
+def test_serialize_workflow__input_variables():
+    # GIVEN a workflow with inputs
+    class Inputs(BaseInputs):
+        input_1: str
+        input_2: Optional[str]
+        input_3: int = 1
+        input_4: Optional[int] = 2
+
+    class MyWorkflow(BaseWorkflow[Inputs, BaseState]):
+        pass
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN the inputs should be serialized correctly
+    assert "input_variables" in data
+    input_variables = data["input_variables"]
+    assert isinstance(input_variables, list)
+    assert len(input_variables) == 4
+
+    input_1 = next(var for var in input_variables if isinstance(var, dict) and var["key"] == "input_1")
+    assert input_1 == {
+        "id": "13bd7980-3fbd-486c-9ebd-a29d84f7bda0",
+        "key": "input_1",
+        "type": "STRING",
+        "default": None,
+        "required": True,
+        "extensions": {"color": None},
+    }
+
+    input_2 = next(var for var in input_variables if isinstance(var, dict) and var["key"] == "input_2")
+    assert input_2 == {
+        "id": "13847952-beab-408d-945e-cfa079e6e124",
+        "key": "input_2",
+        "type": "STRING",
+        "default": None,
+        "required": False,
+        "extensions": {"color": None},
+    }
+
+    input_3 = next(var for var in input_variables if isinstance(var, dict) and var["key"] == "input_3")
+    assert input_3 == {
+        "id": "2e38e1a4-09ff-4bb8-a12e-9bf54d4f3a5e",
+        "key": "input_3",
+        "type": "NUMBER",
+        "default": {"type": "NUMBER", "value": 1.0},
+        "required": False,
+        "extensions": {"color": None},
+    }
+
+    input_4 = next(var for var in input_variables if isinstance(var, dict) and var["key"] == "input_4")
+    assert input_4 == {
+        "id": "d945b6ae-2490-4bfb-9b1c-b1e484dfd4f6",
+        "key": "input_4",
+        "type": "NUMBER",
+        "default": {"type": "NUMBER", "value": 2.0},
+        "required": False,
+        "extensions": {"color": None},
+    }
+
+
+def test_serialize_workflow__state_variables():
+    # GIVEN a workflow with state variables
+    class State(BaseState):
+        state_1: str = "hello"
+        state_2: Optional[str] = None
+        state_3: int = 1
+        state_4: Optional[int] = 2
+
+    class MyWorkflow(BaseWorkflow[BaseInputs, State]):
+        pass
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN the state variables should be serialized correctly
+    assert "state_variables" in data
+    state_variables = data["state_variables"]
+    assert isinstance(state_variables, list)
+    assert len(state_variables) == 4
+
+    state_1 = next(var for var in state_variables if isinstance(var, dict) and var["key"] == "state_1")
+    assert state_1 == {
+        "id": "83c5b71d-56eb-42a5-84df-97e3591370c2",
+        "key": "state_1",
+        "type": "STRING",
+        "default": {"type": "STRING", "value": "hello"},
+        "required": False,
+        "extensions": {"color": None},
+    }
+
+    state_2 = next(var for var in state_variables if isinstance(var, dict) and var["key"] == "state_2")
+    assert state_2 == {
+        "id": "9b0cfeec-aa66-42b3-8f31-aa7eb8ac30ea",
+        "key": "state_2",
+        "type": "STRING",
+        "default": None,
+        "required": False,
+        "extensions": {"color": None},
+    }
+
+    state_3 = next(var for var in state_variables if isinstance(var, dict) and var["key"] == "state_3")
+    assert state_3 == {
+        "id": "3e19c570-6b46-4eab-ad81-d8d97028496f",
+        "key": "state_3",
+        "type": "NUMBER",
+        "default": {"type": "NUMBER", "value": 1.0},
+        "required": False,
+        "extensions": {"color": None},
+    }
+
+    state_4 = next(var for var in state_variables if isinstance(var, dict) and var["key"] == "state_4")
+    assert state_4 == {
+        "id": "50c735de-f269-4d0a-b511-c9a1104451bb",
+        "key": "state_4",
+        "type": "NUMBER",
+        "default": {"type": "NUMBER", "value": 2.0},
+        "required": False,
+        "extensions": {"color": None},
+    }