vellum-ai 0.14.15__py3-none-any.whl → 0.14.17__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/workflows/tests/test_workflow_display.py

@@ -5,9 +5,11 @@ 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.nodes.core.try_node.node import TryNode
 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.nodes.vellum.try_node import BaseTryNodeDisplay
 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
@@ -119,6 +121,29 @@ def test_get_event_display_context__node_display_filled_without_base_display():
     assert StartNode.__output_ids__ == node_event_display.output_display
 
 
+def test_get_event_display_context__node_display_filled_without_output_display():
+    # GIVEN a simple workflow
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+    class StartNodeDisplay(BaseNodeDisplay[StartNode]):
+        pass
+
+    # WHEN we gather the event display context
+    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+
+    # THEN the node display should be included
+    assert str(StartNode.__id__) in display_context.node_displays
+    node_event_display = display_context.node_displays[str(StartNode.__id__)]
+
+    # AND so should their output ids
+    assert node_event_display.output_display.keys() == {"foo"}
+
+
 def test_get_event_display_context__node_display_to_include_subworkflow_display():
     # GIVEN a simple workflow
     class InnerNode(BaseNode):
@@ -145,19 +170,33 @@ def test_get_event_display_context__node_display_to_include_subworkflow_display(
     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)
+@pytest.mark.parametrize(
+    ["AdornmentNode", "AdornmentNodeDisplay", "expected_adornment_output_names"],
+    [
+        [RetryNode, BaseRetryNodeDisplay, {"foo"}],
+        [TryNode, BaseTryNodeDisplay, {"foo", "error"}],
+    ],
+    ids=["retry_node", "try_node"],
+)
+def test_get_event_display_context__node_display_for_adornment_nodes(
+    AdornmentNode,
+    AdornmentNodeDisplay,
+    expected_adornment_output_names,
+):
+    # GIVEN a simple workflow with an adornment
+    @AdornmentNode.wrap()
     class MyNode(BaseNode):
-        pass
+        class Outputs(BaseNode.Outputs):
+            foo: str
 
     class MyWorkflow(BaseWorkflow):
         graph = MyNode
 
     # AND a display class for the node
+    adornment_node_id = uuid4()
     inner_node_id = uuid4()
 
-    @BaseRetryNodeDisplay.wrap()
+    @AdornmentNodeDisplay.wrap(node_id=adornment_node_id)
     class MyNodeDisplay(BaseNodeDisplay[MyNode]):
         node_id = inner_node_id
 
@@ -165,11 +204,16 @@ def test_get_event_display_context__node_display_for_adornment_nodes():
     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 str(adornment_node_id) in display_context.node_displays
+    node_event_display = display_context.node_displays[str(adornment_node_id)]
     assert node_event_display.subworkflow_display is not None
     assert str(inner_node_id) in node_event_display.subworkflow_display.node_displays
 
+    # AND the inner node should have the correct outputs
+    inner_node_display = node_event_display.subworkflow_display.node_displays[str(inner_node_id)]
+    assert inner_node_display.output_display.keys() == {"foo"}
+    assert node_event_display.output_display.keys() == expected_adornment_output_names
+
 
 def test_get_event_display_context__templating_node_input_display():
     # GIVEN a simple workflow with a templating node referencing another node output