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

vellum_ee/workflows/display/nodes/tests/test_base_node_display.py

@@ -4,6 +4,7 @@ from uuid import UUID
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references.constant import ConstantValueReference
+from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -17,7 +18,7 @@ def node_with_implicit_properties():
     class MyNodeDisplay(BaseNodeDisplay[MyNode]):
         pass
 
-    expected_id = UUID("ace7f746-4fe6-45c7-8207-fc8a4d0c7f6f")
+    expected_id = UUID("8628fe5e-620d-4935-92f1-a9ed7ed7fd92")
 
     return MyNodeDisplay, expected_id
 
@@ -55,6 +56,7 @@ def test_get_id(node_info):
 def test_serialize_condition__accessor_expression():
     # GIVEN a node with an accessor expression in a Port
     class MyNode(BaseNode):
+
         class Ports(BaseNode.Ports):
             foo = Port.on_if(ConstantValueReference({"hello": "world"})["hello"])
 
@@ -65,7 +67,7 @@ def test_serialize_condition__accessor_expression():
     # THEN the condition should be serialized correctly
     assert data["ports"] == [
         {
-            "id": "7de6ea94-7f6c-475e-8f38-ec8ac317fd19",
+            "id": "c5a19e2d-3476-4d51-ae8f-a00f101dec30",
             "name": "foo",
             "type": "IF",
             "expression": {
@@ -98,5 +100,149 @@ def test_serialize_display_data():
     # THEN the condition should be serialized correctly
     assert data["display_data"] == {
         "position": {"x": 0.0, "y": 0.0},
-        "comment": {"value": "I hope this works"},
+        "comment": {"expanded": True, "value": "I hope this works"},
     }
+
+
+def test_serialize_node_label_with_pascal_case():
+    """
+    Tests that a node with PascalCase name serializes with proper title case label.
+    """
+
+    # GIVEN a node with a PascalCase name that includes common patterns
+    class MyCustomNode(BaseNode):
+        pass
+
+    # WHEN we serialize the node
+    node_display_class = get_node_display_class(MyCustomNode)
+    data = node_display_class().serialize(WorkflowDisplayContext())
+
+    # THEN the label should be converted to proper title case with spaces
+    assert data["label"] == "My Custom Node"
+
+
+def test_serialize_display_data_with_icon_and_color():
+    """
+    Tests that nodes with icon and color serialize display_data correctly.
+    """
+
+    # GIVEN a node with icon and color in display data
+    class MyNode(BaseNode):
+        pass
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        display_data = NodeDisplayData(
+            position=NodeDisplayPosition(x=100, y=200), icon="vellum:icon:star", color="navy"
+        )
+
+    # WHEN we serialize the node
+    data = MyNodeDisplay().serialize(WorkflowDisplayContext())
+
+    # THEN the display_data should include icon and color
+    assert data["display_data"] == {"position": {"x": 100, "y": 200}, "icon": "vellum:icon:star", "color": "navy"}
+
+
+def test_serialize_display_data_with_various_icon_formats():
+    """
+    Tests that different icon formats are serialized correctly.
+    """
+
+    # GIVEN a node with a vellum icon format
+    class MyNode(BaseNode):
+        pass
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        display_data = NodeDisplayData(icon="vellum:icon:home")
+
+    # WHEN we serialize the node
+    data = MyNodeDisplay().serialize(WorkflowDisplayContext())
+
+    # THEN the icon should be preserved as-is
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["icon"] == "vellum:icon:home"
+
+
+def test_serialize_display_data_with_navy_color():
+    """
+    Tests that navy color values are serialized correctly.
+    """
+
+    # GIVEN a node with a navy color
+    class MyNode(BaseNode):
+        pass
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        display_data = NodeDisplayData(color="navy")
+
+    # WHEN we serialize the node
+    data = MyNodeDisplay().serialize(WorkflowDisplayContext())
+
+    # THEN the color should be preserved as-is
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["color"] == "navy"
+
+
+def test_serialize_basenode_with_display_class():
+    """Tests that BaseNode.Display class attributes serialize correctly."""
+
+    class MyNode(BaseNode):
+        class Display:
+            icon = "vellum:icon:gear"
+            color = "purple"
+
+    node_display_class = get_node_display_class(MyNode)
+    data = node_display_class().serialize(WorkflowDisplayContext())
+
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["icon"] == "vellum:icon:gear"
+    assert display_data["color"] == "purple"
+
+
+def test_serialize_explicit_display_data_overrides_display_class():
+    """Tests that BaseNodeDisplay's explicit display_data takes precedence over Display class."""
+
+    class MyNode(BaseNode):
+        class Display:
+            icon = "vellum:icon:check"
+            color = "gold"
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        display_data = NodeDisplayData(
+            position=NodeDisplayPosition(x=100, y=200),
+            icon="vellum:icon:times",
+            color="navy",
+        )
+
+    data = MyNodeDisplay().serialize(WorkflowDisplayContext())
+
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["icon"] == "vellum:icon:times"  # BaseNodeDisplay overrides
+    assert display_data["color"] == "navy"  # BaseNodeDisplay overrides
+    assert display_data["position"] == {"x": 100, "y": 200}
+
+
+def test_serialize_display_class_used_as_fallback():
+    """Tests that Display class attributes are used when BaseNodeDisplay doesn't specify them."""
+
+    class MyNode(BaseNode):
+        class Display:
+            icon = "vellum:icon:star"
+            color = "green"
+
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        display_data = NodeDisplayData(
+            position=NodeDisplayPosition(x=50, y=75),
+            # Not specifying icon or color - should fall back to Display class
+        )
+
+    data = MyNodeDisplay().serialize(WorkflowDisplayContext())
+
+    display_data = data["display_data"]
+    assert isinstance(display_data, dict)
+    assert display_data["icon"] == "vellum:icon:star"  # Falls back to Display class
+    assert display_data["color"] == "green"  # Falls back to Display class
+    assert display_data["position"] == {"x": 50, "y": 75}  # From BaseNodeDisplay

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

@@ -4,6 +4,7 @@ from .code_execution_node import BaseCodeExecutionNodeDisplay
 from .conditional_node import BaseConditionalNodeDisplay
 from .error_node import BaseErrorNodeDisplay
 from .final_output_node import BaseFinalOutputNodeDisplay
+from .function_node import FunctionNodeDisplay
 from .guardrail_node import BaseGuardrailNodeDisplay
 from .inline_prompt_node import BaseInlinePromptNodeDisplay
 from .inline_subworkflow_node import BaseInlineSubworkflowNodeDisplay
@@ -25,6 +26,7 @@ __all__ = [
     "BaseConditionalNodeDisplay",
     "BaseErrorNodeDisplay",
     "BaseFinalOutputNodeDisplay",
+    "FunctionNodeDisplay",
     "BaseGuardrailNodeDisplay",
     "BaseInlinePromptNodeDisplay",
     "BaseInlineSubworkflowNodeDisplay",

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

@@ -1,8 +1,7 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar, cast
+from typing import ClassVar, Dict, Generic, Optional, TypeVar
 
 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_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
@@ -31,8 +30,21 @@ class BaseAPINodeDisplay(BaseNodeDisplay[_APINodeType], Generic[_APINodeType]):
         APINode.authorization_type,
     }
 
+    # TODO: Only serialize attributes not inputs
+    __unserializable_attributes__ = {
+        APINode.data,
+        APINode.url,
+        APINode.method,
+        APINode.json,
+        APINode.headers,
+        APINode.api_key_header_key,
+        APINode.api_key_header_value,
+        APINode.bearer_token_value,
+        APINode.authorization_type,
+    }
+
     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
@@ -165,13 +177,11 @@ class BaseAPINodeDisplay(BaseNodeDisplay[_APINodeType], Generic[_APINodeType]):
         ]
         inputs.extend(additional_header_inputs)
 
-        _, text_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
-        _, json_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.json)]
-        _, status_code_output_display = display_context.global_node_output_displays[
-            cast(OutputReference, node.Outputs.status_code)
-        ]
+        text_output_display = self.get_node_output_display(node.Outputs.text)
+        json_output_display = self.get_node_output_display(node.Outputs.json)
+        status_code_output_display = self.get_node_output_display(node.Outputs.status_code)
 
-        return {
+        serialized_node: JsonObject = {
             "id": str(node_id),
             "type": "API",
             "inputs": [input.dict() for input in inputs],
@@ -200,8 +210,11 @@ class BaseAPINodeDisplay(BaseNodeDisplay[_APINodeType], Generic[_APINodeType]):
                 "json_output_id": str(json_output_display.id),
                 "status_code_output_id": str(status_code_output_display.id),
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "ports": self.serialize_ports(display_context),
+            **self.serialize_generic_fields(display_context),
         }
+
+        attributes = self._serialize_attributes(display_context)
+        if attributes:
+            serialized_node["attributes"] = attributes
+
+        return serialized_node

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

@@ -6,7 +6,7 @@ from typing import Any, Callable, Dict, Generic, Optional, Type, TypeVar, cast
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.nodes.utils import get_wrapped_node
-from vellum.workflows.types.core import JsonArray, JsonObject
+from vellum.workflows.types.core import 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
@@ -94,8 +94,8 @@ class BaseAdornmentNodeDisplay(BaseNodeDisplay[_BaseAdornmentNodeType], Generic[
         additional_kwargs = get_additional_kwargs(wrapped_node_display.node_id) if get_additional_kwargs else {}
         serialized_wrapped_node = wrapped_node_display.serialize(display_context, **kwargs, **additional_kwargs)
 
-        adornments = cast(JsonArray, serialized_wrapped_node.get("adornments")) or []
-        serialized_wrapped_node["adornments"] = adornments + [adornment] if adornment else adornments
+        adornments = cast(list, serialized_wrapped_node.get("adornments")) or []
+        serialized_wrapped_node["adornments"] = [adornment] + adornments if adornment else adornments
 
         return serialized_wrapped_node
 

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

@@ -1,19 +1,36 @@
 import inspect
+import os
 from uuid import UUID
 from typing import ClassVar, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
-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_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
+from vellum_ee.workflows.display.utils.exceptions import NodeValidationError
+from vellum_ee.workflows.display.utils.expressions import virtual_open
 
 _CodeExecutionNodeType = TypeVar("_CodeExecutionNodeType", bound=CodeExecutionNode)
 
 
+def _read_file_from_path_with_virtual_support(node_filepath: str, script_filepath: str) -> Optional[str]:
+    """
+    Read a file using virtual_open which handles VirtualFileFinder instances.
+    """
+    node_filepath_dir = os.path.dirname(node_filepath)
+    normalized_script_filepath = script_filepath.lstrip("./")
+    full_filepath = os.path.join(node_filepath_dir, normalized_script_filepath)
+
+    try:
+        with virtual_open(full_filepath) as file:
+            return file.read()
+    except (FileNotFoundError, IsADirectoryError):
+        return None
+
+
 class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Generic[_CodeExecutionNodeType]):
     output_id: ClassVar[Optional[UUID]] = None
     log_output_id: ClassVar[Optional[UUID]] = None
@@ -24,7 +41,7 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
     }
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
@@ -36,10 +53,12 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
             code_value = raw_code
         elif filepath:
             node_file_path = inspect.getfile(node)
-            file_code = read_file_from_path(
+            file_code = _read_file_from_path_with_virtual_support(
                 node_filepath=node_file_path,
                 script_filepath=filepath,
             )
+            if not file_code:
+                display_context.add_error(NodeValidationError(f"Filepath '{filepath}' does not exist", node.__name__))
             code_value = file_code
         else:
             code_value = ""
@@ -100,8 +119,5 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
                 "output_id": str(self.output_id) if self.output_id else str(output_display.id),
                 "log_output_id": str(self.log_output_id) if self.log_output_id else str(log_output_display.id),
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "ports": self.serialize_ports(display_context),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }

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, Tuple, TypeVar, Union
+from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, TypeVar, Union, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.and_ import AndExpression
@@ -45,7 +45,7 @@ class BaseConditionalNodeDisplay(BaseNodeDisplay[_ConditionalNodeType], Generic[
     rule_ids: ClassVar[List[RuleIdMap]]
     condition_ids: ClassVar[list[ConditionId]]
 
-    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+    def serialize(self, display_context: WorkflowDisplayContext, **_kwargs) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
@@ -61,6 +61,9 @@ class BaseConditionalNodeDisplay(BaseNodeDisplay[_ConditionalNodeType], Generic[
 but the defined conditions have length {len(condition_ids)}"""
             )
 
+        ports = self.serialize_ports(display_context)
+        ports_by_index = {idx: port for idx, port in enumerate(ports)}
+
         def serialize_rule(
             descriptor: BaseDescriptor, path: List[int], rule_id_map: Optional[RuleIdMap]
         ) -> Dict[str, Any]:
@@ -166,7 +169,7 @@ but the defined conditions have length {len(condition_ids)}"""
             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}"))
+            source_handle_id = self._get_source_handle_id(idx, source_handle_ids, ports_by_index)
 
             if port._condition is None:
                 if port._condition_type == ConditionType.ELSE:
@@ -214,10 +217,7 @@ but the defined conditions have length {len(condition_ids)}"""
                 "conditions": conditions,  # type: ignore
                 "version": "2",
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "ports": self.serialize_ports(display_context),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }
 
     def get_nested_rule_details_by_path(
@@ -258,3 +258,14 @@ but the defined conditions have length {len(condition_ids)}"""
 
     def _get_condition_ids(self) -> List[ConditionId]:
         return self._get_explicit_node_display_attr("condition_ids", List[ConditionId]) or []
+
+    def _get_source_handle_id(
+        self, idx: int, source_handle_ids: Dict[int, UUID], ports_by_index: Dict[int, Any]
+    ) -> str:
+        if source_handle_ids.get(idx):
+            return str(source_handle_ids[idx])
+
+        if idx in ports_by_index:
+            return cast(str, ports_by_index[idx]["id"])
+
+        return str(uuid4_from_hash(f"{self.node_id}|source_handle|{idx}"))

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

@@ -20,7 +20,7 @@ class BaseErrorNodeDisplay(BaseNodeDisplay[_ErrorNodeType], Generic[_ErrorNodeTy
 
     __serializable_inputs__ = {ErrorNode.error}
 
-    def serialize(self, display_context: WorkflowDisplayContext, **kwargs) -> JsonObject:
+    def serialize(self, display_context: WorkflowDisplayContext, **_kwargs) -> JsonObject:
         node_id = self.node_id
         error_source_input_id = self.node_input_ids_by_name.get(
             ErrorNode.error.name,
@@ -47,9 +47,7 @@ class BaseErrorNodeDisplay(BaseNodeDisplay[_ErrorNodeType], Generic[_ErrorNodeTy
                 "target_handle_id": str(self.get_target_handle_id()),
                 "error_source_input_id": str(error_source_input_id),
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }
 
         if self.name:

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

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, ClassVar, Generic, Optional, TypeVar
+from typing import ClassVar, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.types.core import JsonObject
@@ -8,7 +8,6 @@ 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
-from vellum_ee.workflows.display.utils.expressions import serialize_value
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 
 _FinalOutputNodeType = TypeVar("_FinalOutputNodeType", bound=FinalOutputNode)
@@ -19,7 +18,7 @@ NODE_INPUT_KEY = "node_input"
 class BaseFinalOutputNodeDisplay(BaseNodeDisplay[_FinalOutputNodeType], Generic[_FinalOutputNodeType]):
     output_name: ClassVar[Optional[str]] = None
 
-    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+    def serialize(self, display_context: WorkflowDisplayContext, **_kwargs) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
@@ -45,17 +44,7 @@ class BaseFinalOutputNodeDisplay(BaseNodeDisplay[_FinalOutputNodeType], Generic[
                 "node_input_id": str(node_input.id),
             },
             "inputs": [node_input.dict()],
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "outputs": [
-                {
-                    "id": str(self._get_output_id()),
-                    "name": node.Outputs.value.name,
-                    "type": inferred_type,
-                    "value": serialize_value(display_context, node.Outputs.value.instance),
-                }
-            ],
+            **self.serialize_generic_fields(display_context),
         }
 
     def _get_output_id(self) -> UUID:

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

@@ -0,0 +1,14 @@
+from typing import Generic, TypeVar
+
+from vellum.workflows.nodes.displayable.tool_calling_node.utils import FunctionNode
+from vellum_ee.workflows.display.editor.types import NodeDisplayData
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+
+_FunctionNodeType = TypeVar("_FunctionNodeType", bound=FunctionNode)
+
+
+class FunctionNodeDisplay(BaseNodeDisplay[_FunctionNodeType], Generic[_FunctionNodeType]):
+    display_data = NodeDisplayData(
+        icon="vellum:icon:rectangle-code",
+        color="purple",
+    )

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

@@ -15,7 +15,7 @@ class BaseGuardrailNodeDisplay(BaseNodeDisplay[_GuardrailNodeType], Generic[_Gua
     __serializable_inputs__ = {GuardrailNode.metric_inputs}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **_kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
@@ -45,8 +45,5 @@ class BaseGuardrailNodeDisplay(BaseNodeDisplay[_GuardrailNodeType], Generic[_Gua
                 "metric_definition_id": str(raise_if_descriptor(node.metric_definition)),
                 "release_tag": raise_if_descriptor(node.release_tag),
             },
-            "display_data": self.get_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "ports": self.serialize_ports(display_context),
+            **self.serialize_generic_fields(display_context, exclude=["outputs"]),
         }