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

vellum_ee/scripts/generate_node_definitions.py

@@ -0,0 +1,89 @@
+import json
+import logging
+import os
+from typing import Any, Dict, List, Optional, Type
+
+from vellum.workflows.nodes.bases.base import BaseNode
+import vellum.workflows.nodes.displayable as displayable_module
+from vellum.workflows.vellum_client import create_vellum_client
+from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+
+logger = logging.getLogger(__name__)
+
+
+def create_display_context_with_client() -> WorkflowDisplayContext:
+    """Create a WorkflowDisplayContext with Vellum client for serialization."""
+    client = create_vellum_client()
+    return WorkflowDisplayContext(client=client, dry_run=True)
+
+
+def get_all_displayable_node_classes() -> List[Type[BaseNode]]:
+    """Get all displayable node classes dynamically from the displayable module."""
+    node_classes = []
+    for class_name in displayable_module.__all__:
+        node_class = getattr(displayable_module, class_name)
+        node_classes.append(node_class)
+    return node_classes
+
+
+def clean_node_definition(definition: Dict[str, Any]) -> Dict[str, Any]:
+    """Remove unwanted fields from a successfully serialized node definition."""
+    fields_to_remove = ["inputs", "data", "type", "adornments", "should_file_merge"]
+    cleaned = {k: v for k, v in definition.items() if k not in fields_to_remove}
+    return cleaned
+
+
+def serialize_node_definition(
+    node_class: Type[BaseNode], display_context: WorkflowDisplayContext
+) -> Optional[Dict[str, Any]]:
+    """Serialize a single node definition, returning None if it fails."""
+    try:
+        display_class = get_node_display_class(node_class)
+        display_instance = display_class()
+        definition = display_instance.serialize(display_context)
+        return clean_node_definition(definition)
+    except Exception as e:
+        logger.info(f"Warning: Failed to serialize {node_class.__name__}: {e}")
+        return None
+
+
+def main() -> None:
+    """Main function to generate node definitions."""
+    logger.info("Generating node definitions...")
+
+    display_context = create_display_context_with_client()
+    node_classes = get_all_displayable_node_classes()
+
+    successful_nodes = []
+    errors = []
+
+    for node_class in node_classes:
+        logger.info(f"Serializing {node_class.__name__}...")
+        definition = serialize_node_definition(node_class, display_context)
+
+        if definition is not None:
+            successful_nodes.append(definition)
+        else:
+            try:
+                display_class = get_node_display_class(node_class)
+                display_instance = display_class()
+                display_instance.serialize(display_context)
+            except Exception as e:
+                errors.append({"node": node_class.__name__, "error": f"{e.__class__.__name__}: {str(e)}"})
+
+    result = {"nodes": successful_nodes, "errors": errors}
+
+    output_path = "ee/vellum_ee/assets/node-definitions.json"
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    with open(output_path, "w") as f:
+        json.dump(result, f, indent=2)
+
+    logger.info(
+        f"Generated {len(successful_nodes)} successful node definitions and {len(errors)} errors in {output_path}"
+    )
+
+
+if __name__ == "__main__":
+    main()

vellum_ee/workflows/display/base.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
+from enum import Enum
 from uuid import UUID
-from typing import Optional, Type
+from typing import TYPE_CHECKING, Dict, Optional, Type
 
 from pydantic import Field
 
@@ -9,6 +10,32 @@ from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
 
+if TYPE_CHECKING:
+    from vellum.workflows.triggers.base import BaseTrigger
+
+
+class WorkflowTriggerType(Enum):
+    MANUAL = "MANUAL"
+    INTEGRATION = "INTEGRATION"
+    SCHEDULED = "SCHEDULED"
+
+
+def get_trigger_type_mapping() -> Dict[Type["BaseTrigger"], WorkflowTriggerType]:
+    """
+    Get mapping from trigger classes to their WorkflowTriggerType enum values.
+
+    This is a function to avoid circular imports - it imports trigger classes
+    only when called, not at module load time.
+
+    Returns:
+        Dictionary mapping trigger classes to their enum types
+    """
+    from vellum.workflows.triggers.manual import ManualTrigger
+
+    return {
+        ManualTrigger: WorkflowTriggerType.MANUAL,
+    }
+
 
 class WorkflowDisplayDataViewport(UniversalBaseModel):
     x: float = 0.0
@@ -39,15 +66,6 @@ class WorkflowMetaDisplay:
         )
 
 
-@dataclass
-class WorkflowMetaDisplayOverrides(WorkflowMetaDisplay):
-    """
-    DEPRECATED: Use WorkflowMetaDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
 @dataclass
 class WorkflowInputsDisplay:
     id: UUID
@@ -55,15 +73,6 @@ class WorkflowInputsDisplay:
     color: Optional[str] = None
 
 
-@dataclass
-class WorkflowInputsDisplayOverrides(WorkflowInputsDisplay):
-    """
-    DEPRECATED: Use WorkflowInputsDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
 @dataclass
 class StateValueDisplay:
     id: UUID
@@ -71,27 +80,10 @@ class StateValueDisplay:
     color: Optional[str] = None
 
 
-@dataclass
-class StateValueDisplayOverrides(StateValueDisplay):
-    """
-    DEPRECATED: Use StateValueDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
 @dataclass
 class EdgeDisplay:
     id: UUID
-
-
-@dataclass
-class EdgeDisplayOverrides(EdgeDisplay):
-    """
-    DEPRECATED: Use EdgeDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
+    z_index: Optional[int] = None
 
 
 @dataclass
@@ -100,25 +92,7 @@ class EntrypointDisplay:
     edge_display: EdgeDisplay
 
 
-@dataclass
-class EntrypointDisplayOverrides(EntrypointDisplay):
-    """
-    DEPRECATED: Use EntrypointDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass
-
-
 @dataclass
 class WorkflowOutputDisplay:
     id: UUID
     name: str
-
-
-@dataclass
-class WorkflowOutputDisplayOverrides(WorkflowOutputDisplay):
-    """
-    DEPRECATED: Use WorkflowOutputDisplay instead. Will be removed in 0.15.0
-    """
-
-    pass

vellum_ee/workflows/display/editor/types.py

@@ -17,6 +17,9 @@ class NodeDisplayComment(UniversalBaseModel):
 
 class NodeDisplayData(UniversalBaseModel):
     position: NodeDisplayPosition = Field(default_factory=NodeDisplayPosition)
+    z_index: Optional[int] = None
     width: Optional[int] = None
     height: Optional[int] = None
     comment: Optional[NodeDisplayComment] = None
+    icon: Optional[str] = None
+    color: Optional[str] = None

vellum_ee/workflows/display/exceptions.py

@@ -0,0 +1,3 @@
+from vellum_ee.workflows.display.utils.exceptions import NodeValidationError
+
+__all__ = ["NodeValidationError"]

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -8,6 +8,7 @@ from typing import (
     Dict,
     ForwardRef,
     Generic,
+    List,
     Optional,
     Set,
     Tuple,
@@ -162,7 +163,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                     {
                         "id": id,
                         "name": attribute.name,
-                        "value": serialize_value(display_context, attribute.instance),
+                        "value": serialize_value(node_id, display_context, attribute.instance),
                     }
                 )
             except ValueError as e:
@@ -175,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,
             }
@@ -183,46 +184,19 @@ 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 = self.serialize_ports(display_context)
-
-        outputs: JsonArray = []
-        for output in node.Outputs:
-            type = primitive_type_to_vellum_variable_type(output)
-            value = (
-                serialize_value(display_context, output.instance)
-                if output.instance is not None and output.instance != undefined
-                else None
-            )
-
-            outputs.append(
-                {
-                    "id": str(uuid4_from_hash(f"{node_id}|{output.name}")),
-                    "name": output.name,
-                    "type": type,
-                    "value": value,
-                }
-            )
-
         return {
             "id": str(node_id),
-            "label": node.__qualname__,
+            "label": self.label,
             "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,
+            **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:
@@ -233,7 +207,9 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                         "id": id,
                         "name": port.name,
                         "type": port._condition_type.value,
-                        "expression": (serialize_value(display_context, port._condition) if port._condition else None),
+                        "expression": (
+                            serialize_value(node_id, display_context, port._condition) if port._condition else None
+                        ),
                     }
                 )
             else:
@@ -247,6 +223,104 @@ 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:
+            output_type = primitive_type_to_vellum_variable_type(output)
+            value = (
+                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": output_id,
+                    "name": output.name,
+                    "type": output_type,
+                    "value": value,
+                }
+            )
+
+        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
 
@@ -269,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)
@@ -293,6 +367,9 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         if default_port in port_displays:
             return port_displays[default_port].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.")
@@ -310,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
@@ -322,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:
@@ -387,29 +464,60 @@ 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)
+
+        # Step 2: Get explicit BaseNodeDisplay values
+        explicit_value = self._get_explicit_node_display_attr("display_data", NodeDisplayData)
+
+        # 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
+                )
+
+        # Step 4: Return the merged result
+        if base_kwargs:
+            return NodeDisplayData(**base_kwargs)
 
         return NodeDisplayData()