vellum-ai 0.14.21__py3-none-any.whl → 0.14.23__py3-none-any.whl

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

@@ -3,6 +3,7 @@ import types
 from uuid import UUID
 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
@@ -17,7 +18,62 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _BaseAdornmentNodeType = TypeVar("_BaseAdornmentNodeType", bound=BaseAdornmentNode)
 
 
+def _recursively_replace_wrapped_node(node_class: Type[BaseNode], wrapped_node_display_class: Type["BaseNodeDisplay"]):
+    if not issubclass(node_class, BaseAdornmentNode):
+        return
+
+    # We must edit the node display class to use __wrapped_node__ everywhere it
+    # references the adorned node class, which is three places:
+    wrapped_node_class = node_class.__wrapped_node__
+    if not wrapped_node_class:
+        raise ValueError(f"Node {node_class.__name__} must be used as an adornment with the `wrap` method.")
+
+    # 1. The node display class' parameterized type
+    original_base_node_display = get_original_base(wrapped_node_display_class)
+    original_base_node_display.__args__ = (wrapped_node_class,)
+    wrapped_node_display_class._node_display_registry[wrapped_node_class] = wrapped_node_display_class
+    wrapped_node_display_class.__annotate_node__()
+
+    # 2. The node display class' output displays
+    for old_output in node_class.Outputs:
+        new_output = getattr(wrapped_node_class.Outputs, old_output.name, None)
+        if new_output is None:
+            # If the adornment is adding a new output, such as TryNode adding an "error" output,
+            # we skip it, since it should not be included in the adorned node's output displays
+            wrapped_node_display_class.output_display.pop(old_output, None)
+            continue
+
+        if old_output not in wrapped_node_display_class.output_display:
+            # If the adorned node doesn't have an output display defined for this output, we define one
+            wrapped_node_display_class.output_display[new_output] = NodeOutputDisplay(
+                id=wrapped_node_class.__output_ids__[old_output.name],
+                name=old_output.name,
+            )
+        else:
+            wrapped_node_display_class.output_display[new_output] = wrapped_node_display_class.output_display.pop(
+                old_output
+            )
+
+    # 3. The node display class' port displays
+    old_ports = list(wrapped_node_display_class.port_displays.keys())
+    for old_port in old_ports:
+        new_port = getattr(wrapped_node_class.Ports, old_port.name)
+        wrapped_node_display_class.port_displays[new_port] = wrapped_node_display_class.port_displays.pop(old_port)
+
+    # Now we keep drilling down recursively
+    if (
+        issubclass(wrapped_node_display_class, BaseAdornmentNodeDisplay)
+        and wrapped_node_display_class.__wrapped_node_display__
+    ):
+        _recursively_replace_wrapped_node(
+            wrapped_node_class,
+            wrapped_node_display_class.__wrapped_node_display__,
+        )
+
+
 class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Generic[_BaseAdornmentNodeType]):
+    __wrapped_node_display__: Optional[Type[BaseNodeDisplay]] = None
+
     def serialize(
         self,
         display_context: "WorkflowDisplayContext",
@@ -52,10 +108,6 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
             if not issubclass(node_class, BaseAdornmentNode):
                 raise ValueError(f"Node {node_class.__name__} must be wrapped with a {BaseAdornmentNode.__name__}")
 
-            wrapped_node_class = node_class.__wrapped_node__
-            if not wrapped_node_class:
-                raise ValueError(f"Node {node_class.__name__} must be used as an adornment with the `wrap` method.")
-
             # `mypy` is wrong here, `cls` is indexable bc it's Generic
             BaseAdornmentDisplay = cls[node_class]  # type: ignore[index]
 
@@ -64,6 +116,7 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
                     ns[key] = kwarg
 
                 ns["node_id"] = node_id or uuid4_from_hash(node_class.__qualname__)
+                ns["__wrapped_node_display__"] = wrapped_node_display_class
 
             AdornmentDisplay = types.new_class(
                 re.sub(r"^Base", "", cls.__name__), bases=(BaseAdornmentDisplay,), exec_body=exec_body
@@ -71,43 +124,11 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
 
             setattr(wrapped_node_display_class, "__adorned_by__", AdornmentDisplay)
 
-            # We must edit the node display class to use __wrapped_node__ everywhere it
-            # references the adorned node class, which is three places:
-
-            # 1. The node display class' parameterized type
-            original_base_node_display = get_original_base(wrapped_node_display_class)
-            original_base_node_display.__args__ = (wrapped_node_class,)
-            wrapped_node_display_class._node_display_registry[wrapped_node_class] = wrapped_node_display_class
-            wrapped_node_display_class.__annotate_node__()
-
-            # 2. The node display class' output displays
-            for old_output in node_class.Outputs:
-                new_output = getattr(wrapped_node_class.Outputs, old_output.name, None)
-                if new_output is None:
-                    # If the adornment is adding a new output, such as TryNode adding an "error" output,
-                    # we skip it, since it should not be included in the adorned node's output displays
-                    wrapped_node_display_class.output_display.pop(old_output, None)
-                    continue
-
-                if old_output not in wrapped_node_display_class.output_display:
-                    # If the adorned node doesn't have an output display defined for this output, we define one
-                    wrapped_node_display_class.output_display[new_output] = NodeOutputDisplay(
-                        id=wrapped_node_class.__output_ids__[old_output.name],
-                        name=old_output.name,
-                    )
-                else:
-                    wrapped_node_display_class.output_display[new_output] = (
-                        wrapped_node_display_class.output_display.pop(old_output)
-                    )
-
-            # 3. The node display class' port displays
-            old_ports = list(wrapped_node_display_class.port_displays.keys())
-            for old_port in old_ports:
-                new_port = getattr(wrapped_node_class.Ports, old_port.name)
-                wrapped_node_display_class.port_displays[new_port] = wrapped_node_display_class.port_displays.pop(
-                    old_port
-                )
-
-            return wrapped_node_display_class
+            _recursively_replace_wrapped_node(
+                node_class,
+                wrapped_node_display_class,
+            )
+
+            return AdornmentDisplay
 
         return decorator

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

@@ -47,6 +47,14 @@ class BaseFinalOutputNodeDisplay(BaseNodeVellumDisplay[_FinalOutputNodeType], Ge
             "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": self.serialize_value(display_context, node.Outputs.value.instance),
+                }
+            ],
         }
 
     def _get_output_id(self) -> UUID:

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

@@ -99,4 +99,12 @@ def test_serialize_workflow():
             "name": "BasicFinalOutputNode",
             "module": ["tests", "workflows", "basic_final_output_node", "workflow"],
         },
+        "outputs": [
+            {
+                "id": "97349956-d228-4b51-a64b-1331f788373f",
+                "name": "value",
+                "type": "STRING",
+                "value": {"type": "WORKFLOW_INPUT", "input_variable_id": "e39a7b63-de15-490a-ae9b-8112c767aea0"},
+            }
+        ],
     }

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

@@ -123,6 +123,17 @@ def test_serialize_workflow__missing_final_output_node():
                     "name": "FirstFinalOutputNode",
                     "module": ["tests", "workflows", "complex_final_output_node", "missing_final_output_node"],
                 },
+                "outputs": [
+                    {
+                        "id": "5517e50d-7f40-4f7c-acb2-e329d79a25bf",
+                        "name": "value",
+                        "type": "STRING",
+                        "value": {
+                            "type": "WORKFLOW_INPUT",
+                            "input_variable_id": "da086239-d743-4246-b666-5c91e22fb88c",
+                        },
+                    }
+                ],
             },
             {
                 "id": "bb88768d-472e-4997-b7ea-de09163d1b4c",
@@ -156,14 +167,7 @@ def test_serialize_workflow__missing_final_output_node():
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
                 "base": {
                     "name": "FinalOutputNode",
-                    "module": [
-                        "vellum",
-                        "workflows",
-                        "nodes",
-                        "displayable",
-                        "final_output_node",
-                        "node",
-                    ],
+                    "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
                 },
                 "definition": None,
             },

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -10,7 +10,6 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.events.workflow import NodeEventDisplayContext, WorkflowEventDisplayContext
-from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.utils import get_unadorned_node, get_wrapped_node
 from vellum.workflows.ports import Port
@@ -150,8 +149,6 @@ class BaseWorkflowDisplay(
             if node_output in node_output_displays:
                 continue
 
-            # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
-            # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
             node_output_displays[node_output] = node_display.get_node_output_display(node_output)
 
     def _enrich_node_port_displays(
@@ -272,8 +269,8 @@ class BaseWorkflowDisplay(
                 continue
 
             edge_display_overrides = self.edge_displays.get((edge.from_port, edge.to_node))
-            edge_displays[(edge.from_port, edge.to_node)] = self._generate_edge_display(
-                edge, node_displays, overrides=edge_display_overrides
+            edge_displays[(edge.from_port, edge.to_node)] = edge_display_overrides or self._generate_edge_display(
+                edge, node_displays
             )
 
         for edge in self._workflow.get_unused_edges():
@@ -281,8 +278,8 @@ class BaseWorkflowDisplay(
                 continue
 
             edge_display_overrides = self.edge_displays.get((edge.from_port, edge.to_node))
-            edge_displays[(edge.from_port, edge.to_node)] = self._generate_edge_display(
-                edge, node_displays, overrides=edge_display_overrides
+            edge_displays[(edge.from_port, edge.to_node)] = edge_display_overrides or self._generate_edge_display(
+                edge, node_displays
             )
 
         workflow_output_displays: Dict[BaseDescriptor, WorkflowOutputDisplay] = {}
@@ -293,11 +290,6 @@ class BaseWorkflowDisplay(
             if not isinstance(workflow_output, OutputReference):
                 raise ValueError(f"{workflow_output} must be an {OutputReference.__name__}")
 
-            if not workflow_output.instance or not isinstance(
-                workflow_output.instance, (OutputReference, CoalesceExpression)
-            ):
-                raise ValueError("Expected to find a descriptor instance on the workflow output")
-
             workflow_output_display = self.output_displays.get(workflow_output)
             workflow_output_displays[workflow_output] = (
                 workflow_output_display or self._generate_workflow_output_display(workflow_output)
@@ -443,30 +435,20 @@ class BaseWorkflowDisplay(
 
         return additional_node_displays
 
-    def _generate_edge_display(
-        self,
-        edge: Edge,
-        node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
-        overrides: Optional[EdgeDisplay] = None,
-    ) -> EdgeDisplay:
+    def _generate_edge_display(self, edge: Edge, node_displays: Dict[Type[BaseNode], BaseNodeDisplay]) -> EdgeDisplay:
         source_node = get_unadorned_node(edge.from_port.node_class)
         target_node = get_unadorned_node(edge.to_node)
 
         source_node_id = node_displays[source_node].node_id
         target_node_id = node_displays[target_node].node_id
 
-        return self._generate_edge_display_from_source(source_node_id, target_node_id, overrides)
+        return self._generate_edge_display_from_source(source_node_id, target_node_id)
 
     def _generate_edge_display_from_source(
         self,
         source_node_id: UUID,
         target_node_id: UUID,
-        overrides: Optional[EdgeDisplay] = None,
     ) -> EdgeDisplay:
-        edge_id: UUID
-        if overrides:
-            edge_id = overrides.id
-        else:
-            edge_id = uuid4_from_hash(f"{self.workflow_id}|id|{source_node_id}|{target_node_id}")
-
-        return EdgeDisplay(id=edge_id)
+        return EdgeDisplay(
+            id=uuid4_from_hash(f"{self.workflow_id}|id|{source_node_id}|{target_node_id}"),
+        )

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

@@ -235,3 +235,57 @@ def test_get_event_display_context__templating_node_input_display():
     node_event_display = display_context.node_displays[MyNode.__id__]
 
     assert node_event_display.input_display.keys() == {"inputs.foo"}
+
+
+def test_get_event_display_context__node_display_for_mutiple_adornments():
+    # GIVEN a simple workflow with multiple adornments
+    @TryNode.wrap()
+    @RetryNode.wrap()
+    class MyNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+    class MyWorkflow(BaseWorkflow):
+        graph = MyNode
+
+    # AND a display class for the node
+    node_id = uuid4()
+    inner_node_id = uuid4()
+    innermost_node_id = uuid4()
+
+    @BaseTryNodeDisplay.wrap(node_id=node_id)
+    @BaseRetryNodeDisplay.wrap(node_id=inner_node_id)
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        node_id = innermost_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 node_id in display_context.node_displays
+    node_event_display = display_context.node_displays[node_id]
+    assert node_event_display.subworkflow_display
+
+    # AND the inner node should be included
+    assert inner_node_id in node_event_display.subworkflow_display.node_displays
+    inner_node_event_display = node_event_display.subworkflow_display.node_displays[inner_node_id]
+    assert inner_node_event_display.subworkflow_display
+
+    # AND the innermost node should be included
+    assert innermost_node_id in inner_node_event_display.subworkflow_display.node_displays
+    innermost_node_event_display = inner_node_event_display.subworkflow_display.node_displays[innermost_node_id]
+    assert not innermost_node_event_display.subworkflow_display
+
+
+def test_get_event_display_context__workflow_output_display_with_none():
+    # GIVEN a workflow with a workflow output that is None
+    class MyWorkflow(BaseWorkflow):
+        class Outputs(BaseWorkflow.Outputs):
+            foo = None
+            bar = "baz"
+
+    # WHEN we gather the event display context
+    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+
+    # THEN the workflow output display should be included
+    assert display_context.workflow_outputs.keys() == {"foo", "bar"}

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -363,8 +363,8 @@ class VellumWorkflowDisplay(
         target_node_display = node_displays[entrypoint_target]
         target_node_id = target_node_display.node_id
 
-        edge_display = self._generate_edge_display_from_source(
-            entrypoint_node_id, target_node_id, overrides=edge_display_overrides
+        edge_display = edge_display_overrides or self._generate_edge_display_from_source(
+            entrypoint_node_id, target_node_id
         )
 
         return EntrypointVellumDisplay(id=entrypoint_id, edge_display=edge_display)