vellum-ai 0.13.0__tar.gz → 0.13.2__tar.gz

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/PKG-INFO

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

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_cli/__init__.py

@@ -4,6 +4,7 @@ import click
 
 from vellum_cli.aliased_group import ClickAliasedGroup
 from vellum_cli.image_push import image_push_command
+from vellum_cli.ping import ping_command
 from vellum_cli.pull import pull_command
 from vellum_cli.push import push_command
 
@@ -14,6 +15,15 @@ def main() -> None:
     pass
 
 
+@main.command
+def ping() -> None:
+    """
+    Ping Vellum to confirm that requests are correctly authenticated and to return information about the active
+    Workspace/Organization
+    """
+    ping_command()
+
+
 class PushGroup(ClickAliasedGroup):
     def get_command(self, ctx, cmd_name):
         # First try to get the command normally
@@ -54,6 +64,11 @@ def workflows():
     is_flag=True,
     help="Check the Workflow for errors and expected changes, without updating its state in Vellum.",
 )
+@click.option(
+    "--strict",
+    is_flag=True,
+    help="Raises an error if we detect an unexpected discrepancy in the generated artifact.",
+)
 def workflows_push(
     module: Optional[str],
     deploy: Optional[bool],
@@ -62,6 +77,7 @@ def workflows_push(
     deployment_description: Optional[str],
     release_tag: Optional[List[str]],
     dry_run: Optional[bool],
+    strict: Optional[bool],
 ) -> None:
     """
     Push Workflows to Vellum. If a module is provided, only the Workflow for that module will be pushed.
@@ -76,6 +92,7 @@ def workflows_push(
         deployment_description=deployment_description,
         release_tags=release_tag,
         dry_run=dry_run,
+        strict=strict,
     )
 
 
@@ -91,6 +108,11 @@ def workflows_push(
     is_flag=True,
     help="Check the Workflow for errors and expected changes, without updating its state in Vellum.",
 )
+@click.option(
+    "--strict",
+    is_flag=True,
+    help="Raises an error if we detect an unexpected discrepancy in the generated artifact.",
+)
 def push_module(
     ctx: click.Context,
     deploy: Optional[bool],
@@ -99,6 +121,7 @@ def push_module(
     deployment_description: Optional[str],
     release_tag: Optional[List[str]],
     dry_run: Optional[bool],
+    strict: Optional[bool],
 ) -> None:
     """Push a specific module to Vellum"""
 
@@ -111,6 +134,7 @@ def push_module(
             deployment_description=deployment_description,
             release_tags=release_tag,
             dry_run=dry_run,
+            strict=strict,
         )
 
 

vellum_ai-0.13.2/ee/vellum_cli/ping.py

@@ -0,0 +1,28 @@
+from dotenv import load_dotenv
+
+from vellum.workflows.vellum_client import create_vellum_client
+from vellum_cli.logger import load_cli_logger
+
+
+def ping_command():
+    load_dotenv()
+    logger = load_cli_logger()
+
+    client = create_vellum_client()
+
+    workspace = client.workspaces.workspace_identity()
+    organization = client.organizations.organization_identity()
+
+    logger.info(
+        f"""\
+Successfully authenticated with Vellum!
+
+Organization:
+    ID: {organization.id}
+    Name: {organization.name}
+
+Workspace:
+    ID: {workspace.id}
+    Name: {workspace.name}
+"""
+    )

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_cli/push.py

@@ -9,6 +9,7 @@ from typing import List, Optional
 
 from dotenv import load_dotenv
 
+from vellum.client.core.api_error import ApiError
 from vellum.resources.workflows.client import OMIT
 from vellum.types import WorkflowPushDeploymentConfigRequest
 from vellum.workflows.utils.names import snake_to_title_case
@@ -28,6 +29,7 @@ def push_command(
     deployment_description: Optional[str] = None,
     release_tags: Optional[List[str]] = None,
     dry_run: Optional[bool] = None,
+    strict: Optional[bool] = None,
 ) -> None:
     load_dotenv()
     logger = load_cli_logger()
@@ -109,18 +111,66 @@ def push_command(
     artifact.seek(0)
     artifact.name = f"{workflow_config.module.replace('.', '__')}.tar.gz"
 
-    response = client.workflows.push(
-        # Remove this once we could serialize using the artifact in Vembda
-        # https://app.shortcut.com/vellum/story/5585
-        exec_config=json.dumps(exec_config),
-        label=label,
-        workflow_sandbox_id=workflow_config.workflow_sandbox_id,
-        artifact=artifact,
-        # We should check with fern if we could auto-serialize typed object fields for us
-        # https://app.shortcut.com/vellum/story/5568
-        deployment_config=deployment_config_serialized,  # type: ignore[arg-type]
-        dry_run=dry_run,
-    )
+    try:
+        response = client.workflows.push(
+            # Remove this once we could serialize using the artifact in Vembda
+            # https://app.shortcut.com/vellum/story/5585
+            exec_config=json.dumps(exec_config),
+            label=label,
+            workflow_sandbox_id=workflow_config.workflow_sandbox_id,
+            artifact=artifact,
+            # We should check with fern if we could auto-serialize typed object fields for us
+            # https://app.shortcut.com/vellum/story/5568
+            deployment_config=deployment_config_serialized,  # type: ignore[arg-type]
+            dry_run=dry_run,
+            strict=strict,
+        )
+    except ApiError as e:
+        if e.status_code != 400 or not isinstance(e.body, dict) or "diffs" not in e.body:
+            raise e
+
+        diffs: dict = e.body["diffs"]
+        generated_only = diffs.get("generated_only", [])
+        generated_only_str = (
+            "\n".join(
+                ["Files that were generated but not found in the original project:"]
+                + [f"- {file}" for file in generated_only]
+            )
+            if generated_only
+            else ""
+        )
+
+        original_only = diffs.get("original_only", [])
+        original_only_str = (
+            "\n".join(
+                ["Files that were found in the original project but not generated:"]
+                + [f"- {file}" for file in original_only]
+            )
+            if original_only
+            else ""
+        )
+
+        modified = diffs.get("modified", {})
+        modified_str = (
+            "\n\n".join(
+                ["Files that were different between the original project and the generated artifact:"]
+                + ["\n".join(line.strip() for line in lines) for lines in modified.values()]
+            )
+            if modified
+            else ""
+        )
+
+        reported_diffs = f"""\
+{e.body.get("detail")}
+
+{generated_only_str}
+
+{original_only_str}
+
+{modified_str}
+"""
+        logger.error(reported_diffs)
+        return
 
     if dry_run:
         error_messages = [str(e) for e in workflow_display.errors]

vellum_ai-0.13.2/ee/vellum_cli/tests/test_ping.py

@@ -0,0 +1,47 @@
+from datetime import datetime
+
+from click.testing import CliRunner
+
+from vellum.client.types.organization_read import OrganizationRead
+from vellum.client.types.workspace_read import WorkspaceRead
+from vellum_cli import main as cli_main
+
+
+def test_ping__happy_path(vellum_client):
+    # GIVEN a cli
+    runner = CliRunner()
+
+    # AND a valid response from both the workflow and organization API calls
+    vellum_client.workspaces.workspace_identity.return_value = WorkspaceRead(
+        id="1234567890",
+        name="Test Workspace",
+        label="Test Workspace",
+        created=datetime.now(),
+    )
+    vellum_client.organizations.organization_identity.return_value = OrganizationRead(
+        id="1234567890",
+        name="Test Organization",
+        new_member_join_behavior="AUTO_ACCEPT_FROM_SHARED_DOMAIN",
+        allow_staff_access=True,
+    )
+
+    # WHEN calling `vellum ping`
+    result = runner.invoke(cli_main, ["ping"])
+
+    # THEN it should return Status information about the user's workspace
+    assert result.exit_code == 0
+    assert (
+        result.output
+        == """\x1b[38;20m\
+Successfully authenticated with Vellum!
+
+Organization:
+    ID: 1234567890
+    Name: Test Organization
+
+Workspace:
+    ID: 1234567890
+    Name: Test Workspace
+\x1b[0m
+"""
+    )

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_cli/tests/test_push.py

@@ -7,6 +7,7 @@ from uuid import uuid4
 
 from click.testing import CliRunner
 
+from vellum.client.core.api_error import ApiError
 from vellum.client.types.workflow_push_response import WorkflowPushResponse
 from vellum.utils.uuid import is_valid_uuid
 from vellum_cli import main as cli_main
@@ -242,3 +243,78 @@ class ExampleWorkflow(BaseWorkflow):
     assert "Serialization is not supported." in result.output
     assert "## Proposed Diffs" in result.output
     assert "iterable_item_added" in result.output
+
+
+def test_push__strict_option_returns_diffs(mock_module, vellum_client):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    base_dir = os.path.join(temp_dir, *module.split("."))
+    os.makedirs(base_dir, exist_ok=True)
+    workflow_py_file_content = """\
+from vellum.workflows import BaseWorkflow
+
+class ExampleWorkflow(BaseWorkflow):
+    pass
+"""
+    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
+        f.write(workflow_py_file_content)
+
+    # AND the push API call returns a 4xx response with diffs
+    vellum_client.workflows.push.side_effect = ApiError(
+        status_code=400,
+        body={
+            "detail": "Failed to push workflow due to unexpected detected differences in the generated artifact.",
+            "diffs": {
+                "generated_only": ["state.py"],
+                "modified": {
+                    "workflow.py": [
+                        "--- a/workflow.py\n",
+                        "+++ b/workflow.py\n",
+                        "@@ -1 +1 @@\n",
+                        "-print('hello')",
+                        "+print('foo')",
+                    ]
+                },
+                "original_only": ["inputs.py"],
+            },
+        },
+    )
+
+    # WHEN calling `vellum push` on strict mode
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module, "--strict"])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # AND we should have called the push API with the strict option
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert call_args["strict"] is True
+
+    # AND the report should be in the output
+    assert (
+        result.output
+        == """\
+\x1b[38;20mLoading workflow from examples.mock.test_push__strict_option_returns_diffs\x1b[0m
+\x1b[31;20mFailed to push workflow due to unexpected detected differences in the generated artifact.
+
+Files that were generated but not found in the original project:
+- state.py
+
+Files that were found in the original project but not generated:
+- inputs.py
+
+Files that were different between the original project and the generated artifact:
+
+--- a/workflow.py
++++ b/workflow.py
+@@ -1 +1 @@
+-print('hello')
++print('foo')
+\x1b[0m
+"""
+    )

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_ee/workflows/display/nodes/vellum/base_node.py

@@ -1,12 +1,11 @@
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, TypeVar, cast
 
+from vellum.workflows.constants import UNDEF
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.expressions.and_ import AndExpression
 from vellum.workflows.expressions.between import BetweenExpression
 from vellum.workflows.expressions.is_not_null import IsNotNullExpression
 from vellum.workflows.expressions.is_null import IsNullExpression
 from vellum.workflows.expressions.not_between import NotBetweenExpression
-from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.output import OutputReference
@@ -14,6 +13,7 @@ from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.references.workflow_input import WorkflowInputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 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.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import convert_descriptor_to_operator
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -28,6 +28,18 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
         node = self._node
         node_id = self.node_id
 
+        attributes: JsonArray = []
+        for attribute in node:
+            id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+
+            attributes.append(
+                {
+                    "id": id,
+                    "name": attribute.name,
+                    "value": self.serialize_value(display_context, cast(BaseDescriptor, attribute.instance)),
+                }
+            )
+
         ports: JsonArray = []
         for idx, port in enumerate(node.Ports):
             id = str(uuid4_from_hash(f"{node_id}|{idx}"))
@@ -36,6 +48,7 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
                 ports.append(
                     {
                         "id": id,
+                        "name": port.name,
                         "type": port._condition_type.value,
                         "expression": (
                             self.serialize_condition(display_context, port._condition) if port._condition else None
@@ -46,10 +59,29 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
                 ports.append(
                     {
                         "id": id,
+                        "name": port.name,
                         "type": "DEFAULT",
                     }
                 )
 
+        outputs: JsonArray = []
+        for output in node.Outputs:
+            type = primitive_type_to_vellum_variable_type(output)
+            value = (
+                self.serialize_value(display_context, output.instance)
+                if output.instance is not None and output.instance != UNDEF
+                else None
+            )
+
+            outputs.append(
+                {
+                    "id": str(uuid4_from_hash(f"{node_id}|{output.name}")),
+                    "name": output.name,
+                    "type": type,
+                    "value": value,
+                }
+            )
+
         return {
             "id": str(node_id),
             "label": node.__qualname__,
@@ -63,7 +95,8 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
             },
             "ports": ports,
             "adornments": None,
-            "attributes": [],
+            "attributes": attributes,
+            "outputs": outputs,
         }
 
     def get_generic_node_display_data(self) -> GenericNodeDisplayData:
@@ -71,12 +104,25 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
         return explicit_value if explicit_value else GenericNodeDisplayData()
 
     def serialize_condition(self, display_context: WorkflowDisplayContext, condition: BaseDescriptor) -> JsonObject:
-        if isinstance(condition, (AndExpression, OrExpression)):
-            return {}
-        elif isinstance(condition, (IsNullExpression, IsNotNullExpression)):
-            return {}
+        if isinstance(condition, (IsNullExpression, IsNotNullExpression)):
+            lhs = self.serialize_value(display_context, condition._expression)  # type: ignore[attr-defined]
+            return {
+                "type": "UNARY_EXPRESSION",
+                "lhs": lhs,
+                "operator": convert_descriptor_to_operator(condition),
+            }
         elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
-            return {}
+            base = self.serialize_value(display_context, condition._value)  # type: ignore[attr-defined]
+            lhs = self.serialize_value(display_context, condition._start)  # type: ignore[attr-defined]
+            rhs = self.serialize_value(display_context, condition._end)  # type: ignore[attr-defined]
+
+            return {
+                "type": "TERNARY_EXPRESSION",
+                "base": base,
+                "operator": convert_descriptor_to_operator(condition),
+                "lhs": lhs,
+                "rhs": rhs,
+            }
         else:
             lhs = self.serialize_value(display_context, condition._lhs)  # type: ignore[attr-defined]
             rhs = self.serialize_value(display_context, condition._rhs)  # type: ignore[attr-defined]
@@ -113,7 +159,7 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
             }
 
         if isinstance(value, ExecutionCountReference):
-            node_class_display = display_context.node_displays[value.node_class]
+            node_class_display = display_context.global_node_displays[value.node_class]
 
             return {
                 "type": "EXECUTION_COUNTER",
@@ -127,4 +173,6 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
                 "value": vellum_value.dict(),
             }
 
-        raise ValueError(f"Unsupported descriptor type: {value.__class__.__name__}")
+        # If it's not any of the references we know about,
+        # then try to serialize it as a nested value
+        return self.serialize_condition(display_context, value)

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -74,6 +74,9 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
         node_inputs: List[NodeInput] = []
         prompt_inputs: List[VellumVariable] = []
 
+        if not value:
+            return node_inputs, prompt_inputs
+
         for variable_name, variable_value in value.items():
             node_input = create_node_input(
                 node_id=node_id,

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_ee/workflows/display/nodes/vellum/map_node.py

@@ -69,7 +69,7 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
                 "workflow_raw_data": serialized_subworkflow["workflow_raw_data"],
                 "input_variables": cast(JsonObject, renamed_input_variables),
                 "output_variables": serialized_subworkflow["output_variables"],
-                "concurrency": raise_if_descriptor(node.concurrency),
+                "concurrency": raise_if_descriptor(node.max_concurrency),
                 "items_input_id": items_workflow_input_id,
                 "item_input_id": item_workflow_input_id,
                 "index_input_id": index_workflow_input_id,

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py

@@ -27,16 +27,20 @@ class BasePromptDeploymentNodeDisplay(
         node_id = self.node_id
 
         prompt_inputs = raise_if_descriptor(node.prompt_inputs)
-        node_inputs = [
-            create_node_input(
-                node_id=node_id,
-                input_name=variable_name,
-                value=variable_value,
-                display_context=display_context,
-                input_id=self.prompt_input_ids_by_name.get(variable_name),
-            )
-            for variable_name, variable_value in prompt_inputs.items()
-        ]
+        node_inputs = (
+            [
+                create_node_input(
+                    node_id=node_id,
+                    input_name=variable_name,
+                    value=variable_value,
+                    display_context=display_context,
+                    input_id=self.prompt_input_ids_by_name.get(variable_name),
+                )
+                for variable_name, variable_value in prompt_inputs.items()
+            ]
+            if prompt_inputs
+            else []
+        )
 
         _, output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
         _, array_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.results)]

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -2,6 +2,7 @@ import pytest
 from uuid import UUID, uuid4
 from typing import List, cast
 
+from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
@@ -19,7 +20,6 @@ from vellum_ee.workflows.display.vellum import (
     NodeInputValuePointerRule,
     NodeOutputData,
     NodeOutputPointer,
-    StringVellumValue,
     WorkflowInputsVellumDisplayOverrides,
     WorkflowMetaVellumDisplay,
 )
@@ -74,7 +74,7 @@ class MyNodeB(BaseNode):
                     type="INPUT_VARIABLE",
                     data=InputVariableData(input_variable_id="a154c29d-fac0-4cd0-ba88-bc52034f5470"),
                 ),
-                ConstantValuePointer(type="CONSTANT_VALUE", data=StringVellumValue(type="STRING", value="fallback")),
+                ConstantValuePointer(type="CONSTANT_VALUE", data=StringVellumValue(value="fallback")),
             ],
         ),
     ],

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_ee/workflows/display/nodes/vellum/utils.py

@@ -2,6 +2,7 @@ from uuid import UUID
 from typing import Any, List, Optional, Type, Union, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.expressions.and_ import AndExpression
 from vellum.workflows.expressions.begins_with import BeginsWithExpression
 from vellum.workflows.expressions.between import BetweenExpression
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
@@ -25,11 +26,13 @@ from vellum.workflows.expressions.less_than import LessThanExpression
 from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
 from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.expressions.not_in import NotInExpression
+from vellum.workflows.expressions.or_ import OrExpression
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
 from vellum.workflows.references import NodeReference, OutputReference
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule, primitive_to_vellum_value
+from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule
 from vellum_ee.workflows.display.vellum import (
     ConstantValuePointer,
     ExecutionCounterData,
@@ -170,5 +173,9 @@ def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> str:
         return "between"
     elif isinstance(descriptor, NotBetweenExpression):
         return "notBetween"
+    elif isinstance(descriptor, AndExpression):
+        return "and"
+    elif isinstance(descriptor, OrExpression):
+        return "or"
     else:
         raise ValueError(f"Unsupported descriptor type: {descriptor}")

{vellum_ai-0.13.0 → vellum_ai-0.13.2}/ee/vellum_ee/workflows/display/tests/test_vellum_workflow_display.py

@@ -1,4 +1,10 @@
+from uuid import UUID
+
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes import BaseNode
+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.workflows.get_vellum_workflow_display_class import get_workflow_display
 from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
@@ -40,3 +46,45 @@ def test_vellum_workflow_display__serialize_empty_workflow():
             ],
         },
     }
+
+
+def test_vellum_workflow_display__serialize_input_variables_with_capitalized_variable_override():
+    # GIVEN a workflow with input variables
+    class Inputs(BaseInputs):
+        foo: str
+
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            output = Inputs.foo
+
+    class ExampleWorkflow(BaseWorkflow[Inputs, BaseState]):
+        graph = StartNode
+
+    class ExampleWorkflowDisplay(VellumWorkflowDisplay[ExampleWorkflow]):
+        inputs_display = {
+            Inputs.foo: WorkflowInputsVellumDisplayOverrides(
+                id=UUID("97b63d71-5413-417f-9cf5-49e1b4fd56e4"), name="Foo", required=True
+            )
+        }
+
+    display = get_workflow_display(
+        base_display_class=ExampleWorkflowDisplay,
+        workflow_class=ExampleWorkflow,
+    )
+
+    # WHEN serializing the workflow
+    exec_config = display.serialize()
+
+    # THEN the input variables are what we expect
+    input_variables = exec_config["input_variables"]
+
+    assert input_variables == [
+        {
+            "id": "97b63d71-5413-417f-9cf5-49e1b4fd56e4",
+            "key": "Foo",
+            "type": "STRING",
+            "default": None,
+            "required": True,
+            "extensions": {"color": None},
+        }
+    ]