vellum-ai 1.3.8__py3-none-any.whl → 1.3.9__py3-none-any.whl

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -162,7 +162,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:
@@ -187,7 +187,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         for output in node.Outputs:
             type = primitive_type_to_vellum_variable_type(output)
             value = (
-                serialize_value(display_context, output.instance)
+                serialize_value(node_id, display_context, output.instance)
                 if output.instance is not None and output.instance != undefined
                 else None
             )
@@ -214,6 +214,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     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:
@@ -224,7 +225,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:
@@ -248,21 +251,23 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
     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 self._node:
+        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"{self.node_id}|{attribute.name}"))
+                else str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
             )
             try:
                 attribute_dict: JsonObject = {
                     "id": id,
                     "name": attribute.name,
-                    "value": serialize_value(display_context, attribute.instance),
+                    "value": serialize_value(node_id, display_context, attribute.instance),
                 }
                 attributes.append(attribute_dict)
             except ValueError as e:
@@ -272,7 +277,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
     def serialize_generic_fields(self, display_context: "WorkflowDisplayContext") -> JsonObject:
         """Serialize generic fields that are common to all nodes."""
-        return {
+        result: JsonObject = {
             "display_data": self.get_display_data().dict(),
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
@@ -280,6 +285,22 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             "ports": self.serialize_ports(display_context),
         }
 
+        # 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
+
+        return result
+
     def get_base(self) -> CodeResourceDefinition:
         node = self._node
 
@@ -429,13 +450,20 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
                 if explicit_value.comment.expanded is not None
                 else NodeDisplayComment(value=docstring, expanded=True)
             )
-            return NodeDisplayData(
-                position=explicit_value.position,
-                z_index=explicit_value.z_index,
-                width=explicit_value.width,
-                height=explicit_value.height,
-                comment=comment,
-            )
+            from typing import Any, Dict
+
+            kwargs: Dict[str, Any] = {
+                "position": explicit_value.position,
+                "z_index": explicit_value.z_index,
+                "width": explicit_value.width,
+                "height": explicit_value.height,
+                "comment": comment,
+            }
+            if explicit_value.icon is not None:
+                kwargs["icon"] = explicit_value.icon
+            if explicit_value.color is not None:
+                kwargs["color"] = explicit_value.color
+            return NodeDisplayData(**kwargs)
 
         if explicit_value:
             return explicit_value

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
@@ -117,3 +118,66 @@ def test_serialize_node_label_with_pascal_case():
 
     # 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"

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

@@ -51,7 +51,7 @@ class BaseFinalOutputNodeDisplay(BaseNodeDisplay[_FinalOutputNodeType], Generic[
                     "id": str(self._get_output_id()),
                     "name": node.Outputs.value.name,
                     "type": inferred_type,
-                    "value": serialize_value(display_context, node.Outputs.value.instance),
+                    "value": serialize_value(node_id, display_context, node.Outputs.value.instance),
                 }
             ],
         }

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

@@ -33,7 +33,7 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
                 {
                     "id": id,
                     "name": attribute.name,
-                    "value": serialize_value(display_context, attribute.instance),
+                    "value": serialize_value(node_id, display_context, attribute.instance),
                 }
             )
 

vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py

@@ -197,7 +197,7 @@ def test_serialize_node__prompt_inputs__state_reference():
             "type": "DICTIONARY_REFERENCE",
             "entries": [
                 {
-                    "id": "feadab18-07c6-4790-adf0-28be20a5447a",
+                    "id": "52559b9e-4e8e-438a-8246-cfa30c98d5d1",
                     "key": "foo",
                     "value": {
                         "type": "WORKFLOW_STATE",
@@ -205,7 +205,7 @@ def test_serialize_node__prompt_inputs__state_reference():
                     },
                 },
                 {
-                    "id": "4f252277-d0a5-48ad-8cc5-be03bf0dbd19",
+                    "id": "3750feb9-5d5c-4150-b62d-a9924f466888",
                     "key": "bar",
                     "value": {
                         "type": "CONSTANT_VALUE",
@@ -325,52 +325,52 @@ def test_serialize_node__prompt_parameters__dynamic_references():
     assert parameters_attribute["value"]["type"] == "DICTIONARY_REFERENCE"
     assert parameters_attribute["value"]["entries"] == [
         {
-            "id": "6b63ff96-a2eb-4c6e-bad1-bde01605fa86",
+            "id": "24703d3a-ee6c-4b1b-80f8-6c19ef16723a",
             "key": "stop",
             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
         },
         {
-            "id": "265a1c17-2089-4ac1-b2ce-361b6b9a3335",
+            "id": "88a3bf5d-f42b-4895-850e-ad843945a003",
             "key": "temperature",
             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
         },
         {
-            "id": "699976ec-8ec2-476a-a011-7cf810a8a307",
+            "id": "ede3e0c2-3033-4d0a-bd72-e52595bdc916",
             "key": "max_tokens",
             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
         },
         {
-            "id": "a87e23da-9794-41ff-ba80-c3a77e976e75",
+            "id": "0013cd8f-7658-4908-80fc-b8995d8ca4cc",
             "key": "top_p",
             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
         },
         {
-            "id": "18eb53c2-ec1a-4115-9f21-083af430df67",
+            "id": "98eb2e57-d4ec-4c27-b39b-0b8086918a0f",
             "key": "top_k",
             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
         },
         {
-            "id": "295509a2-5837-452c-893d-f47b67c63c8a",
+            "id": "04accc66-888c-4145-8b4f-d8ff99e38172",
             "key": "frequency_penalty",
             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
         },
         {
-            "id": "5fc64379-5566-426a-a909-dd56c3305aa5",
+            "id": "9236d564-0637-48de-8423-cdf3617dd6b4",
             "key": "presence_penalty",
             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
         },
         {
-            "id": "5d326da0-c096-4425-8bf1-3a18764e96e3",
+            "id": "74f3e80a-3935-45af-a9b3-d49e310a4c03",
             "key": "logit_bias",
             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
         },
         {
-            "id": "cd1a0e1b-6667-48a0-9964-257e1ec8851d",
+            "id": "69a7ebf7-d21a-44e9-a0fa-43eb9b2815df",
             "key": "custom_parameters",
             "value": {
                 "entries": [
                     {
-                        "id": "a9a3092e-dd18-4533-b6b5-24588ebd8f7f",
+                        "id": "e709dc4d-f2db-4dc9-b912-401b52fbb7b4",
                         "key": "json_schema",
                         "value": {
                             "input_variable_id": "c02d1201-86d1-4364-b3b3-4fc6824db8a4",

vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py

@@ -92,7 +92,7 @@ def test_serialize_node__prompt_inputs__input_reference():
             "type": "DICTIONARY_REFERENCE",
             "entries": [
                 {
-                    "id": "ab7902ef-de14-4edc-835c-366d3ef6a70e",
+                    "id": "845009c8-03f8-4de4-b956-841309457d37",
                     "key": "foo",
                     "value": {"type": "WORKFLOW_INPUT", "input_variable_id": "e3657390-fd3c-4fea-8cdd-fc5ea79f3278"},
                 }
@@ -134,12 +134,12 @@ def test_serialize_node__prompt_inputs__mixed_values():
             "type": "DICTIONARY_REFERENCE",
             "entries": [
                 {
-                    "id": "0fc7e25e-075c-4849-b89b-9729d1aeada1",
+                    "id": "a4016385-3cab-4c01-b9d2-7865cd54bdb0",
                     "key": "foo",
                     "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "bar"}},
                 },
                 {
-                    "id": "bba42c89-fa7b-4cb7-bc16-0d21ce060a4b",
+                    "id": "828928b1-24e3-4457-9d6f-4f0692dfa355",
                     "key": "baz",
                     "value": {"type": "WORKFLOW_INPUT", "input_variable_id": "8d57cf1d-147c-427b-9a5e-e5f6ab76e2eb"},
                 },
@@ -540,7 +540,7 @@ def test_serialize_tool_prompt_node_with_inline_workflow():
             "type": "DICTIONARY_REFERENCE",
             "entries": [
                 {
-                    "id": "76ceec7b-ec37-474f-ba38-2bfd27cecc5d",
+                    "id": "b1dfaf2b-b9fb-4fea-ad04-a988e5223d06",
                     "key": "chat_history",
                     "value": {
                         "type": "BINARY_EXPRESSION",

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

@@ -37,7 +37,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
                 {
                     "id": id,
                     "name": attribute.name,
-                    "value": serialize_value(display_context, attribute.instance),
+                    "value": serialize_value(node_id, display_context, attribute.instance),
                 }
             )
 

vellum_ee/workflows/display/tests/test_base_workflow_display.py

@@ -476,3 +476,49 @@ def test_serialize_workflow_with_node_display_data():
 
     assert test_node is not None, "TestNode not found in serialized nodes"
     assert test_node["display_data"] == {"position": {"x": 100, "y": 200}, "z_index": 10, "width": 300, "height": 150}
+
+
+def test_serialize_workflow_with_node_icon_and_color():
+    """
+    Tests that nodes with icon and color serialize correctly in workflow context.
+    """
+
+    # GIVEN a workflow with a node that has icon and color
+    class TestNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    class TestWorkflow(BaseWorkflow):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_result = TestNode.Outputs.result
+
+    class TestNodeDisplay(BaseNodeDisplay[TestNode]):
+        display_data = NodeDisplayData(position=NodeDisplayPosition(x=100, y=200), icon="vellum:icon:cog", color="navy")
+
+    class TestWorkflowDisplay(BaseWorkflowDisplay[TestWorkflow]):
+        pass
+
+    # WHEN we serialize the workflow
+    display = get_workflow_display(
+        base_display_class=TestWorkflowDisplay,
+        workflow_class=TestWorkflow,
+    )
+    serialized_workflow = display.serialize()
+
+    # THEN the node should include icon and color in display_data
+    workflow_raw_data = cast(Dict[str, Any], serialized_workflow["workflow_raw_data"])
+    nodes = cast(List[Dict[str, Any]], workflow_raw_data["nodes"])
+
+    test_node = None
+    for node in nodes:
+        if node.get("type") == "GENERIC":
+            definition = node.get("definition")
+            if isinstance(definition, dict) and definition.get("name") == "TestNode":
+                test_node = node
+                break
+
+    assert test_node is not None, "TestNode not found in serialized nodes"
+    assert test_node["display_data"]["icon"] == "vellum:icon:cog"
+    assert test_node["display_data"]["color"] == "navy"

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -293,7 +293,7 @@ def test_serialize_node__workflow_input_as_nested_chat_history():
                         "type": "DICTIONARY_REFERENCE",
                         "entries": [
                             {
-                                "id": "52d98e78-5c2a-488e-be3e-f3b487d94ad3",
+                                "id": "07513ab1-cf47-490e-8b43-5da226332a00",
                                 "key": "hello",
                                 "value": {
                                     "type": "WORKFLOW_INPUT",

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

@@ -219,7 +219,7 @@ def test_serialize_workflow():
                         "type": "DICTIONARY_REFERENCE",
                         "entries": [
                             {
-                                "id": "0bfa70a2-164f-460c-9e9a-4d62221eadf4",
+                                "id": "6eb6687c-f894-4398-8e62-7dc89e96a0a4",
                                 "key": "noun",
                                 "value": {
                                     "type": "WORKFLOW_INPUT",
@@ -521,37 +521,37 @@ def test_serialize_workflow_with_nested_descriptor_blocks():
         {
             "entries": [
                 {
-                    "id": "24a203be-3cba-4b20-bc84-9993a476c120",
+                    "id": "4e61fbcf-13b3-4d5f-b5fb-2bf919a92045",
                     "key": "block_type",
                     "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "CHAT_MESSAGE"}},
                 },
                 {
-                    "id": "c06269e6-f74c-4860-8fa5-22dcbdc89399",
+                    "id": "79dd757e-46db-4c36-9ffc-ddb763d14f27",
                     "key": "state",
                     "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
                 },
                 {
-                    "id": "dd9c0d43-b931-4dc8-8b3a-a7507ddff0c1",
+                    "id": "2f8164e8-5495-4b9c-8268-d75618cd0842",
                     "key": "cache_config",
                     "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
                 },
                 {
-                    "id": "bef22f2b-0b6e-4910-88cc-6df736d2e20e",
+                    "id": "0e8dc132-de9a-40dc-9845-336bc957df5a",
                     "key": "chat_role",
                     "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "SYSTEM"}},
                 },
                 {
-                    "id": "c0beec30-f85e-4a78-a3fb-baee54a692f8",
+                    "id": "755a45d2-2420-4414-b318-5790880f84ec",
                     "key": "chat_source",
                     "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
                 },
                 {
-                    "id": "f601f4f2-62fe-4697-9fe0-99ca8aa64500",
+                    "id": "3a563cdb-d130-497f-bac6-c324a4349a3c",
                     "key": "chat_message_unterminated",
                     "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
                 },
                 {
-                    "id": "ad550008-64e3-44a3-a32a-84ec226db31c",
+                    "id": "2d0c084e-c54f-48f5-9444-a17f8aeb8f76",
                     "key": "blocks",
                     "value": {
                         "items": [

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

@@ -131,6 +131,7 @@ def test_serialize_workflow():
                             "id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
                             "label": "Start Node",
                             "type": "GENERIC",
+                            "should_file_merge": True,
                             "display_data": {"position": {"x": 200.0, "y": -50.0}},
                             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
                             "definition": {

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

@@ -118,6 +118,7 @@ def test_serialize_workflow():
                             "id": "baf6d316-dc75-41e8-96c0-015aede96309",
                             "label": "Iteration",
                             "type": "GENERIC",
+                            "should_file_merge": True,
                             "display_data": {"position": {"x": 200.0, "y": -50.0}},
                             "base": {
                                 "name": "BaseNode",

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

@@ -155,6 +155,7 @@ def test_serialize_workflow():
                                                 "id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
                                                 "label": "Start Node",
                                                 "type": "GENERIC",
+                                                "should_file_merge": True,
                                                 "display_data": {"position": {"x": 200.0, "y": -50.0}},
                                                 "base": {
                                                     "name": "BaseNode",
@@ -400,7 +401,7 @@ def test_serialize_workflow():
                     "type": "DICTIONARY_REFERENCE",
                     "entries": [
                         {
-                            "id": "d5c4d578-6ef1-4786-88f6-1ab0892d0798",
+                            "id": "8eb8b551-9b48-43b3-861f-52adb5c585a8",
                             "key": "question",
                             "value": {
                                 "type": "WORKFLOW_INPUT",

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

@@ -42,6 +42,7 @@ def test_serialize_workflow():
         "id": "21f29cac-da87-495f-bba1-093d423f4e46",
         "label": "Get Current Weather Node",
         "type": "GENERIC",
+        "should_file_merge": True,
         "display_data": {
             "position": {"x": 200.0, "y": -50.0},
             "comment": {
@@ -169,7 +170,7 @@ def test_serialize_workflow():
                     "type": "DICTIONARY_REFERENCE",
                     "entries": [
                         {
-                            "id": "b6d4427d-16dd-478a-9780-f88d60d2263d",
+                            "id": "8eb8b551-9b48-43b3-861f-52adb5c585a8",
                             "key": "question",
                             "value": {
                                 "type": "WORKFLOW_INPUT",

vellum_ee/workflows/display/utils/events.py

@@ -1,3 +1,6 @@
+from typing import Optional
+
+from vellum import Vellum
 from vellum.workflows.events.workflow import WorkflowExecutionInitiatedEvent
 from vellum_ee.workflows.display.utils.registry import (
     get_parent_display_context_from_event,
@@ -25,7 +28,9 @@ def _should_mark_workflow_dynamic(event: WorkflowExecutionInitiatedEvent) -> boo
     return True
 
 
-def event_enricher(event: WorkflowExecutionInitiatedEvent) -> WorkflowExecutionInitiatedEvent:
+def event_enricher(
+    event: WorkflowExecutionInitiatedEvent, client: Optional[Vellum] = None
+) -> WorkflowExecutionInitiatedEvent:
     if event.name != "workflow.execution.initiated":
         return event
 
@@ -33,6 +38,7 @@ def event_enricher(event: WorkflowExecutionInitiatedEvent) -> WorkflowExecutionI
     workflow_display = get_workflow_display(
         workflow_class=workflow_definition,
         parent_display_context=get_parent_display_context_from_event(event),
+        client=client,
         dry_run=True,
     )
     register_workflow_display_context(event.span_id, workflow_display.display_context)