vellum-ai 0.14.11__py3-none-any.whl → 0.14.13__py3-none-any.whl

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -59,11 +59,27 @@ _NodeDisplayAttrType = TypeVar("_NodeDisplayAttrType")
 class BaseNodeDisplayMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         cls = super().__new__(mcs, name, bases, dct)
-        if isinstance(dct.get("node_id"), UUID):
+        base_node_display_class = cast(Type["BaseNodeDisplay"], cls)
+        node_class = base_node_display_class.infer_node_class()
+        if not issubclass(node_class, BaseNode):
+            return cls
+
+        display_node_id = dct.get("node_id")
+        if isinstance(display_node_id, UUID):
             # Display classes are able to override the id of the node class it's parameterized by
-            base_node_display_class = cast(Type["BaseNodeDisplay"], cls)
-            node_class = base_node_display_class.infer_node_class()
-            node_class.__id__ = dct["node_id"]
+            node_class.__id__ = display_node_id
+
+        output_display = dct.get("output_display")
+        if isinstance(output_display, dict):
+            # And the node class' output ids
+            for reference, node_output_display in output_display.items():
+                if not isinstance(reference, OutputReference):
+                    continue
+                if not isinstance(node_output_display, NodeOutputDisplay):
+                    continue
+
+                node_class.__output_ids__[reference.name] = node_output_display.id
+
         return cls
 
 

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -2,6 +2,7 @@ import types
 from typing import TYPE_CHECKING, Optional, Type
 
 from vellum.workflows.types.generics import NodeType
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import NodeDisplayType
@@ -29,4 +30,12 @@ def get_node_display_class(
         f"{node_class.__name__}Display",
         bases=(NodeDisplayBaseClass,),
     )
+    output_display = {
+        ref: NodeOutputDisplay(id=node_class.__output_ids__[ref.name], name=ref.name)
+        for ref in node_class.Outputs
+        if ref.name in node_class.__output_ids__
+    }
+    if output_display:
+        setattr(NodeDisplayClass, "output_display", output_display)
+
     return NodeDisplayClass

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

@@ -115,7 +115,8 @@ def test_serialize_node__retry(serialize_node):
     )
 
 
-def test_serialize_node__retry__no_display():  # GIVEN an adornment node
+def test_serialize_node__retry__no_display():
+    # GIVEN an adornment node
     @RetryNode.wrap(max_attempts=5)
     class StartNode(BaseNode):
         pass
@@ -212,6 +213,28 @@ def test_serialize_node__try(serialize_node):
     )
 
 
+def test_serialize_node__try__no_display():
+    # GIVEN an adornment node
+    @TryNode.wrap()
+    class StartNode(BaseNode):
+        pass
+
+    # AND a workflow that uses the adornment node
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=MyWorkflow,
+    )
+
+    exec_config = workflow_display.serialize()
+
+    # THEN the workflow display is created successfully
+    assert exec_config is not None
+
+
 def test_serialize_node__stacked():
     @TryNode.wrap()
     @RetryNode.wrap(max_attempts=5)

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -411,7 +411,7 @@ class BaseWorkflowDisplay(
             input_display = {}
             if isinstance(current_node_display, BaseNodeVellumDisplay):
                 input_display = current_node_display.node_input_ids_by_name
-            node_display_meta = {
+            output_display = {
                 output.name: current_node_display.output_display[output].id
                 for output in current_node_display.output_display
             }
@@ -434,7 +434,7 @@ class BaseWorkflowDisplay(
 
             node_event_displays[node_id] = NodeEventDisplayContext(
                 input_display=input_display,
-                output_display=node_display_meta,
+                output_display=output_display,
                 port_display=port_display_meta,
                 subworkflow_display=subworkflow_display_context,
             )

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

@@ -95,6 +95,26 @@ def test_serialize_workflow__node_display_class_not_registered():
     assert data is not None
 
 
+def test_get_event_display_context__node_display_filled_without_base_display():
+    # GIVEN a simple workflow
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+    # 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 StartNode.__output_ids__ == node_event_display.output_display
+
+
 def test_get_event_display_context__node_display_to_include_subworkflow_display():
     # GIVEN a simple workflow
     class InnerNode(BaseNode):