vellum-ai 0.14.10__py3-none-any.whl → 0.14.12__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.14.10",
+            "X-Fern-SDK-Version": "0.14.12",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/workflows/events/types.py

@@ -75,6 +75,7 @@ class BaseDeploymentParentContext(BaseParentContext):
     release_tag_id: UUID
     release_tag_name: str
     external_id: Optional[str]
+    metadata: Optional[dict]
 
 
 class WorkflowDeploymentParentContext(BaseDeploymentParentContext):
@@ -104,6 +105,10 @@ class WorkflowSandboxParentContext(BaseParentContext):
     scenario_id: UUID
 
 
+class APIRequestParentContext(BaseParentContext):
+    type: Literal["API_REQUEST"] = "API_REQUEST"
+
+
 # Define the discriminated union
 ParentContext = Annotated[
     Union[
@@ -112,6 +117,7 @@ ParentContext = Annotated[
         WorkflowDeploymentParentContext,
         PromptDeploymentParentContext,
         WorkflowSandboxParentContext,
+        APIRequestParentContext,
     ],
     Field(discriminator="type"),
 ]

vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -458,16 +458,7 @@ def main(word: str) -> int:
         node.run()
 
     # THEN the node should have produced the exception we expected
-    assert (
-        exc_info.value.message
-        == """\
-2 validation errors for FunctionCall
-arguments
-  Field required [type=missing, input_value={'n': 'hello', 'a': {}}, input_type=dict]
-name
-  Field required [type=missing, input_value={'n': 'hello', 'a': {}}, input_type=dict]\
-"""
-    )
+    assert exc_info.value.message == "Expected an output of type 'FunctionCall', but received 'dict'"
 
 
 def test_run_node__run_inline__valid_dict_to_pydantic_any_type():
@@ -743,3 +734,46 @@ Node.js v21.7.3
 
     # AND the error should contain the execution error details
     assert exc_info.value.message == message
+
+
+def test_run_node__execute_code__list_extends():
+    # GIVEN a node that will return a list with output type Json
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Json]):
+        code = """\
+def main(left, right):
+    all = []
+    all.extend(left)
+    all.extend(right)
+    return all
+"""
+        code_inputs = {
+            "left": [1, 2, 3],
+            "right": [4, 5, 6],
+        }
+        runtime = "PYTHON_3_11_6"
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # AND the result should be the correct output
+    assert outputs == {"result": [1, 2, 3, 4, 5, 6], "log": ""}
+
+
+def test_run_node__execute_code__non_str_print():
+    # GIVEN a node that will print a non-string value
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = """\
+def main():
+    print(type(1))
+    return "hello"
+"""
+        code_inputs = {}
+        runtime = "PYTHON_3_11_6"
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # AND the result should be the correct output
+    assert outputs == {"result": "hello", "log": "<class 'int'>\n"}

vellum/workflows/nodes/displayable/code_execution_node/utils.py

@@ -1,12 +1,12 @@
 import io
 import os
-import re
-from typing import Any, Tuple, Union, get_args, get_origin
+from typing import Any, Tuple, Union
 
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel
 
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.utils import cast_to_output_type
 from vellum.workflows.types.core import EntityInputsInterface
 
 
@@ -67,49 +67,6 @@ def _clean_for_dict_wrapper(obj):
     return obj
 
 
-def _get_type_name(obj: Any) -> str:
-    if isinstance(obj, type):
-        return obj.__name__
-
-    if get_origin(obj) is Union:
-        children = [_get_type_name(child) for child in get_args(obj)]
-        return " | ".join(children)
-
-    return str(obj)
-
-
-def _cast_to_output_type(result: Any, output_type: Any) -> Any:
-    is_valid_output_type = isinstance(output_type, type)
-    if get_origin(output_type) is Union:
-        allowed_types = get_args(output_type)
-        for allowed_type in allowed_types:
-            try:
-                return _cast_to_output_type(result, allowed_type)
-            except NodeException:
-                continue
-    elif get_origin(output_type) is list:
-        allowed_item_type = get_args(output_type)[0]
-        if isinstance(result, list):
-            return [_cast_to_output_type(item, allowed_item_type) for item in result]
-    elif is_valid_output_type and issubclass(output_type, BaseModel) and not isinstance(result, output_type):
-        try:
-            return output_type.model_validate(result)
-        except ValidationError as e:
-            raise NodeException(
-                code=WorkflowErrorCode.INVALID_OUTPUTS,
-                message=re.sub(r"\s+For further information visit [^\s]+", "", str(e)),
-            ) from e
-    elif is_valid_output_type and isinstance(result, output_type):
-        return result
-
-    output_type_name = _get_type_name(output_type)
-    result_type_name = _get_type_name(type(result))
-    raise NodeException(
-        code=WorkflowErrorCode.INVALID_OUTPUTS,
-        message=f"Expected an output of type '{output_type_name}', but received '{result_type_name}'",
-    )
-
-
 def run_code_inline(
     code: str,
     inputs: EntityInputsInterface,
@@ -117,6 +74,11 @@ def run_code_inline(
 ) -> Tuple[str, Any]:
     log_buffer = io.StringIO()
 
+    def _inline_print(*args: Any, **kwargs: Any) -> None:
+        str_args = [str(arg) for arg in args]
+        print_line = f"{' '.join(str_args)}\n"
+        log_buffer.write(print_line)
+
     def wrap_value(value):
         if isinstance(value, list):
             return ListWrapper(
@@ -135,7 +97,7 @@ def run_code_inline(
     exec_globals = {
         "__arg__inputs": {name: wrap_value(value) for name, value in inputs.items()},
         "__arg__out": None,
-        "print": lambda *args, **kwargs: log_buffer.write(f"{' '.join(args)}\n"),
+        "print": _inline_print,
     }
     run_args = [f"{name}=__arg__inputs['{name}']" for name in inputs.keys()]
     execution_code = f"""\
@@ -154,7 +116,6 @@ __arg__out = main({", ".join(run_args)})
     logs = log_buffer.getvalue()
     result = exec_globals["__arg__out"]
 
-    if output_type != Any:
-        result = _cast_to_output_type(result, output_type)
+    result = cast_to_output_type(result, output_type)
 
     return logs, result

vellum/workflows/nodes/displayable/final_output_node/node.py

@@ -1,7 +1,9 @@
 from typing import Any, Dict, Generic, Tuple, Type, TypeVar, get_args
 
+from vellum.workflows.constants import undefined
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
+from vellum.workflows.nodes.utils import cast_to_output_type
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.types.utils import get_original_base
@@ -48,4 +50,14 @@ class FinalOutputNode(BaseNode[StateType], Generic[StateType, _OutputType], meta
     class Outputs(BaseNode.Outputs):
         # We use our mypy plugin to override the _OutputType with the actual output type
         # for downstream references to this output.
-        value: _OutputType  # type: ignore[valid-type]
+        value: _OutputType = undefined  # type: ignore[valid-type]
+
+    def run(self) -> Outputs:
+        original_outputs = self.Outputs()
+
+        return self.Outputs(
+            value=cast_to_output_type(
+                original_outputs.value,
+                self.__class__.get_output_type(),
+            )
+        )

vellum/workflows/nodes/displayable/final_output_node/tests/test_node.py

@@ -0,0 +1,20 @@
+import pytest
+
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.state.base import BaseState
+
+
+def test_final_output_node__mismatched_output_type():
+    # GIVEN a FinalOutputNode with a mismatched output type
+    class StringOutputNode(FinalOutputNode[BaseState, str]):
+        class Outputs(FinalOutputNode.Outputs):
+            value = {"foo": "bar"}
+
+    # WHEN the node is run
+    node = StringOutputNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN an error is raised
+    assert str(exc_info.value) == "Expected an output of type 'str', but received 'dict'"

vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py

@@ -70,8 +70,10 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                         value=input_value,
                     )
                 )
-            elif isinstance(input_value, list) and all(
-                isinstance(message, (ChatMessage, ChatMessageRequest)) for message in input_value
+            elif (
+                isinstance(input_value, list)
+                and len(input_value) > 0
+                and all(isinstance(message, (ChatMessage, ChatMessageRequest)) for message in input_value)
             ):
                 chat_history = [
                     (
@@ -95,7 +97,7 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                         value=cast(Dict[str, Any], input_value),
                     )
                 )
-            elif isinstance(input_value, float):
+            elif isinstance(input_value, (int, float)):
                 compiled_inputs.append(
                     WorkflowRequestNumberInputRequest(
                         name=input_name,
@@ -110,7 +112,6 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                         message=f"Failed to serialize input '{input_name}' of type '{input_value.__class__}': {e}",
                         code=WorkflowErrorCode.INVALID_INPUTS,
                     )
-
                 compiled_inputs.append(
                     WorkflowRequestJsonInputRequest(
                         name=input_name,

vellum/workflows/nodes/displayable/subworkflow_deployment_node/tests/test_node.py

@@ -10,6 +10,7 @@ from vellum.client.types.workflow_execution_workflow_result_event import Workflo
 from vellum.client.types.workflow_output_string import WorkflowOutputString
 from vellum.client.types.workflow_request_chat_history_input_request import WorkflowRequestChatHistoryInputRequest
 from vellum.client.types.workflow_request_json_input_request import WorkflowRequestJsonInputRequest
+from vellum.client.types.workflow_request_number_input_request import WorkflowRequestNumberInputRequest
 from vellum.client.types.workflow_result_event import WorkflowResultEvent
 from vellum.client.types.workflow_stream_event import WorkflowStreamEvent
 from vellum.workflows.errors import WorkflowErrorCode
@@ -134,6 +135,116 @@ def test_run_workflow__any_array(vellum_client):
     ]
 
 
+def test_run_workflow__empty_array(vellum_client):
+    # GIVEN a Subworkflow Deployment Node
+    class ExampleSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "fruits": [],
+        }
+
+    # AND we know what the Subworkflow Deployment will respond with
+    def generate_subworkflow_events(*args: Any, **kwargs: Any) -> Iterator[WorkflowStreamEvent]:
+        execution_id = str(uuid4())
+        expected_events: List[WorkflowStreamEvent] = [
+            WorkflowExecutionWorkflowResultEvent(
+                execution_id=execution_id,
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="INITIATED",
+                    ts=datetime.now(),
+                ),
+            ),
+            WorkflowExecutionWorkflowResultEvent(
+                execution_id=execution_id,
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="FULFILLED",
+                    ts=datetime.now(),
+                    outputs=[
+                        WorkflowOutputString(
+                            id=str(uuid4()),
+                            name="greeting",
+                            value="Great!",
+                        )
+                    ],
+                ),
+            ),
+        ]
+        yield from expected_events
+
+    vellum_client.execute_workflow_stream.side_effect = generate_subworkflow_events
+
+    # WHEN we run the node
+    node = ExampleSubworkflowDeploymentNode()
+    events = list(node.run())
+
+    # THEN the node should have completed successfully
+    assert events[-1].name == "greeting"
+    assert events[-1].value == "Great!"
+
+    # AND we should have invoked the Subworkflow Deployment with the expected inputs
+    call_kwargs = vellum_client.execute_workflow_stream.call_args.kwargs
+    assert call_kwargs["inputs"] == [
+        WorkflowRequestJsonInputRequest(name="fruits", value=[]),
+    ]
+
+
+def test_run_workflow__int_input(vellum_client):
+    # GIVEN a Subworkflow Deployment Node
+    class ExampleSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "number": 42,
+        }
+
+    # AND we know what the Subworkflow Deployment will respond with
+    def generate_subworkflow_events(*args: Any, **kwargs: Any) -> Iterator[WorkflowStreamEvent]:
+        execution_id = str(uuid4())
+        expected_events: List[WorkflowStreamEvent] = [
+            WorkflowExecutionWorkflowResultEvent(
+                execution_id=execution_id,
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="INITIATED",
+                    ts=datetime.now(),
+                ),
+            ),
+            WorkflowExecutionWorkflowResultEvent(
+                execution_id=execution_id,
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="FULFILLED",
+                    ts=datetime.now(),
+                    outputs=[
+                        WorkflowOutputString(
+                            id=str(uuid4()),
+                            name="greeting",
+                            value="Great!",
+                        )
+                    ],
+                ),
+            ),
+        ]
+        yield from expected_events
+
+    vellum_client.execute_workflow_stream.side_effect = generate_subworkflow_events
+
+    # WHEN we run the node
+    node = ExampleSubworkflowDeploymentNode()
+    events = list(node.run())
+
+    # THEN the node should have completed successfully
+    assert events[-1].name == "greeting"
+    assert events[-1].value == "Great!"
+
+    # AND we should have invoked the Subworkflow Deployment with the expected inputs
+    call_kwargs = vellum_client.execute_workflow_stream.call_args.kwargs
+    assert call_kwargs["inputs"] == [
+        WorkflowRequestNumberInputRequest(name="number", value=42),
+    ]
+
+
 def test_run_workflow__no_deployment():
     """Confirm that we raise error when running a subworkflow deployment node with no deployment attribute set"""
 

vellum/workflows/nodes/utils.py

@@ -2,10 +2,12 @@ from functools import cache
 import json
 import sys
 from types import ModuleType
-from typing import Any, Callable, Optional, Type, TypeVar, Union, get_args, get_origin
+from typing import Any, Callable, Dict, ForwardRef, List, Optional, Type, TypeVar, Union, get_args, get_origin
 
-from pydantic import BaseModel
+from pydantic import BaseModel, create_model
 
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes import BaseNode
 from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.ports.port import Port
@@ -141,3 +143,55 @@ def parse_type_from_str(result_as_str: str, output_type: Any) -> Any:
             raise ValueError("Invalid JSON format for result_as_str")
 
     raise ValueError(f"Unsupported output type: {output_type}")
+
+
+def _get_type_name(obj: Any) -> str:
+    if isinstance(obj, type):
+        return obj.__name__
+
+    if get_origin(obj) is Union:
+        children = [_get_type_name(child) for child in get_args(obj)]
+        return " | ".join(children)
+
+    return str(obj)
+
+
+def _clean_output_type(output_type: Any) -> Any:
+    """
+    pydantic currently has a bug where it doesn't support forward references in the create_model function. It will
+    fail due to a max recursion depth error. This negatively impacts our `Json` type, but could also apply to other
+    user defined ForwardRef types.
+    """
+    if get_origin(output_type) is Union:
+        clean_args = [_clean_output_type(child) for child in get_args(output_type)]
+        return Union[tuple(clean_args)]
+
+    if isinstance(output_type, ForwardRef):
+        # Here is where we prevent the max recursion depth error
+        return Any
+
+    if get_origin(output_type) is list:
+        clean_args = [_clean_output_type(child) for child in get_args(output_type)]
+        return List[clean_args[0]]  # type: ignore[valid-type]
+
+    if get_origin(output_type) is dict:
+        clean_args = [_clean_output_type(child) for child in get_args(output_type)]
+        return Dict[clean_args[0], clean_args[1]]  # type: ignore[valid-type]
+
+    return output_type
+
+
+def cast_to_output_type(result: Any, output_type: Any) -> Any:
+    clean_output_type = _clean_output_type(output_type)
+    DynamicModel = create_model("Output", output_type=(clean_output_type, ...))
+
+    try:
+        # mypy doesn't realize that this dynamic model has the output_type field defined above
+        return DynamicModel.model_validate({"output_type": result}).output_type  # type: ignore[attr-defined]
+    except Exception:
+        output_type_name = _get_type_name(output_type)
+        result_type_name = _get_type_name(type(result))
+        raise NodeException(
+            code=WorkflowErrorCode.INVALID_OUTPUTS,
+            message=f"Expected an output of type '{output_type_name}', but received '{result_type_name}'",
+        )

vellum/workflows/outputs/base.py

@@ -204,13 +204,11 @@ class BaseOutputs(metaclass=_BaseOutputsMeta):
         if not isinstance(other, dict):
             return super().__eq__(other)
 
-        outputs = {
-            name: value for name, value in vars(self).items() if not name.startswith("_") and value is not undefined
-        }
+        outputs = {ref.name: value for ref, value in self if value is not undefined}
         return outputs == other
 
     def __repr__(self) -> str:
-        values = f"{', '.join(f'{k}={v}' for k, v in vars(self).items() if not k.startswith('_'))}"
+        values = f"{', '.join(f'{ref.name}={value}' for ref, value in self if value is not undefined)}"
         return f"{self.__class__.__name__}({values})"
 
     def __iter__(self) -> Iterator[Tuple[OutputReference, Any]]:

vellum/workflows/workflows/base.py

@@ -24,7 +24,6 @@ from typing import (
     get_args,
 )
 
-from vellum.workflows.context import get_execution_context
 from vellum.workflows.edges import Edge
 from vellum.workflows.emitters.base import BaseWorkflowEmitter
 from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
@@ -481,7 +480,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         return self.get_inputs_class()()
 
     def get_default_state(self, workflow_inputs: Optional[InputsType] = None) -> StateType:
-        execution_context = get_execution_context()
+        execution_context = self._execution_context
         return self.get_state_class()(
             meta=(
                 StateMeta(

{vellum_ai-0.14.10.dist-info → vellum_ai-0.14.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.10
+Version: 0.14.12
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -32,7 +32,7 @@ Requires-Dist: pydantic (>=1.9.2)
 Requires-Dist: pydantic-core (>=2.18.2,<3.0.0)
 Requires-Dist: pydash (==7.0.6)
 Requires-Dist: python-dotenv (==1.0.1)
-Requires-Dist: pytz (==2022.6)
+Requires-Dist: pytz (==2025.1)
 Requires-Dist: pyyaml (==6.0.1)
 Requires-Dist: requests (==2.32.3)
 Requires-Dist: tomli (==2.0.2)

{vellum_ai-0.14.10.dist-info → vellum_ai-0.14.12.dist-info}/RECORD

@@ -3,7 +3,7 @@ vellum_cli/README.md,sha256=2NudRoLzWxNKqnuVy1JuQ7DerIaxWGYkrH8kMd-asIE,90
 vellum_cli/__init__.py,sha256=7aO9XFnaEVRiVshn86cFudebFUccT-gV8xIARJWqKYo,12257
 vellum_cli/aliased_group.py,sha256=ugW498j0yv4ALJ8vS9MsO7ctDW7Jlir9j6nE_uHAP8c,3363
 vellum_cli/config.py,sha256=aKnhvM5B8QdPA4cQC5Sqg7ImP-vNcVdSkZmk_OBpQTw,9309
-vellum_cli/image_push.py,sha256=SJwhwWJsLjwGNezNVd_oCVpFMfPsAB3dfLWmriZZUtw,4419
+vellum_cli/image_push.py,sha256=4auU15Pb6c8DTGvT-AQ5HHXXrvIvEDs6L02d4OvJYI8,5199
 vellum_cli/init.py,sha256=WpnMXPItPmh0f0bBGIer3p-e5gu8DUGwSArT_FuoMEw,5093
 vellum_cli/logger.py,sha256=PuRFa0WCh4sAGFS5aqWB0QIYpS6nBWwPJrIXpWxugV4,1022
 vellum_cli/ping.py,sha256=lWyJw6sziXjyTopTYRdFF5hV-sYPVDdX0yVbG5fzcY4,585
@@ -44,13 +44,14 @@ vellum_ee/workflows/display/nodes/vellum/map_node.py,sha256=8CPnn06HIBxBOiECevUf
 vellum_ee/workflows/display/nodes/vellum/merge_node.py,sha256=HkNMgdQELiON42jdO-xDLmqrEKdGx1RVqrz2DXNTLS8,3239
 vellum_ee/workflows/display/nodes/vellum/note_node.py,sha256=TMb8txILu2uWjzoxaghjgjlzeBAgzn4vkP_8zSh2qoE,1151
 vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py,sha256=LFjLUrH6sJ4czPnExdRqFr0PB_yKBMLXLvK5GAzIAgc,3273
-vellum_ee/workflows/display/nodes/vellum/retry_node.py,sha256=hVqTtuHjlw_cYJ3ydNAvUHfGEoQi5YocVONZUo4i_Gs,1717
+vellum_ee/workflows/display/nodes/vellum/retry_node.py,sha256=-eH41LUcpUeZ8Jt5f40-e7bn2nymeKnsYsmuprVSo7g,5805
 vellum_ee/workflows/display/nodes/vellum/search_node.py,sha256=TxcAGZDl_hvJ7Y1hUi9YVEVrj9Ie0hKkASdpfRL4_cs,9227
 vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py,sha256=62baAElKoRKIoba0lLhnrXGWWx96B73VxKGxh7BaIxc,2612
 vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=JVIMPR3WpveOCWZubHKZkE04mavnTdb_9QY_r3XliRg,3424
 vellum_ee/workflows/display/nodes/vellum/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/nodes/vellum/tests/test_error_node.py,sha256=ulrpoYUW-5kIxfG4Lf5F2p0k_EoYKhmahEbF3P_eruM,1648
 vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py,sha256=bg9INsXiWfyK047u8TD1oEOFYrqDq8GC7Hvgz69n7BE,1988
+vellum_ee/workflows/display/nodes/vellum/tests/test_retry_node.py,sha256=NuIw8Yb42KUdoGi3Ur8_7VPg50IC4hNrwAkCociwqNk,2091
 vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py,sha256=mtzB8LJlFCHVFM4H5AanLp29gQfaVmnN4A4iaRGJHoI,2427
 vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=4YUaTeD_OWF-UaPMyOTBTu9skGC1jgSHlAYrzbH7Z04,5039
 vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=HBfGz4yt9GlmMW9JxzaCacPnHBDNIeXE8Jhqr9DqLLw,6191
@@ -60,7 +61,7 @@ vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=cdpUoDN
 vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py,sha256=A1-tIpC5KIKG9JA_rkd1nLS8zUG3Kb4QiVdvb3boFxE,2509
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=9bIAEXXZQDqsUrDJqmHEeWYiZsYkVTQ4jBY-dPFVXEc,15054
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=jw8keqLNwqtaRSAqStsMqSH_OuuaTQm2MvIx8NTWAa4,14905
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=1cszL6N6FNGVm61MOa7AEiHnF0QjZWqDQuPOp4yiG94,18277
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py,sha256=-12ZkZb3f5gyoNASV2yeQtMo5HmNsVEo8nXwL6IC-I8,6261
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=6th6kCwzql6lddjkTQx4Jbvvs4ChqtJwctW-B4QuBhI,37352
@@ -122,7 +123,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=tKtdM1_GqmGq1gpi9ydWD_T-MM7fPn8QdHh8ww19cNI,117564
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=6EXwOd5Ka1mFWVlZmNuffO99dtW8DTAiRFes4vDlQFY,1869
+vellum/client/core/client_wrapper.py,sha256=MP_Nnw3gVGA7BJKaD9TBDDwVxta-T1ASzBcOv7dF4nU,1869
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -1314,7 +1315,7 @@ vellum/workflows/events/__init__.py,sha256=6pxxceJo2dcaRkWtkDAYlUQZ-PHBQSZytIoyu
 vellum/workflows/events/node.py,sha256=uHT6If0esgZ3nLjrjmUPTKf3qbjGhoV_x5YKpjDBDcU,5280
 vellum/workflows/events/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/events/tests/test_event.py,sha256=uRfMwSOqU-ROeZKCEngGvvJYlOZuxBhnC3qH5AGi3fM,15399
-vellum/workflows/events/types.py,sha256=cjRE8WL8tYCFradd9NOGl_H0mN3LiWWnA1uHmyT2Q0Q,3412
+vellum/workflows/events/types.py,sha256=AeTJaQt_fNHDLI4nyBzo7XrW9QQybRC09AKzu3kEYEE,3575
 vellum/workflows/events/workflow.py,sha256=sLO29djAeHGVd4hLhaNtOQ48uwUjfl-DotZQt06PxQA,6033
 vellum/workflows/exceptions.py,sha256=NiBiR3ggfmPxBVqD-H1SqmjI-7mIn0EStSN1BqApvCM,1213
 vellum/workflows/expressions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1411,12 +1412,14 @@ vellum/workflows/nodes/displayable/code_execution_node/node.py,sha256=pdDrjI8wdq
 vellum/workflows/nodes/displayable/code_execution_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/code_execution_node/tests/fixtures/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/code_execution_node/tests/fixtures/main.py,sha256=5QsbmkzSlSbcbWTG_JmIqcP-JNJzOPTKxGzdHos19W4,79
-vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py,sha256=OTnw5jgX90xWVB3vKryL5QSIr2YnGBz-W0q9C9LpNoc,22471
-vellum/workflows/nodes/displayable/code_execution_node/utils.py,sha256=BQraIN4I3DCzXLEuBlRYCyp7ote7hQmnnKHu4jFHCCA,5174
+vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py,sha256=6_RMWedAWU8Zrl0HpQorQgDfjV9khnhSouZeKadClyI,23441
+vellum/workflows/nodes/displayable/code_execution_node/utils.py,sha256=-7YdoF-NO6sYRFLrdkd0n1FhQFiZbJBK5jFcLNghgVo,3635
 vellum/workflows/nodes/displayable/conditional_node/__init__.py,sha256=AS_EIqFdU1F9t8aLmbZU-rLh9ry6LCJ0uj0D8F0L5Uw,72
 vellum/workflows/nodes/displayable/conditional_node/node.py,sha256=Qjfl33gZ3JEgxBA1EgzSUebboGvsARthIxxcQyvx5Gg,1152
 vellum/workflows/nodes/displayable/final_output_node/__init__.py,sha256=G7VXM4OWpubvSJtVkGmMNeqgb9GkM7qZT838eL18XU4,72
-vellum/workflows/nodes/displayable/final_output_node/node.py,sha256=TFbDZLGEtmg2cOQeJ56pUQdAkuHRa_qjBRIOGZU7Fy4,1990
+vellum/workflows/nodes/displayable/final_output_node/node.py,sha256=PuQ0RvtAmoSIZ5En_92tym7gpSMEoiHgwu20-UDbC7o,2368
+vellum/workflows/nodes/displayable/final_output_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vellum/workflows/nodes/displayable/final_output_node/tests/test_node.py,sha256=E6LQ74qZjY4Xi4avx2qdOCgGhF8pEcNLBh8cqYRkzMI,709
 vellum/workflows/nodes/displayable/guardrail_node/__init__.py,sha256=Ab5eXmOoBhyV4dMWdzh32HLUmnPIBEK_zFCT38C4Fng,68
 vellum/workflows/nodes/displayable/guardrail_node/node.py,sha256=h5nIBzQxbXTrdTq1wjDcekk1RV4-rKUNCshqdBAiJJY,4025
 vellum/workflows/nodes/displayable/inline_prompt_node/__init__.py,sha256=gSUOoEZLlrx35-tQhSAd3An8WDwBqyiQh-sIebLU9wU,74
@@ -1436,9 +1439,9 @@ vellum/workflows/nodes/displayable/search_node/node.py,sha256=_VHHuTNN4icZBgc7O5
 vellum/workflows/nodes/displayable/search_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/search_node/tests/test_node.py,sha256=2-QCV7Vk_-YMls33p0GOUtCv3f2uPNZCjkB2CRjek7o,6562
 vellum/workflows/nodes/displayable/subworkflow_deployment_node/__init__.py,sha256=9yYM6001YZeqI1VOk1QuEM_yrffk_EdsO7qaPzINKds,92
-vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py,sha256=epzcTald7En5vcIunCBXPT2t21FfR6V154rvGutkbLA,9596
+vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py,sha256=DGjvfusXCH5F98oPT3S6Xn0yedLuCq6EXb2DTzu2oDM,9661
 vellum/workflows/nodes/displayable/subworkflow_deployment_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/nodes/displayable/subworkflow_deployment_node/tests/test_node.py,sha256=Eo66jE2eHGD94U0XEfhpvGNG9kFDOBADQ1_5-1EFmIc,11735
+vellum/workflows/nodes/displayable/subworkflow_deployment_node/tests/test_node.py,sha256=7Yt8rmN7xoZ91BrcA1pDGwG_t0NYtym4ORXrAIIX-P0,15861
 vellum/workflows/nodes/displayable/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/displayable/tests/test_inline_text_prompt_node.py,sha256=LaxohBcKfSW2PSiBBlx67FdW_q4YC2BM2ouH-vuGPAA,4700
 vellum/workflows/nodes/displayable/tests/test_search_node_wth_text_output.py,sha256=VepO5z1277c1y5N6LLIC31nnWD1aak2m5oPFplfJHHs,6935
@@ -1448,9 +1451,9 @@ vellum/workflows/nodes/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQ
 vellum/workflows/nodes/experimental/openai_chat_completion_node/__init__.py,sha256=lsyD9laR9p7kx5-BXGH2gUTM242UhKy8SMV0SR6S2iE,90
 vellum/workflows/nodes/experimental/openai_chat_completion_node/node.py,sha256=1EGeiaT-Zoo6pttQFKKBcdf3dmhAbjKGaErYD5FFwlc,10185
 vellum/workflows/nodes/mocks.py,sha256=gvM2tyoe-V84jFbFdhQsyGAPyQBzmjn_CkhT_yxccgY,499
-vellum/workflows/nodes/utils.py,sha256=fq49n624n_nBMIFcc0JiNkHy3qna-bsA1HZRPa5AOJE,4918
+vellum/workflows/nodes/utils.py,sha256=uaTPGYp4utenz_QDghqQ23Q1iCsGHQ40nNZh1g9H9WI,7117
 vellum/workflows/outputs/__init__.py,sha256=AyZ4pRh_ACQIGvkf0byJO46EDnSix1ZCAXfvh-ms1QE,94
-vellum/workflows/outputs/base.py,sha256=2MtTlyzePopMZDpBWQIV8HRV-km1-z0vI8Gm012q9OQ,8665
+vellum/workflows/outputs/base.py,sha256=W5KL9FaWfSbZuF7lOQ677giHO839Do-MXmuzkDuzPqk,8607
 vellum/workflows/ports/__init__.py,sha256=bZuMt-R7z5bKwpu4uPW7LlJeePOQWmCcDSXe5frUY5g,101
 vellum/workflows/ports/node_ports.py,sha256=g4A-8iUAvEJSkaWppbvzAR8XU02R9U-qLN4rP2Kq4Aw,2743
 vellum/workflows/ports/port.py,sha256=eI2SOZPZ5rsC3jMsxW6Rbn0NpaYQsrR7AapiIbbiy8Q,3635
@@ -1501,13 +1504,13 @@ vellum/workflows/utils/uuids.py,sha256=DFzPv9RCvsKhvdTEIQyfSek2A31D6S_QcmeLPbgrg
 vellum/workflows/utils/vellum_variables.py,sha256=fC2aSLvlS31D15dOWu43LBRR0QsgUKNXBiCUvvaLXSs,3231
 vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528t75s,683
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
-vellum/workflows/workflows/base.py,sha256=ggnrtDWSScVRuH3PEs2A81qXX0LngYuvS62SkDS2SOE,22724
+vellum/workflows/workflows/base.py,sha256=TSS2BHC8LAi-N5GdEa75BeChwzwTzL7yldFnTlLINro,22665
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=NRteiICyJvDM5zrtUfq2fZoXcGQVaWC9xmNlLLVW0cU,7979
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-0.14.10.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.14.10.dist-info/METADATA,sha256=XNLaqqk85P8uRcr65-_IDfNPVi1dpNP9YvZgm43wAGs,5408
-vellum_ai-0.14.10.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.14.10.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.14.10.dist-info/RECORD,,
+vellum_ai-0.14.12.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.14.12.dist-info/METADATA,sha256=1EKavjMZSN9SKB-BwKDVMb4xfc-TnvoUWiYD7JMCqLk,5408
+vellum_ai-0.14.12.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.14.12.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.14.12.dist-info/RECORD,,

vellum_cli/image_push.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import re
 import subprocess
 from typing import List, Optional
 
@@ -24,9 +25,11 @@ def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
     # listing all of the architectures of the image instead of just the one that matches the machine. We can fall back
     # to using normal inspect which returns the machine image for this case though. And in the future we could figure
     # out how to call the docker host directly to do this.
+    logger.info("Pre-validating image...")
     docker_client = docker.from_env()
     check_architecture(docker_client, image, logger)
 
+    logger.info("Authenticating...")
     auth = vellum_client.container_images.docker_service_token()
 
     docker_client.login(
@@ -70,9 +73,29 @@ def image_push_command(image: str, tags: Optional[List[str]] = None) -> None:
             except Exception:
                 continue
 
-    logger.info("Updating Vellum metadata and enforcing the first law of robotics...")
-    image_details = docker_client.api.inspect_image(image)
-    sha = image_details["Id"]
+    result = subprocess.run(
+        ["docker", "inspect", "--format='{{index .RepoDigests 0}}'", image],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+
+    sha = ""
+    if result.returncode == 0:
+        match = re.search(r"sha256[^']*", result.stdout.decode("utf-8"))
+        if match and match.group(0):
+            sha = match.group(0)
+
+    if not sha:
+        # Fallback to using docker client if command line fails, at least on some systems
+        # this appears to give a bad sha.
+        logger.warning(
+            "Could not determine image hash with command line docker falling back to python docker client..."
+        )
+
+        image_details = docker_client.api.inspect_image(image)
+        sha = image_details["Id"]
+
+    logger.info(f"Updating Vellum metadata and validating image works in our system with image digest: {sha}...")
 
     vellum_client.container_images.push_container_image(
         name=image_name,
@@ -95,12 +118,12 @@ def check_architecture(docker_client: DockerClient, image: str, logger: logging.
         manifest = json.loads(result.stdout)
         architectures = [manifest_item["platform"]["architecture"] for manifest_item in manifest["manifests"]]
     except Exception:
-        logger.warning("Error parsing manifest response")
+        logger.debug("Error parsing manifest response")
         manifest_parse_failed = True
 
     # Fall back to inspect image if we errored out using docker command line
     if result.returncode != 0 or manifest_parse_failed:
-        logger.warning(f"Error inspecting manifest: {result.stderr.decode('utf-8').strip()}")
+        logger.debug(f"Error inspecting manifest: {result.stderr.decode('utf-8').strip()}")
         image_details = docker_client.api.inspect_image(image)
 
         if image_details["Architecture"] != _SUPPORTED_ARCHITECTURE:

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

@@ -1,11 +1,19 @@
 import inspect
-from typing import Any, Generic, TypeVar, cast
+from typing import Any, Callable, Generic, Optional, Tuple, Type, TypeVar, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
+from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
+from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.workflows.base import BaseWorkflow
+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
 from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
@@ -39,4 +47,86 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
             "attributes": attributes,
         }
 
-        return super().serialize(display_context, adornment=adornment)
+        serialized_node = super().serialize(
+            display_context,
+            adornment=adornment,
+        )
+
+        if serialized_node["type"] == "GENERIC":
+            return serialized_node
+
+        serialized_node_definition = serialized_node.get("definition")
+        if isinstance(serialized_node_definition, dict):
+            serialized_node_definition_module = serialized_node_definition.get("module")
+            if isinstance(serialized_node_definition_module, list):
+                serialized_node_definition_module.extend(
+                    [
+                        serialized_node_definition["name"],
+                        ADORNMENT_MODULE_NAME,
+                    ]
+                )
+                serialized_node_definition["name"] = node.__name__
+
+        return serialized_node
+
+    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+        inner_node = self._node.__wrapped_node__
+        if not inner_node:
+            return super().get_node_output_display(output)
+
+        node_display_class = get_node_display_class(BaseNodeDisplay, inner_node)
+        node_display = node_display_class()
+
+        inner_output = getattr(inner_node.Outputs, output.name)
+        return node_display.get_node_output_display(inner_output)
+
+    @classmethod
+    def wrap(
+        cls,
+        max_attempts: int,
+        delay: Optional[float] = None,
+        retry_on_error_code: Optional[WorkflowErrorCode] = None,
+        retry_on_condition: Optional[BaseDescriptor] = None,
+    ) -> Callable[..., Type["BaseRetryNodeDisplay"]]:
+        _max_attempts = max_attempts
+        _delay = delay
+        _retry_on_error_code = retry_on_error_code
+        _retry_on_condition = retry_on_condition
+
+        NodeDisplayType = TypeVar("NodeDisplayType", bound=BaseNodeDisplay)
+
+        def decorator(inner_cls: Type[NodeDisplayType]) -> Type[NodeDisplayType]:
+            node_class = inner_cls.infer_node_class()
+            wrapped_node_class = cast(Type[BaseNode], node_class.__wrapped_node__)
+
+            class RetryNodeDisplay(BaseRetryNodeDisplay[node_class]):  # type: ignore[valid-type]
+                max_attempts = _max_attempts
+                delay = _delay
+                retry_on_error_code = _retry_on_error_code
+                retry_on_condition = _retry_on_condition
+
+            setattr(inner_cls, "__adorned_by__", RetryNodeDisplay)
+
+            # We must edit the node display class to use __wrapped_node__ everywhere it
+            # references the adorned node class, which is three places:
+
+            # 1. The node display class' parameterized type
+            original_base_node_display = get_original_base(inner_cls)
+            original_base_node_display.__args__ = (wrapped_node_class,)
+            inner_cls._node_display_registry[wrapped_node_class] = inner_cls
+
+            # 2. The node display class' output displays
+            old_outputs = list(inner_cls.output_display.keys())
+            for old_output in old_outputs:
+                new_output = getattr(wrapped_node_class.Outputs, old_output.name)
+                inner_cls.output_display[new_output] = inner_cls.output_display.pop(old_output)
+
+            # 3. The node display class' port displays
+            old_ports = list(inner_cls.port_displays.keys())
+            for old_port in old_ports:
+                new_port = getattr(wrapped_node_class.Ports, old_port.name)
+                inner_cls.port_displays[new_port] = inner_cls.port_displays.pop(old_port)
+
+            return inner_cls
+
+        return decorator

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

@@ -0,0 +1,47 @@
+from typing import Any, Dict, cast
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.core.retry_node.node import RetryNode
+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_retry_node_parameters():
+    """Test that RetryNode parameters are correctly serialized."""
+
+    # GIVEN a RetryNode with specific parameters
+    @RetryNode.wrap(max_attempts=5, delay=2.5, retry_on_error_code=WorkflowErrorCode.INVALID_INPUTS)
+    class MyRetryNode(BaseNode):
+        pass
+
+    # AND a workflow using the node
+    class MyWorkflow(BaseWorkflow):
+        graph = MyRetryNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
+
+    # THEN the correct inputs should be serialized on the node
+    serialized_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["type"] != "ENTRYPOINT"
+    )
+
+    retry_adornment = next(
+        adornment for adornment in serialized_node["adornments"] if adornment["label"] == "RetryNode"
+    )
+
+    max_attempts_attribute = next(attr for attr in retry_adornment["attributes"] if attr["name"] == "max_attempts")
+    assert max_attempts_attribute["value"]["value"]["value"] == 5
+
+    delay_attribute = next(attr for attr in retry_adornment["attributes"] if attr["name"] == "delay")
+    assert delay_attribute["value"]["value"]["value"] == 2.5
+
+    retry_on_error_code_attribute = next(
+        attr for attr in retry_adornment["attributes"] if attr["name"] == "retry_on_error_code"
+    )
+
+    assert retry_on_error_code_attribute["value"]["value"]["type"] == "STRING"
+    assert retry_on_error_code_attribute["value"]["value"]["value"] == "INVALID_INPUTS"

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

@@ -29,11 +29,8 @@ class InnerRetryGenericNode(BaseNode):
         output: str
 
 
-class InnerRetryGenericNodeDisplay(BaseNodeDisplay[InnerRetryGenericNode.__wrapped_node__]):  # type: ignore
-    pass
-
-
-class OuterRetryNodeDisplay(BaseRetryNodeDisplay[InnerRetryGenericNode]):  # type: ignore
+@BaseRetryNodeDisplay.wrap(max_attempts=3)
+class InnerRetryGenericNodeDisplay(BaseNodeDisplay[InnerRetryGenericNode]):
     pass
 
 
@@ -44,7 +41,6 @@ def test_serialize_node__retry(serialize_node):
         global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)},
         global_node_displays={
             InnerRetryGenericNode.__wrapped_node__: InnerRetryGenericNodeDisplay,
-            InnerRetryGenericNode: OuterRetryNodeDisplay,
         },
     )