vellum-ai 0.14.14__py3-none-any.whl → 0.14.16__py3-none-any.whl

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

@@ -9,16 +9,15 @@ class Inputs(BaseInputs):
     input: str
 
 
-class BasicGenericNode(BaseNode):
-    pass
-
-
 def test_serialize_node__basic(serialize_node):
+    class BasicGenericNode(BaseNode):
+        pass
+
     serialized_node = serialize_node(BasicGenericNode)
     assert not DeepDiff(
         {
-            "id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
-            "label": "BasicGenericNode",
+            "id": "8d7cbfe4-72ca-4367-a401-8d28723d2f00",
+            "label": "test_serialize_node__basic.<locals>.BasicGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -34,10 +33,10 @@ def test_serialize_node__basic(serialize_node):
                     "test_trigger_serialization",
                 ],
             },
-            "trigger": {"id": "9d3a1b3d-4a38-4f2e-bbf1-dd8be152bce8", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "trigger": {"id": "be19c63b-3492-46b1-be9d-16f8d2e6410b", "merge_behavior": "AWAIT_ATTRIBUTES"},
             "ports": [
                 {
-                    "id": "89dccfa5-cc1a-4612-bd87-86cb444f6dd4",
+                    "id": "8bec8d0c-113f-4110-afcb-4a6e566e7236",
                     "name": "default",
                     "type": "DEFAULT",
                 }
@@ -51,17 +50,16 @@ def test_serialize_node__basic(serialize_node):
     )
 
 
-class AwaitAnyGenericNode(BaseNode):
-    class Trigger(BaseNode.Trigger):
-        merge_behavior = MergeBehavior.AWAIT_ANY
-
-
 def test_serialize_node__await_any(serialize_node):
+    class AwaitAnyGenericNode(BaseNode):
+        class Trigger(BaseNode.Trigger):
+            merge_behavior = MergeBehavior.AWAIT_ANY
+
     serialized_node = serialize_node(AwaitAnyGenericNode)
     assert not DeepDiff(
         {
-            "id": "0ba67f76-aaff-4bd4-a20f-73a32ef5810d",
-            "label": "AwaitAnyGenericNode",
+            "id": "42e17f0e-8496-415f-9c72-f85250ba6f0b",
+            "label": "test_serialize_node__await_any.<locals>.AwaitAnyGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -77,10 +75,10 @@ def test_serialize_node__await_any(serialize_node):
                     "test_trigger_serialization",
                 ],
             },
-            "trigger": {"id": "ffa72187-9a18-453f-ae55-b77aad332630", "merge_behavior": "AWAIT_ANY"},
+            "trigger": {"id": "5bb6bb4c-4374-44c8-a7b5-7bb6c1060a5b", "merge_behavior": "AWAIT_ANY"},
             "ports": [
                 {
-                    "id": "38b83138-cb07-40cb-82d2-83982ded6883",
+                    "id": "d9a84db7-8bd6-4a15-9e3c-c2e898c26d16",
                     "name": "default",
                     "type": "DEFAULT",
                 }
@@ -94,17 +92,16 @@ def test_serialize_node__await_any(serialize_node):
     )
 
 
-class AwaitAllGenericNode(BaseNode):
-    class Trigger(BaseNode.Trigger):
-        merge_behavior = MergeBehavior.AWAIT_ALL
-
-
 def test_serialize_node__await_all(serialize_node):
+    class AwaitAllGenericNode(BaseNode):
+        class Trigger(BaseNode.Trigger):
+            merge_behavior = MergeBehavior.AWAIT_ALL
+
     serialized_node = serialize_node(AwaitAllGenericNode)
     assert not DeepDiff(
         {
-            "id": "09d06cd3-06ea-40cc-afd8-17ad88542271",
-            "label": "AwaitAllGenericNode",
+            "id": "b3e1145a-5f41-456b-9382-6d0a1e828c2f",
+            "label": "test_serialize_node__await_all.<locals>.AwaitAllGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -120,10 +117,10 @@ def test_serialize_node__await_all(serialize_node):
                     "test_trigger_serialization",
                 ],
             },
-            "trigger": {"id": "62074276-c817-476d-b59d-da523ae3f218", "merge_behavior": "AWAIT_ALL"},
+            "trigger": {"id": "124ba9cf-a30e-41ef-81bf-143708f8b1c3", "merge_behavior": "AWAIT_ALL"},
             "ports": [
                 {
-                    "id": "9edef179-a1c2-4624-b185-593bb84f08a0",
+                    "id": "fa73da35-0bf9-4f02-bf5b-0b0d1a6f1494",
                     "name": "default",
                     "type": "DEFAULT",
                 }

vellum_ee/workflows/display/utils/vellum.py

@@ -28,6 +28,7 @@ from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualTo
 from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.expressions.not_in import NotInExpression
 from vellum.workflows.expressions.or_ import OrExpression
+from vellum.workflows.expressions.parse_json import ParseJsonExpression
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.references import OutputReference, WorkflowInputReference
@@ -163,5 +164,7 @@ def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperato
         return "or"
     elif isinstance(descriptor, CoalesceExpression):
         return "coalesce"
+    elif isinstance(descriptor, ParseJsonExpression):
+        return "parseJson"
     else:
         raise ValueError(f"Unsupported descriptor type: {descriptor}")

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -212,8 +212,6 @@ class BaseWorkflowDisplay(
 
         port_displays: Dict[Port, PortDisplay] = {}
 
-        # TODO: We should still serialize nodes that are in the workflow's directory but aren't used in the graph.
-        # https://app.shortcut.com/vellum/story/5394
         for node in self._workflow.get_nodes():
             extracted_node_displays = self._extract_node_displays(node)
 
@@ -403,11 +401,11 @@ class BaseWorkflowDisplay(
             for input in display_context.workflow_input_displays
         }
         node_displays = {
-            str(node.__id__): display_context.node_displays[node] for node in display_context.node_displays
+            node.__id__: (node, display_context.node_displays[node]) for node in display_context.node_displays
         }
         node_event_displays = {}
         for node_id in node_displays:
-            current_node_display = node_displays[node_id]
+            node, current_node_display = node_displays[node_id]
             input_display = {}
             if isinstance(current_node_display, BaseNodeVellumDisplay):
                 input_display = current_node_display.node_input_ids_by_name
@@ -418,7 +416,6 @@ class BaseWorkflowDisplay(
             port_display_meta = {
                 port.name: current_node_display.port_displays[port].id for port in current_node_display.port_displays
             }
-            node = current_node_display._node
             subworkflow_display_context: Optional[WorkflowEventDisplayContext] = None
             if hasattr(node, "subworkflow"):
                 # All nodes that have a subworkflow attribute are currently expected to call them `subworkflow`
@@ -432,7 +429,7 @@ class BaseWorkflowDisplay(
                     )
                     subworkflow_display_context = subworkflow_display.get_event_display_context()
 
-            node_event_displays[node_id] = NodeEventDisplayContext(
+            node_event_displays[str(node_id)] = NodeEventDisplayContext(
                 input_display=input_display,
                 output_display=output_display,
                 port_display=port_display_meta,

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -1,9 +1,13 @@
 import pytest
+from uuid import uuid4
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.inline_subworkflow_node.node import InlineSubworkflowNode
+from vellum.workflows.nodes.core.retry_node.node import RetryNode
+from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.vellum.retry_node import BaseRetryNodeDisplay
 from vellum_ee.workflows.display.vellum import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -139,3 +143,51 @@ def test_get_event_display_context__node_display_to_include_subworkflow_display(
 
     assert node_event_display.subworkflow_display is not None
     assert str(InnerNode.__id__) in node_event_display.subworkflow_display.node_displays
+
+
+def test_get_event_display_context__node_display_for_adornment_nodes():
+    # GIVEN a simple workflow with a retry node adornment
+    @RetryNode.wrap(max_attempts=4)
+    class MyNode(BaseNode):
+        pass
+
+    class MyWorkflow(BaseWorkflow):
+        graph = MyNode
+
+    # AND a display class for the node
+    inner_node_id = uuid4()
+
+    @BaseRetryNodeDisplay.wrap()
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        node_id = inner_node_id
+
+    # WHEN we gather the event display context
+    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+
+    # THEN the subworkflow display should be included
+    assert str(MyNode.__id__) in display_context.node_displays
+    node_event_display = display_context.node_displays[str(MyNode.__id__)]
+    assert node_event_display.subworkflow_display is not None
+    assert str(inner_node_id) in node_event_display.subworkflow_display.node_displays
+
+
+def test_get_event_display_context__templating_node_input_display():
+    # GIVEN a simple workflow with a templating node referencing another node output
+    class DataNode(BaseNode):
+        class Outputs:
+            bar: str
+
+    class MyNode(TemplatingNode):
+        inputs = {"foo": DataNode.Outputs.bar}
+
+    class MyWorkflow(BaseWorkflow):
+        graph = DataNode >> MyNode
+
+    # WHEN we gather the event display context
+    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+
+    # THEN the subworkflow display should be included
+    assert str(MyNode.__id__) in display_context.node_displays
+    node_event_display = display_context.node_displays[str(MyNode.__id__)]
+
+    assert node_event_display.input_display.keys() == {"inputs.foo"}