vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -8,7 +8,9 @@ from typing import (
     Dict,
     ForwardRef,
     Generic,
+    List,
     Optional,
+    Set,
     Tuple,
     Type,
     TypeVar,
@@ -20,44 +22,11 @@ from typing import (
 from vellum.client.types.code_resource_definition import CodeResourceDefinition
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.constants import undefined
-from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.expressions.accessor import AccessorExpression
-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
-from vellum.workflows.expressions.contains import ContainsExpression
-from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
-from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
-from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
-from vellum.workflows.expressions.ends_with import EndsWithExpression
-from vellum.workflows.expressions.equals import EqualsExpression
-from vellum.workflows.expressions.greater_than import GreaterThanExpression
-from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
-from vellum.workflows.expressions.in_ import InExpression
-from vellum.workflows.expressions.is_nil import IsNilExpression
-from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
-from vellum.workflows.expressions.is_not_null import IsNotNullExpression
-from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
-from vellum.workflows.expressions.is_null import IsNullExpression
-from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
-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.expressions.parse_json import ParseJsonExpression
 from vellum.workflows.nodes.bases.base import BaseNode
-from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.utils import get_unadorned_node, get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference
-from vellum.workflows.references.constant import ConstantValueReference
-from vellum.workflows.references.execution_count import ExecutionCountReference
-from vellum.workflows.references.lazy import LazyReference
-from vellum.workflows.references.state_value import StateValueReference
-from vellum.workflows.references.vellum_secret import VellumSecretReference
-from vellum.workflows.references.workflow_input import WorkflowInputReference
+from vellum.workflows.references.node import NodeReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.types.utils import get_original_base
@@ -67,10 +36,8 @@ from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_var
 from vellum_ee.workflows.display.editor.types import NodeDisplayComment, NodeDisplayData
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
-from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
-from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
+from vellum_ee.workflows.display.utils.expressions import serialize_value
 from vellum_ee.workflows.display.utils.registry import register_node_display_class
-from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -78,18 +45,40 @@ if TYPE_CHECKING:
 _NodeDisplayAttrType = TypeVar("_NodeDisplayAttrType")
 
 
+def _get_node_input_ids_by_ref(node_class: Type[BaseNode], path: str, inst: Any):
+    if isinstance(inst, dict):
+        node_input_ids_by_name: Dict[str, UUID] = {}
+        for key, value in inst.items():
+            node_input_ids_by_name.update(_get_node_input_ids_by_ref(node_class, f"{path}.{key}", value))
+        return node_input_ids_by_name
+
+    return {path: uuid4_from_hash(f"{node_class.__id__}|{path}")}
+
+
 class BaseNodeDisplayMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         cls = cast(Type["BaseNodeDisplay"], super().__new__(mcs, name, bases, dct))
+        # This cast shouldn't be necessary, but it's a workaround for a mypy bug
+        node_class = cast(Type[BaseNode], cls.infer_node_class() if name != "BaseNodeDisplay" else BaseNode)
 
         if not dct.get("output_display"):
-            node_class = cls.infer_node_class()
             cls.output_display = {
                 ref: NodeOutputDisplay(id=node_class.__output_ids__[ref.name], name=ref.name)
                 for ref in node_class.Outputs
                 if ref.name in node_class.__output_ids__
             }
 
+        if not dct.get("node_input_ids_by_name"):
+            node_input_ids_by_name: Dict[str, UUID] = {}
+            for ref in node_class:
+                if ref not in cls.__serializable_inputs__:
+                    continue
+
+                node_input_ids_by_name.update(_get_node_input_ids_by_ref(node_class, ref.name, ref.instance))
+
+            if node_input_ids_by_name:
+                cls.node_input_ids_by_name = node_input_ids_by_name
+
         return cls.__annotate_node__()
 
     def __annotate_node__(cls):
@@ -134,11 +123,22 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     # Default values set by the metaclass
     output_display: Dict[OutputReference, NodeOutputDisplay]
     port_displays: Dict[Port, PortDisplayOverrides] = {}
-    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
+    attribute_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
+    # START: Attributes for backwards compatible serialization
     # Used to explicitly set the target handle id for a node
     # Once all nodes are Generic Nodes, we may replace this with a trigger_id or trigger attribute
     target_handle_id: ClassVar[Optional[UUID]] = None
+    # Used to explicitly set the input ids for each node input
+    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
+    # Used by each class extending BaseNodeDisplay to specify which attributes are meant to be serialized
+    # as the former `"inputs"` field
+    __serializable_inputs__: Set[NodeReference] = set()
+    # Used by each class extending BaseNodeDisplay to specify which attributes are meant to be opted out
+    # of serialization. It's possible that we keep this one as a user facing api in the future, but
+    # don't want to commit to that decision just yet
+    __unserializable_attributes__: Set[NodeReference] = set()
+    # END: Attributes for backwards compatible serialization
 
     def serialize(self, display_context: "WorkflowDisplayContext", **kwargs: Any) -> JsonObject:
         node = self._node
@@ -150,13 +150,20 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 # We don't need to serialize generic node attributes containing a subworkflow
                 continue
 
-            id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            if attribute in self.__unserializable_attributes__:
+                continue
+
+            id = (
+                str(self.attribute_ids_by_name[attribute.name])
+                if self.attribute_ids_by_name.get(attribute.name)
+                else str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            )
             try:
                 attributes.append(
                     {
                         "id": id,
                         "name": attribute.name,
-                        "value": self.serialize_value(display_context, attribute.instance),
+                        "value": serialize_value(node_id, display_context, attribute.instance),
                     }
                 )
             except ValueError as e:
@@ -169,7 +176,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
             adornment: JsonObject = {
                 "id": str(node_id),
-                "label": node.__qualname__,
+                "label": self.label,
                 "base": self.get_base().dict(),
                 "attributes": attributes,
             }
@@ -177,10 +184,23 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             existing_adornments = adornments if adornments is not None else []
             return display_class().serialize(display_context, adornments=existing_adornments + [adornment])
 
+        return {
+            "id": str(node_id),
+            "label": self.label,
+            "type": "GENERIC",
+            "adornments": adornments,
+            "attributes": attributes,
+            **self.serialize_generic_fields(display_context),
+        }
+
+    def serialize_ports(self, display_context: "WorkflowDisplayContext") -> JsonArray:
+        """Serialize the ports of the node."""
+        node = self._node
+        node_id = self.node_id
         ports: JsonArray = []
+
         for port in node.Ports:
             id = str(self.get_node_port_display(port).id)
-
             if port._condition_type:
                 ports.append(
                     {
@@ -188,7 +208,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                         "name": port.name,
                         "type": port._condition_type.value,
                         "expression": (
-                            self.serialize_condition(display_context, port._condition) if port._condition else None
+                            serialize_value(node_id, display_context, port._condition) if port._condition else None
                         ),
                     }
                 )
@@ -201,41 +221,106 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                     }
                 )
 
+        return ports
+
+    def serialize_trigger(self) -> JsonObject:
+        """Serialize the trigger of the node."""
+        node = self._node
+        return {
+            "id": str(self.get_target_handle_id()),
+            "merge_behavior": node.Trigger.merge_behavior.value,
+        }
+
+    def _serialize_attributes(self, display_context: "WorkflowDisplayContext") -> JsonArray:
+        """Serialize node attributes, skipping unserializable ones."""
+        node = self._node
+        node_id = self.node_id
+        attributes: JsonArray = []
+        for attribute in node:
+            if attribute in self.__unserializable_attributes__:
+                continue
+
+            id = (
+                str(self.attribute_ids_by_name[attribute.name])
+                if self.attribute_ids_by_name.get(attribute.name)
+                else str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            )
+            try:
+                attribute_dict: JsonObject = {
+                    "id": id,
+                    "name": attribute.name,
+                    "value": serialize_value(node_id, display_context, attribute.instance),
+                }
+                attributes.append(attribute_dict)
+            except ValueError as e:
+                raise ValueError(f"Failed to serialize attribute '{attribute.name}': {e}")
+
+        return attributes
+
+    def _serialize_outputs(self, display_context: "WorkflowDisplayContext") -> JsonArray:
+        """Generate outputs array from node output displays or node.Outputs."""
         outputs: JsonArray = []
+        node = self._node
+
         for output in node.Outputs:
-            type = primitive_type_to_vellum_variable_type(output)
+            output_type = primitive_type_to_vellum_variable_type(output)
             value = (
-                self.serialize_value(display_context, output.instance)
+                serialize_value(self.node_id, display_context, output.instance)
                 if output.instance is not None and output.instance != undefined
                 else None
             )
 
+            output_id = (
+                str(self.output_display[output].id)
+                if output in self.output_display
+                else str(uuid4_from_hash(f"{self.node_id}|{output.name}"))
+            )
+
             outputs.append(
                 {
-                    "id": str(uuid4_from_hash(f"{node_id}|{output.name}")),
+                    "id": output_id,
                     "name": output.name,
-                    "type": type,
+                    "type": output_type,
                     "value": value,
                 }
             )
 
-        return {
-            "id": str(node_id),
-            "label": node.__qualname__,
-            "type": "GENERIC",
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "trigger": {
-                "id": str(self.get_target_handle_id()),
-                "merge_behavior": node.Trigger.merge_behavior.value,
-            },
-            "ports": ports,
-            "adornments": adornments,
-            "attributes": attributes,
-            "outputs": outputs,
+        return outputs
+
+    def serialize_generic_fields(
+        self, display_context: "WorkflowDisplayContext", exclude: Optional[List[str]] = None
+    ) -> JsonObject:
+        """Serialize generic fields that are common to all nodes."""
+        exclude = exclude or []
+
+        result: JsonObject = {
+            "display_data": self.get_display_data().dict() if "display_data" not in exclude else None,
+            "base": self.get_base().dict() if "base" not in exclude else None,
+            "definition": self.get_definition().dict() if "definition" not in exclude else None,
+            "trigger": self.serialize_trigger() if "trigger" not in exclude else None,
+            "ports": self.serialize_ports(display_context) if "ports" not in exclude else None,
+            "outputs": self._serialize_outputs(display_context) if "outputs" not in exclude else None,
         }
 
+        # Only include should_file_merge if there are custom methods defined
+        try:
+            node_class = self.__class__.infer_node_class()
+            has_custom_methods = any(
+                callable(getattr(node_class, name, None)) and inspect.isfunction(getattr(node_class, name, None))
+                for name in node_class.__dict__.keys()
+                if not name.startswith("__")
+            )
+
+            if has_custom_methods:
+                result["should_file_merge"] = True
+        except Exception:
+            pass
+
+        for key in exclude:
+            result.pop(key, None)
+
+        return result
+
     def get_base(self) -> CodeResourceDefinition:
         node = self._node
 
@@ -258,12 +343,12 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         )
         return node_definition
 
-    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
         explicit_display = self.output_display.get(output)
         if explicit_display:
-            return self._node, explicit_display
+            return explicit_display
 
-        return (self._node, NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name))
+        return NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name)
 
     def get_node_port_display(self, port: Port) -> PortDisplay:
         overrides = self.port_displays.get(port)
@@ -278,10 +363,18 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
     def get_source_handle_id(self, port_displays: Dict[Port, PortDisplay]) -> UUID:
         unadorned_node = get_unadorned_node(self._node)
-        default_port = unadorned_node.Ports.default
+        default_port = next((port for port in unadorned_node.Ports if port.default), None)
+        if default_port in port_displays:
+            return port_displays[default_port].id
 
-        default_port_display = port_displays[default_port]
-        return default_port_display.id
+        if default_port:
+            return self.get_node_port_display(default_port).id
+
+        first_port = next((port for port in unadorned_node.Ports), None)
+        if not first_port:
+            raise ValueError(f"Node {self._node.__name__} must have at least one port.")
+
+        return port_displays[first_port].id
 
     def get_trigger_id(self) -> UUID:
         return self.get_target_handle_id()
@@ -294,7 +387,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         return self._get_node_display_uuid("target_handle_id")
 
-    def get_target_handle_id_by_source_node_id(self, source_node_id: UUID) -> UUID:
+    def get_target_handle_id_by_source_node_id(self, _source_node_id: UUID) -> UUID:
         """
         In the vast majority of cases, nodes will only have a single target handle and can be retrieved independently
         of the source node. However, in rare cases (such as legacy Merge nodes), this method can be overridden to
@@ -306,7 +399,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     @cached_property
     def node_id(self) -> UUID:
         """Can be overridden as a class attribute to specify a custom node id."""
-        return uuid4_from_hash(self._node.__qualname__)
+        return self._node.__id__
 
     @cached_property
     def label(self) -> str:
@@ -371,162 +464,62 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         register_node_display_class(node_class=node_class, node_display_class=cls)
 
     def get_display_data(self) -> NodeDisplayData:
-        explicit_value = self._get_explicit_node_display_attr("display_data", NodeDisplayData)
+        """
+        Get display data by first building from BaseNode.Display, then overriding with BaseNodeDisplay.
+        This approach is more scalable as we add more attributes to Display class.
+        """
         docstring = self._node.__doc__
 
-        if explicit_value and explicit_value.comment and docstring:
-            comment = (
-                NodeDisplayComment(value=docstring, expanded=explicit_value.comment.expanded)
-                if explicit_value.comment.expanded
-                else NodeDisplayComment(value=docstring)
-            )
-            return NodeDisplayData(
-                position=explicit_value.position,
-                width=explicit_value.width,
-                height=explicit_value.height,
-                comment=comment,
-            )
+        # Step 1: Start with BaseNode.Display attributes as base
+        base_kwargs: Dict[str, Any] = {}
+        base_node = next((base for base in self._node.__bases__ if issubclass(base, BaseNode)), BaseNode)
 
-        if explicit_value:
-            return explicit_value
+        # Add Display class attributes if they exist
+        if self._node.Display.icon is not None and self._node.Display.icon != base_node.Display.icon:
+            base_kwargs["icon"] = self._node.Display.icon
+        if self._node.Display.color is not None and self._node.Display.color != base_node.Display.color:
+            base_kwargs["color"] = self._node.Display.color
 
+        # Add docstring as comment if present
         if docstring:
-            return NodeDisplayData(
-                comment=NodeDisplayComment(value=docstring),
-            )
+            base_kwargs["comment"] = NodeDisplayComment(value=docstring, expanded=True)
 
-        return NodeDisplayData()
-
-    def serialize_condition(self, display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
-        if isinstance(
-            condition,
-            (
-                IsNullExpression,
-                IsNotNullExpression,
-                IsNilExpression,
-                IsNotNilExpression,
-                IsUndefinedExpression,
-                IsNotUndefinedExpression,
-                ParseJsonExpression,
-            ),
-        ):
-            lhs = self.serialize_value(display_context, condition._expression)
-            return {
-                "type": "UNARY_EXPRESSION",
-                "lhs": lhs,
-                "operator": convert_descriptor_to_operator(condition),
-            }
-        elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
-            base = self.serialize_value(display_context, condition._value)
-            lhs = self.serialize_value(display_context, condition._start)
-            rhs = self.serialize_value(display_context, condition._end)
-
-            return {
-                "type": "TERNARY_EXPRESSION",
-                "base": base,
-                "operator": convert_descriptor_to_operator(condition),
-                "lhs": lhs,
-                "rhs": rhs,
-            }
-        elif isinstance(
-            condition,
-            (
-                AndExpression,
-                BeginsWithExpression,
-                CoalesceExpression,
-                ContainsExpression,
-                DoesNotBeginWithExpression,
-                DoesNotContainExpression,
-                DoesNotEndWithExpression,
-                EndsWithExpression,
-                EqualsExpression,
-                GreaterThanExpression,
-                GreaterThanOrEqualToExpression,
-                InExpression,
-                LessThanExpression,
-                LessThanOrEqualToExpression,
-                NotInExpression,
-                OrExpression,
-            ),
-        ):
-            lhs = self.serialize_value(display_context, condition._lhs)
-            rhs = self.serialize_value(display_context, condition._rhs)
-
-            return {
-                "type": "BINARY_EXPRESSION",
-                "lhs": lhs,
-                "operator": convert_descriptor_to_operator(condition),
-                "rhs": rhs,
-            }
-        elif isinstance(condition, AccessorExpression):
-            return {
-                "type": "BINARY_EXPRESSION",
-                "lhs": self.serialize_value(display_context, condition._base),
-                "operator": "accessField",
-                "rhs": self.serialize_value(display_context, condition._field),
-            }
-
-        raise UnsupportedSerializationException(f"Unsupported condition type: {condition.__class__.__name__}")
-
-    def serialize_value(self, display_context: "WorkflowDisplayContext", value: Any) -> JsonObject:
-        if isinstance(value, ConstantValueReference):
-            return self.serialize_value(display_context, value._value)
-
-        if isinstance(value, LazyReference):
-            child_descriptor = get_child_descriptor(value, display_context)
-            return self.serialize_value(display_context, child_descriptor)
-
-        if isinstance(value, WorkflowInputReference):
-            workflow_input_display = display_context.global_workflow_input_displays[value]
-            return {
-                "type": "WORKFLOW_INPUT",
-                "input_variable_id": str(workflow_input_display.id),
-            }
-
-        if isinstance(value, StateValueReference):
-            state_value_display = display_context.global_state_value_displays[value]
-            return {
-                "type": "STATE_VALUE",
-                "state_variable_id": str(state_value_display.id),
-            }
-
-        if isinstance(value, OutputReference):
-            upstream_node, output_display = display_context.global_node_output_displays[value]
-            upstream_node_display = display_context.global_node_displays[upstream_node]
-
-            return {
-                "type": "NODE_OUTPUT",
-                "node_id": str(upstream_node_display.node_id),
-                "node_output_id": str(output_display.id),
-            }
-
-        if isinstance(value, VellumSecretReference):
-            return {
-                "type": "VELLUM_SECRET",
-                "vellum_secret_name": value.name,
-            }
-
-        if isinstance(value, ExecutionCountReference):
-            node_class_display = display_context.global_node_displays[value.node_class]
-
-            return {
-                "type": "EXECUTION_COUNTER",
-                "node_id": str(node_class_display.node_id),
-            }
+        # Step 2: Get explicit BaseNodeDisplay values
+        explicit_value = self._get_explicit_node_display_attr("display_data", NodeDisplayData)
 
-        if isinstance(value, dict) and any(isinstance(v, BaseDescriptor) for v in value.values()):
-            raise ValueError("Nested references are not supported.")
+        # Step 3: Override with explicit BaseNodeDisplay values if present
+        if explicit_value:
+            # Get fields that were explicitly set for z_index handling
+            fields_set = explicit_value.model_fields_set
+
+            # Override simple attributes (only if not None)
+            for attr in ("position", "width", "height", "icon", "color"):
+                value = getattr(explicit_value, attr, None)
+                if value is not None:
+                    base_kwargs[attr] = value
+
+            # Include z_index if explicitly set (even if None)
+            if "z_index" in fields_set:
+                base_kwargs["z_index"] = explicit_value.z_index
+
+            # Special handling for comment: merge docstring with explicit comment's expanded state
+            if explicit_value.comment:
+                base_kwargs["comment"] = (
+                    NodeDisplayComment(
+                        value=docstring,
+                        expanded=(
+                            explicit_value.comment.expanded if explicit_value.comment.expanded is not None else True
+                        ),
+                    )
+                    if docstring
+                    else explicit_value.comment
+                )
 
-        if not isinstance(value, BaseDescriptor):
-            vellum_value = primitive_to_vellum_value(value)
-            return {
-                "type": "CONSTANT_VALUE",
-                "value": vellum_value.dict(),
-            }
+        # Step 4: Return the merged result
+        if base_kwargs:
+            return NodeDisplayData(**base_kwargs)
 
-        # 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)
+        return NodeDisplayData()
 
 
 register_node_display_class(node_class=BaseNode, node_display_class=BaseNodeDisplay)

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,10 +1,7 @@
 import types
-from uuid import UUID
-from typing import TYPE_CHECKING, Any, Dict, Generic, Type, TypeVar
+from typing import TYPE_CHECKING, Generic, Type, TypeVar
 
-from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.types.generics import NodeType
-from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.registry import get_from_node_display_registry
 
 if TYPE_CHECKING:
@@ -23,30 +20,9 @@ def get_node_display_class(node_class: Type[NodeType]) -> Type["BaseNodeDisplay"
     # `base_node_display_class` is always a Generic class, so it's safe to index into it
     NodeDisplayBaseClass = base_node_display_class[_NodeClassType]  # type: ignore[index]
 
-    def _get_node_input_ids_by_ref(path: str, inst: Any):
-        if isinstance(inst, dict):
-            node_input_ids_by_name: Dict[str, UUID] = {}
-            for key, value in inst.items():
-                node_input_ids_by_name.update(_get_node_input_ids_by_ref(f"{path}.{key}", value))
-            return node_input_ids_by_name
-
-        if isinstance(inst, BaseDescriptor):
-            return {path: uuid4_from_hash(f"{node_class.__id__}|{path}")}
-
-        return {}
-
-    def exec_body(ns: Dict):
-        node_input_ids_by_name: Dict[str, UUID] = {}
-        for ref in node_class:
-            node_input_ids_by_name.update(_get_node_input_ids_by_ref(ref.name, ref.instance))
-
-        if node_input_ids_by_name:
-            ns["node_input_ids_by_name"] = node_input_ids_by_name
-
     NodeDisplayClass = types.new_class(
         f"{node_class.__name__}Display",
         bases=(NodeDisplayBaseClass, Generic[_NodeClassType]),
-        exec_body=exec_body,
     )
 
     return NodeDisplayClass