vellum-ai 0.14.40__tar.gz → 0.14.42__tar.gz

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/PKG-INFO

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

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -30,6 +30,7 @@ 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.does_not_equal import DoesNotEqualExpression
 from vellum.workflows.expressions.ends_with import EndsWithExpression
 from vellum.workflows.expressions.equals import EqualsExpression
 from vellum.workflows.expressions.greater_than import GreaterThanExpression
@@ -49,7 +50,7 @@ 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_wrapped_node
+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
@@ -64,7 +65,7 @@ from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.names import pascal_to_title_case
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
-from vellum_ee.workflows.display.editor.types import NodeDisplayData
+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
@@ -177,29 +178,7 @@ 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])
 
-        ports: JsonArray = []
-        for port in node.Ports:
-            id = str(self.get_node_port_display(port).id)
-
-            if port._condition_type:
-                ports.append(
-                    {
-                        "id": id,
-                        "name": port.name,
-                        "type": port._condition_type.value,
-                        "expression": (
-                            self.serialize_condition(display_context, port._condition) if port._condition else None
-                        ),
-                    }
-                )
-            else:
-                ports.append(
-                    {
-                        "id": id,
-                        "name": port.name,
-                        "type": "DEFAULT",
-                    }
-                )
+        ports = self.serialize_ports(display_context)
 
         outputs: JsonArray = []
         for output in node.Outputs:
@@ -223,7 +202,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             "id": str(node_id),
             "label": node.__qualname__,
             "type": "GENERIC",
-            "display_data": self._get_generic_node_display_data().dict(),
+            "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
             "trigger": {
@@ -236,6 +215,35 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             "outputs": outputs,
         }
 
+    def serialize_ports(self, display_context: "WorkflowDisplayContext") -> JsonArray:
+        """Serialize the ports of the node."""
+        node = self._node
+        ports: JsonArray = []
+
+        for port in node.Ports:
+            id = str(self.get_node_port_display(port).id)
+            if port._condition_type:
+                ports.append(
+                    {
+                        "id": id,
+                        "name": port.name,
+                        "type": port._condition_type.value,
+                        "expression": (
+                            self.serialize_condition(display_context, port._condition) if port._condition else None
+                        ),
+                    }
+                )
+            else:
+                ports.append(
+                    {
+                        "id": id,
+                        "name": port.name,
+                        "type": "DEFAULT",
+                    }
+                )
+
+        return ports
+
     def get_base(self) -> CodeResourceDefinition:
         node = self._node
 
@@ -276,6 +284,13 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         return PortDisplay(id=port_id, node_id=self.node_id)
 
+    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_display = port_displays[default_port]
+        return default_port_display.id
+
     def get_trigger_id(self) -> UUID:
         return self.get_target_handle_id()
 
@@ -363,9 +378,32 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         register_node_display_class(node_class=node_class, node_display_class=cls)
 
-    def _get_generic_node_display_data(self) -> NodeDisplayData:
+    def get_display_data(self) -> NodeDisplayData:
         explicit_value = self._get_explicit_node_display_attr("display_data", NodeDisplayData)
-        return explicit_value if explicit_value else NodeDisplayData()
+        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,
+            )
+
+        if explicit_value:
+            return explicit_value
+
+        if docstring:
+            return NodeDisplayData(
+                comment=NodeDisplayComment(value=docstring),
+            )
+
+        return NodeDisplayData()
 
     def serialize_condition(self, display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
         if isinstance(
@@ -408,6 +446,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 DoesNotBeginWithExpression,
                 DoesNotContainExpression,
                 DoesNotEndWithExpression,
+                DoesNotEqualExpression,
                 EndsWithExpression,
                 EqualsExpression,
                 GreaterThanExpression,

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -82,3 +82,21 @@ def test_serialize_condition__accessor_expression():
             },
         }
     ]
+
+
+def test_serialize_display_data():
+    # GIVEN a node with an accessor expression in a Port
+    class MyNode(BaseNode):
+        """I hope this works"""
+
+        pass
+
+    # WHEN we serialize the node
+    node_display_class = get_node_display_class(MyNode)
+    data = node_display_class().serialize(WorkflowDisplayContext())
+
+    # THEN the condition should be serialized correctly
+    assert data["display_data"] == {
+        "position": {"x": 0.0, "y": 0.0},
+        "comment": {"value": "I hope this works"},
+    }

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/api_node.py

@@ -4,7 +4,7 @@ from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar, cast
 from vellum.workflows.nodes.displayable import APINode
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -13,7 +13,7 @@ from vellum_ee.workflows.display.utils.vellum import WorkspaceSecretPointer
 _APINodeType = TypeVar("_APINodeType", bound=APINode)
 
 
-class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeType]):
+class BaseAPINodeDisplay(BaseNodeDisplay[_APINodeType], Generic[_APINodeType]):
     # A mapping between node input keys and their ids for inputs representing additional header keys
     additional_header_key_input_ids: ClassVar[Optional[Dict[str, UUID]]] = None
 
@@ -192,4 +192,5 @@ class BaseAPINodeDisplay(BaseNodeVellumDisplay[_APINodeType], Generic[_APINodeTy
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/base_adornment_node.py

@@ -10,7 +10,6 @@ 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_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
-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.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -72,7 +71,7 @@ def _recursively_replace_wrapped_node(node_class: Type[BaseNode], wrapped_node_d
         )
 
 
-class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Generic[_BaseAdornmentNodeType]):
+class BaseAdornmentNodeDisplay(BaseNodeDisplay[_BaseAdornmentNodeType], Generic[_BaseAdornmentNodeType]):
     __wrapped_node_display__: Optional[Type[BaseNodeDisplay]] = None
 
     def serialize(

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -6,7 +6,7 @@ from vellum.workflows.nodes.displayable.code_execution_node import CodeExecution
 from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -14,7 +14,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _CodeExecutionNodeType = TypeVar("_CodeExecutionNodeType", bound=CodeExecutionNode)
 
 
-class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType], Generic[_CodeExecutionNodeType]):
+class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Generic[_CodeExecutionNodeType]):
     output_id: ClassVar[Optional[UUID]] = None
     log_output_id: ClassVar[Optional[UUID]] = None
 
@@ -98,4 +98,5 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/conditional_node.py

@@ -16,7 +16,7 @@ from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.nodes.displayable import ConditionalNode
 from vellum.workflows.types.core import ConditionType, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 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 convert_descriptor_to_operator
@@ -40,7 +40,7 @@ class ConditionId:
     rule_group_id: Optional[str]
 
 
-class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Generic[_ConditionalNodeType]):
+class BaseConditionalNodeDisplay(BaseNodeDisplay[_ConditionalNodeType], Generic[_ConditionalNodeType]):
     source_handle_ids: ClassVar[Dict[int, UUID]]
     rule_ids: ClassVar[List[RuleIdMap]]
     condition_ids: ClassVar[list[ConditionId]]
@@ -217,6 +217,7 @@ but the defined conditions have length {len(condition_ids)}"""
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }
 
     def get_nested_rule_details_by_path(

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/error_node.py

@@ -3,7 +3,7 @@ from typing import ClassVar, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes import ErrorNode
 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.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -11,7 +11,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _ErrorNodeType = TypeVar("_ErrorNodeType", bound=ErrorNode)
 
 
-class BaseErrorNodeDisplay(BaseNodeVellumDisplay[_ErrorNodeType], Generic[_ErrorNodeType]):
+class BaseErrorNodeDisplay(BaseNodeDisplay[_ErrorNodeType], Generic[_ErrorNodeType]):
     error_output_id: ClassVar[Optional[UUID]] = None
 
     name: ClassVar[str] = "error-node"

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/final_output_node.py

@@ -4,7 +4,7 @@ from typing import Any, ClassVar, Generic, Optional, TypeVar
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import to_kebab_case
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -12,11 +12,11 @@ from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 
 _FinalOutputNodeType = TypeVar("_FinalOutputNodeType", bound=FinalOutputNode)
 
+NODE_INPUT_KEY = "node_input"
 
-class BaseFinalOutputNodeDisplay(BaseNodeVellumDisplay[_FinalOutputNodeType], Generic[_FinalOutputNodeType]):
-    output_id: ClassVar[Optional[UUID]] = None
+
+class BaseFinalOutputNodeDisplay(BaseNodeDisplay[_FinalOutputNodeType], Generic[_FinalOutputNodeType]):
     output_name: ClassVar[Optional[str]] = None
-    node_input_id: ClassVar[Optional[UUID]] = None
 
     def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
         node = self._node
@@ -24,7 +24,7 @@ class BaseFinalOutputNodeDisplay(BaseNodeVellumDisplay[_FinalOutputNodeType], Ge
 
         node_input = create_node_input(
             node_id,
-            "node_input",
+            NODE_INPUT_KEY,
             # Get the pointer that the Terminal Node's output is referencing
             node.Outputs.value.instance,
             display_context,
@@ -58,13 +58,13 @@ class BaseFinalOutputNodeDisplay(BaseNodeVellumDisplay[_FinalOutputNodeType], Ge
         }
 
     def _get_output_id(self) -> UUID:
-        explicit_value = self._get_explicit_node_display_attr("output_id", UUID)
-        return explicit_value if explicit_value else uuid4_from_hash(f"{self.node_id}|output_id")
+        explicit_value = self.output_display.get(self._node.Outputs.value)
+        return explicit_value.id if explicit_value else uuid4_from_hash(f"{self.node_id}|output_id")
 
     def _get_output_name(self) -> str:
         explicit_value = self._get_explicit_node_display_attr("output_name", str)
         return explicit_value if explicit_value else to_kebab_case(self._node.__name__)
 
     def _get_node_input_id(self) -> UUID:
-        explicit_value = self._get_explicit_node_display_attr("node_input_id", UUID)
+        explicit_value = self.node_input_ids_by_name.get(NODE_INPUT_KEY)
         return explicit_value if explicit_value else uuid4_from_hash(f"{self.node_id}|node_input_id")

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/guardrail_node.py

@@ -3,7 +3,7 @@ from typing import 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.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -11,7 +11,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _GuardrailNodeType = TypeVar("_GuardrailNodeType", bound=GuardrailNode)
 
 
-class BaseGuardrailNodeDisplay(BaseNodeVellumDisplay[_GuardrailNodeType], Generic[_GuardrailNodeType]):
+class BaseGuardrailNodeDisplay(BaseNodeDisplay[_GuardrailNodeType], Generic[_GuardrailNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
@@ -46,4 +46,5 @@ class BaseGuardrailNodeDisplay(BaseNodeVellumDisplay[_GuardrailNodeType], Generi
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -6,7 +6,7 @@ from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.functions import compile_function_definition
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -16,7 +16,7 @@ from vellum_ee.workflows.display.vellum import NodeInput
 _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
 
 
-class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType], Generic[_InlinePromptNodeType]):
+class BaseInlinePromptNodeDisplay(BaseNodeDisplay[_InlinePromptNodeType], Generic[_InlinePromptNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
@@ -72,6 +72,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
                 {"id": str(output_display.id), "name": "text", "type": "STRING", "value": None},
                 {"id": str(array_display.id), "name": "results", "type": "ARRAY", "value": None},
             ],
+            "ports": self.serialize_ports(display_context),
         }
 
     def _generate_node_and_prompt_inputs(

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py

@@ -5,7 +5,7 @@ from vellum import VellumVariable
 from vellum.workflows.inputs.base import BaseInputs
 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.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -17,7 +17,7 @@ _InlineSubworkflowNodeType = TypeVar("_InlineSubworkflowNodeType", bound=InlineS
 
 
 class BaseInlineSubworkflowNodeDisplay(
-    BaseNodeVellumDisplay[_InlineSubworkflowNodeType], Generic[_InlineSubworkflowNodeType]
+    BaseNodeDisplay[_InlineSubworkflowNodeType], Generic[_InlineSubworkflowNodeType]
 ):
     workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
@@ -54,6 +54,7 @@ class BaseInlineSubworkflowNodeDisplay(
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }
 
     def _generate_node_and_workflow_inputs(

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/map_node.py

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Dict, Generic, List, Optional, TypeVar, cast
+from typing import Generic, Optional, TypeVar, cast
 
 from vellum.workflows.nodes import MapNode
 from vellum.workflows.types.core import JsonObject
@@ -36,24 +36,26 @@ class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNod
             parent_display_context=display_context,
         )
         serialized_subworkflow = subworkflow_display.serialize()
+        if not isinstance(serialized_subworkflow["input_variables"], list):
+            raise ValueError("input_variables must be a list")
 
-        renamed_input_variables = []
-        for input_variable in cast(List[Dict[str, str]], serialized_subworkflow["input_variables"]):
-            if input_variable["key"] == "all_items":
-                renamed_item = {**input_variable, "key": "items"}
-                renamed_input_variables.append(renamed_item)
-            else:
-                renamed_input_variables.append(input_variable)
+        input_variables = serialized_subworkflow["input_variables"]
 
         # Note: This must match the items input ID for the map node's node input
         items_workflow_input_id = next(
-            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "items"
+            input_variable["id"]
+            for input_variable in input_variables
+            if isinstance(input_variable, dict) and input_variable["key"] == "items"
         )
         item_workflow_input_id = next(
-            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "item"
+            input_variable["id"]
+            for input_variable in input_variables
+            if isinstance(input_variable, dict) and input_variable["key"] == "item"
         )
         index_workflow_input_id = next(
-            input_variable["id"] for input_variable in renamed_input_variables if input_variable["key"] == "index"
+            input_variable["id"]
+            for input_variable in input_variables
+            if isinstance(input_variable, dict) and input_variable["key"] == "index"
         )
 
         return {
@@ -67,7 +69,7 @@ class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNod
                 "target_handle_id": str(self.get_target_handle_id()),
                 "variant": "INLINE",
                 "workflow_raw_data": serialized_subworkflow["workflow_raw_data"],
-                "input_variables": cast(JsonObject, renamed_input_variables),
+                "input_variables": cast(JsonObject, input_variables),
                 "output_variables": serialized_subworkflow["output_variables"],
                 "concurrency": raise_if_descriptor(node.max_concurrency),
                 "items_input_id": items_workflow_input_id,
@@ -77,4 +79,5 @@ class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNod
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/merge_node.py

@@ -4,13 +4,13 @@ from typing import Any, ClassVar, Generic, List, Optional, TypeVar
 from vellum.workflows.nodes.displayable import MergeNode
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _MergeNodeType = TypeVar("_MergeNodeType", bound=MergeNode)
 
 
-class BaseMergeNodeDisplay(BaseNodeVellumDisplay[_MergeNodeType], Generic[_MergeNodeType]):
+class BaseMergeNodeDisplay(BaseNodeDisplay[_MergeNodeType], Generic[_MergeNodeType]):
     target_handle_ids: ClassVar[List[UUID]]
 
     def __init__(self):
@@ -50,6 +50,7 @@ class BaseMergeNodeDisplay(BaseNodeVellumDisplay[_MergeNodeType], Generic[_Merge
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }
 
     def get_target_handle_ids(self) -> Optional[List[UUID]]:

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/note_node.py

@@ -2,13 +2,13 @@ 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.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _NoteNodeType = TypeVar("_NoteNodeType", bound=NoteNode)
 
 
-class BaseNoteNodeDisplay(BaseNodeVellumDisplay[_NoteNodeType], Generic[_NoteNodeType]):
+class BaseNoteNodeDisplay(BaseNodeDisplay[_NoteNodeType], Generic[_NoteNodeType]):
     text: ClassVar[str] = ""
     style: ClassVar[Union[Dict[str, Any], None]] = None
 

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py

@@ -5,7 +5,7 @@ from vellum.workflows.nodes.displayable.prompt_deployment_node import PromptDepl
 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.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -13,9 +13,7 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _PromptDeploymentNodeType = TypeVar("_PromptDeploymentNodeType", bound=PromptDeploymentNode)
 
 
-class BasePromptDeploymentNodeDisplay(
-    BaseNodeVellumDisplay[_PromptDeploymentNodeType], Generic[_PromptDeploymentNodeType]
-):
+class BasePromptDeploymentNodeDisplay(BaseNodeDisplay[_PromptDeploymentNodeType], Generic[_PromptDeploymentNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
@@ -68,4 +66,5 @@ class BasePromptDeploymentNodeDisplay(
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/search_node.py

@@ -11,7 +11,7 @@ from vellum.workflows.nodes.displayable.search_node import SearchNode
 from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -28,7 +28,7 @@ class VariableIdMap:
     rhs: Optional["VariableIdMap"]
 
 
-class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_SearchNodeType]):
+class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNodeType]):
     # A mapping between the id of the operand (e.g. "lhs_variable_id" or "rhs_variable_id") and the id of the node input
     # that the operand is pointing to.
     metadata_filter_input_id_by_operand_id: Dict[UUID, UUID] = {}
@@ -68,6 +68,7 @@ class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_Sea
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }
 
     def _generate_search_node_inputs(

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py

@@ -4,7 +4,7 @@ from typing import Generic, Optional, TypeVar
 from vellum.workflows.nodes import SubworkflowDeploymentNode
 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.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -13,7 +13,7 @@ _SubworkflowDeploymentNodeType = TypeVar("_SubworkflowDeploymentNodeType", bound
 
 
 class BaseSubworkflowDeploymentNodeDisplay(
-    BaseNodeVellumDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
+    BaseNodeDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
 ):
 
     def serialize(
@@ -57,4 +57,5 @@ class BaseSubworkflowDeploymentNodeDisplay(
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }

{vellum_ai-0.14.40 → vellum_ai-0.14.42}/ee/vellum_ee/workflows/display/nodes/vellum/templating_node.py

@@ -4,7 +4,7 @@ from typing import Generic, Optional, TypeVar
 from vellum.workflows.nodes.core.templating_node import TemplatingNode
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -14,7 +14,7 @@ _TemplatingNodeType = TypeVar("_TemplatingNodeType", bound=TemplatingNode)
 TEMPLATE_INPUT_NAME = TemplatingNode.template.name
 
 
-class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Generic[_TemplatingNodeType]):
+class BaseTemplatingNodeDisplay(BaseNodeDisplay[_TemplatingNodeType], Generic[_TemplatingNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
@@ -67,4 +67,5 @@ class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Gene
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
+            "ports": self.serialize_ports(display_context),
         }

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

@@ -9,7 +9,7 @@ from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.references import LazyReference
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input_value_pointer_rules
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -36,7 +36,7 @@ class MyNodeA(BaseNode):
         output: str
 
 
-class MyNodeADisplay(BaseNodeVellumDisplay[MyNodeA]):
+class MyNodeADisplay(BaseNodeDisplay[MyNodeA]):
     pass