vellum-ai 0.10.4__py3-none-any.whl → 0.10.6__py3-none-any.whl

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

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, List, Optional, TypeVar
+from typing import Any, ClassVar, Dict, Generic, List, Optional, TypeVar, Union, Tuple
 
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -29,7 +29,7 @@ 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 import ConditionalNode
-from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.core import JsonObject, ConditionType
 
 _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)
 
@@ -39,11 +39,13 @@ class RuleIdMap:
     id: UUID
     lhs: Optional["RuleIdMap"]
     rhs: Optional["RuleIdMap"]
+    field_node_input_id: Optional[UUID]
+    value_node_input_id: Optional[UUID]
 
 @dataclass
 class ConditionId:
     id: UUID
-    rule_group_id: UUID
+    rule_group_id: Optional[UUID]
 
 
 class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Generic[_ConditionalNodeType]):
@@ -73,8 +75,11 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
                 if rule_id_map is not None
                 else str(uuid4_from_hash(f"{node_id}|rule|{','.join(str(p) for p in path)}"))
             )
-            field_node_input_id = None
-            value_node_input_id = None
+
+            result = self.get_nested_rule_details_by_path(rule_ids, path) if rule_ids else None
+            if result is None:
+                result = self._generate_hash_for_rule_ids(node_id, rule_id)
+            current_id, field_node_input_id, value_node_input_id = result
 
             # Recursive step: Keep recursing until we hit the other descriptors
             if isinstance(descriptor, (AndExpression, OrExpression)):
@@ -97,79 +102,89 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
             # Base cases for other descriptors
             elif isinstance(descriptor, (IsNullExpression, IsNotNullExpression)):
                 expression_node_input = create_node_input(
-                    node_id, f"{rule_id}.field", descriptor._expression, display_context, None
+                    node_id, f"{current_id}.field", descriptor._expression, display_context, field_node_input_id
                 )
                 node_inputs.append(expression_node_input)
-                field_node_input_id = expression_node_input.id
+                field_node_input_id = UUID(expression_node_input.id)
                 operator = self._convert_descriptor_to_operator(descriptor)
 
             elif isinstance(descriptor, (BetweenExpression, NotBetweenExpression)):
                 field_node_input = create_node_input(
-                    node_id, f"{rule_id}.field", descriptor._value, display_context, None
+                    node_id, f"{current_id}.field", descriptor._value, display_context, field_node_input_id
                 )
                 value_node_input = create_node_input(
-                    node_id, f"{rule_id}.value", f"{descriptor._start},{descriptor._end}", display_context, None
+                    node_id, f"{current_id}.value", f"{descriptor._start},{descriptor._end}", display_context, value_node_input_id
                 )
                 node_inputs.extend([field_node_input, value_node_input])
                 operator = self._convert_descriptor_to_operator(descriptor)
-                field_node_input_id = field_node_input.id
-                value_node_input_id = value_node_input.id
+                field_node_input_id = UUID(field_node_input.id)
+                value_node_input_id = UUID(value_node_input.id)
 
             else:
                 lhs = descriptor._lhs  # type: ignore[attr-defined]
                 rhs = descriptor._rhs  # type: ignore[attr-defined]
 
-                lhs_node_input = create_node_input(node_id, f"{rule_id}.field", lhs, display_context, None)
+                lhs_node_input = create_node_input(node_id, f"{current_id}.field", lhs, display_context, field_node_input_id)
                 node_inputs.append(lhs_node_input)
 
                 if descriptor._rhs is not None:  # type: ignore[attr-defined]
-                    rhs_node_input = create_node_input(node_id, f"{rule_id}.value", rhs, display_context, None)
+                    rhs_node_input = create_node_input(node_id, f"{current_id}.value", rhs, display_context, value_node_input_id)
                     node_inputs.append(rhs_node_input)
-                    value_node_input_id = rhs_node_input.id
+                    value_node_input_id = UUID(rhs_node_input.id)
 
                 operator = self._convert_descriptor_to_operator(descriptor)
-                field_node_input_id = lhs_node_input.id
+                field_node_input_id = UUID(lhs_node_input.id)
 
             return {
-                "id": rule_id,
+                "id": str(current_id),
                 "rules": None,
                 "combinator": None,
                 "negated": False,
-                "field_node_input_id": field_node_input_id,
+                "field_node_input_id": str(field_node_input_id),
                 "operator": operator,
-                "value_node_input_id": value_node_input_id,
+                "value_node_input_id": str(value_node_input_id),
             }
 
         conditions = []
         for idx, port in enumerate(node.Ports):
-            if port._condition is None or port._condition_type is None:
-                continue
-
 
             condition_id = str(condition_ids[idx].id if condition_ids else uuid4_from_hash(f"{node_id}|conditions|{idx}"))
             rule_group_id = str(condition_ids[idx].rule_group_id if condition_ids else uuid4_from_hash(f"{condition_id}|rule_group"))
             source_handle_id = str(source_handle_ids.get(idx) or uuid4_from_hash(f"{node_id}|handles|{idx}"))
 
-            rule_ids = self._get_rule_ids()
-            condition_rule = serialize_rule(port._condition, [idx], rule_ids[idx] if len(rule_ids) > idx else None)
-            rules = condition_rule["rules"] if condition_rule["rules"] else [condition_rule]
-
-            conditions.append(
-                {
-                    "id": condition_id,
-                    "type": port._condition_type.value,
-                    "source_handle_id": source_handle_id,
-                    "data": {
-                        "id": rule_group_id,
-                        "rules": rules,
-                        "combinator": "AND",
-                        "negated": False,
-                        "field_node_input_id": None,
-                        "operator": None,
-                        "value_node_input_id": None,
-                    },
-                }
-            )
+            if port._condition is None:
+                if port._condition_type == ConditionType.ELSE:
+                    conditions.append({
+                        "id": condition_id,
+                        "type": ConditionType.ELSE.value,
+                        "source_handle_id": source_handle_id,
+                        "data": None
+                    })
+                else:
+                    continue
+
+            else:
+                rule_ids = self._get_rule_ids()
+                condition_rule = serialize_rule(port._condition, [idx], rule_ids[idx] if len(rule_ids) > idx else None)
+                rules = condition_rule["rules"] if condition_rule["rules"] else [condition_rule]
+                if port._condition_type is None:
+                    raise ValueError("Condition type is None, but a valid ConditionType is expected.")
+                conditions.append(
+                    {
+                        "id": condition_id,
+                        "type": port._condition_type.value,
+                        "source_handle_id": source_handle_id,
+                        "data": {
+                            "id": rule_group_id,
+                            "rules": rules,
+                            "combinator": "AND",
+                            "negated": False,
+                            "field_node_input_id": None,
+                            "operator": None,
+                            "value_node_input_id": None,
+                        },
+                    }
+                )
 
         return {
             "id": str(node_id),
@@ -225,6 +240,36 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
         else:
             raise ValueError(f"Unsupported descriptor type: {descriptor}")
 
+    def get_nested_rule_details_by_path(
+            self, rule_ids: List[RuleIdMap], path: List[int]
+    ) -> Union[Tuple[UUID, Optional[UUID], Optional[UUID]], None]:
+        current_rule = rule_ids[path[0]]
+
+        for step in path[1:]:
+            if step == 0 and current_rule.lhs:
+                current_rule = current_rule.lhs
+            elif step == 1 and current_rule.rhs:
+                current_rule = current_rule.rhs
+            else:
+                return None
+
+        # This is essentially a leaf
+        if current_rule.lhs and current_rule.lhs.lhs is None and current_rule.lhs.rhs is None:
+            return (
+                current_rule.lhs.id,
+                current_rule.lhs.field_node_input_id,
+                current_rule.lhs.value_node_input_id,
+            )
+
+        return None
+
+    def _generate_hash_for_rule_ids(self, node_id, rule_id) -> Tuple[UUID, UUID, UUID]:
+        return (
+            uuid4_from_hash(f"{node_id}|{rule_id}|current"),
+            uuid4_from_hash(f"{node_id}|{rule_id}||field"),
+            uuid4_from_hash(f"{node_id}|{rule_id}||value")
+        )
+
     def _get_source_handle_ids(self) -> Dict[int, UUID]:
         return self._get_explicit_node_display_attr("source_handle_ids", Dict[int, UUID]) or {}
 

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

@@ -1,12 +1,12 @@
 from uuid import UUID
 from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
 
+from vellum.workflows.nodes import GuardrailNode
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 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.workflows.nodes import GuardrailNode
-from vellum.workflows.types.core import JsonObject
 
 _GuardrailNodeType = TypeVar("_GuardrailNodeType", bound=GuardrailNode)
 
@@ -15,7 +15,7 @@ class BaseGuardrailNodeDisplay(BaseNodeVellumDisplay[_GuardrailNodeType], Generi
     metric_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id

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

@@ -2,7 +2,9 @@ from uuid import UUID
 from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
 
 from vellum import PromptBlock, RichTextChildBlock, VellumVariable
-
+from vellum.workflows.nodes import InlinePromptNode
+from vellum.workflows.references import OutputReference
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -10,9 +12,6 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
-from vellum.workflows.nodes import InlinePromptNode
-from vellum.workflows.references import OutputReference
-from vellum.workflows.types.core import JsonObject
 
 _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
 
@@ -23,7 +22,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
     prompt_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id

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

@@ -1,17 +1,16 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, cast
 
 from vellum import VellumVariable
-
+from vellum.workflows.nodes import InlineSubworkflowNode
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 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.vellum import infer_vellum_variable_type
-from vellum_ee.workflows.display.vellum import NodeInput
+from vellum_ee.workflows.display.vellum import NodeInput, WorkflowOutputVellumDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum.workflows.nodes import InlineSubworkflowNode
-from vellum.workflows.types.core import JsonObject
 
 _InlineSubworkflowNodeType = TypeVar("_InlineSubworkflowNodeType", bound=InlineSubworkflowNode)
 
@@ -22,19 +21,19 @@ class BaseInlineSubworkflowNodeDisplay(
     workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
         node_inputs, workflow_inputs = self._generate_node_and_workflow_inputs(node_id, node, display_context)
-        workflow_outputs = self._generate_workflow_outputs(node)
 
         subworkflow_display = get_workflow_display(
             base_display_class=display_context.workflow_display_class,
             workflow_class=raise_if_descriptor(node.subworkflow),
             parent_display_context=display_context,
         )
+        workflow_outputs = self._generate_workflow_outputs(node, subworkflow_display.display_context)
         serialized_subworkflow = subworkflow_display.serialize()
 
         return {
@@ -87,13 +86,14 @@ class BaseInlineSubworkflowNodeDisplay(
     def _generate_workflow_outputs(
         self,
         node: Type[InlineSubworkflowNode],
+        display_context: WorkflowDisplayContext,
     ) -> List[VellumVariable]:
         workflow_outputs: List[VellumVariable] = []
         for output_descriptor in raise_if_descriptor(node.subworkflow).Outputs:  # type: ignore[union-attr]
-            node_output_display = self.get_node_output_display(output_descriptor)
+            workflow_output_display = cast(WorkflowOutputVellumDisplay, display_context.workflow_output_displays[output_descriptor])
             output_type = infer_vellum_variable_type(output_descriptor)
             workflow_outputs.append(
-                VellumVariable(id=str(node_output_display.id), key=node_output_display.name, type=output_type)
+                VellumVariable(id=str(workflow_output_display.id), key=workflow_output_display.name, type=output_type)
             )
 
         return workflow_outputs

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

@@ -2,7 +2,8 @@ from uuid import UUID
 from typing import Any, ClassVar, Dict, Generic, List, Optional, Type, TypeVar
 
 from vellum import VellumVariable
-
+from vellum.workflows.nodes import MapNode
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -10,8 +11,6 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum.workflows.nodes import MapNode
-from vellum.workflows.types.core import JsonObject
 
 _MapNodeType = TypeVar("_MapNodeType", bound=MapNode)
 
@@ -20,7 +19,7 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
     workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
@@ -32,7 +31,8 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
             # In Vellum it's always 'items'
             variable_name = descriptor.name if descriptor.name != "all_items" else "items"
             variable_id = str(
-                self.workflow_input_ids_by_name.get(variable_name) or uuid4_from_hash(f"{self.node_id}|{variable_name}")
+                self.workflow_input_ids_by_name.get(variable_name)
+                or uuid4_from_hash(f"{self.node_id}|{variable_name}")
             )
             workflow_inputs.append(
                 VellumVariable(

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

@@ -0,0 +1,32 @@
+import json
+from typing import Any, ClassVar, Dict, Generic, TypeVar, Union
+
+from vellum.workflows.nodes import NoteNode
+from vellum.workflows.types.core import JsonObject
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+
+_NoteNodeType = TypeVar("_NoteNodeType", bound=NoteNode)
+
+
+class BaseNoteNodeDisplay(BaseNodeVellumDisplay[_NoteNodeType], Generic[_NoteNodeType]):
+    text: ClassVar[str] = ""
+    style: ClassVar[Union[Dict[str, Any], None]] = None
+
+    def serialize(
+        self, display_context: WorkflowDisplayContext, **kwargs: Any
+    ) -> JsonObject:
+        node_id = self.node_id
+
+        return {
+            "id": str(node_id),
+            "type": "NOTE",
+            "inputs": [],
+            "data": {
+                "label": self.label,
+                "text": self.text,
+                "style": json.dumps(self.style) if self.style else None,
+            },
+            "display_data": self.get_display_data().dict(),
+            "definition": self.get_definition().dict(),
+        }

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

@@ -1,14 +1,14 @@
 from uuid import UUID
 from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar, cast
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
-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.workflows.nodes.displayable.prompt_deployment_node import PromptDeploymentNode
 from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.vellum_client import create_vellum_client
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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
 
 _PromptDeploymentNodeType = TypeVar("_PromptDeploymentNodeType", bound=PromptDeploymentNode)
 
@@ -21,7 +21,7 @@ class BasePromptDeploymentNodeDisplay(
     prompt_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id

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

@@ -31,7 +31,7 @@ class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_Sea
     variable_ids: Optional[VariableIdMap] = None
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id

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

@@ -16,9 +16,9 @@ class BaseSubworkflowDeploymentNodeDisplay(
     BaseNodeVellumDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
 ):
     subworkflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
-    
+
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id

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

@@ -17,7 +17,7 @@ class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Gene
     input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id

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

@@ -1,13 +1,13 @@
 from uuid import UUID
 from typing import Any, ClassVar, Generic, Optional, TypeVar
 
+from vellum.workflows.nodes.core.try_node.node import TryNode
+from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME, get_wrapped_node
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
-from vellum.workflows.nodes.core.try_node.node import TryNode
-from vellum.workflows.nodes.utils import get_wrapped_node
-from vellum.workflows.types.core import JsonObject
 
 _TryNodeType = TypeVar("_TryNodeType", bound=TryNode)
 
@@ -32,7 +32,19 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
 
         serialized_node = node_display.serialize(
             display_context,
-            error_output_id=str(self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id")),
+            error_output_id=self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id"),
         )
 
+        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

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

@@ -1,10 +1,12 @@
 from deepdiff import DeepDiff
 
-from tests.workflows.basic_code_execution_node.try_workflow import TrySimpleCodeExecutionWorkflow
-from tests.workflows.basic_code_execution_node.workflow import SimpleCodeExecutionWorkflow
+from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
+from tests.workflows.basic_code_execution_node.try_workflow import TrySimpleCodeExecutionWorkflow
+from tests.workflows.basic_code_execution_node.workflow import SimpleCodeExecutionWorkflow
+
 
 def test_serialize_workflow():
     # GIVEN a Workflow with a code execution node
@@ -428,8 +430,10 @@ def test_serialize_workflow__try_wrapped():
                 "workflows",
                 "basic_code_execution_node",
                 "try_workflow",
+                "SimpleCodeExecutionNode",
+                ADORNMENT_MODULE_NAME,
             ],
-            "name": "SimpleCodeExecutionNode",
+            "name": "TryNode",
         },
     }