vellum-ai 0.14.4__py3-none-any.whl → 0.14.6__py3-none-any.whl

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

@@ -0,0 +1,39 @@
+from uuid import UUID
+from typing import Any, Callable, Generic, Optional, TypeVar, cast
+
+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
+from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
+
+_BaseAdornmentNodeType = TypeVar("_BaseAdornmentNodeType", bound=BaseAdornmentNode)
+
+
+class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Generic[_BaseAdornmentNodeType]):
+    def serialize(
+        self,
+        display_context: "WorkflowDisplayContext",
+        **kwargs: Any,
+    ) -> dict:
+        node = self._node
+        adornment = cast(Optional[JsonObject], kwargs.get("adornment"))
+        get_additional_kwargs = cast(Optional[Callable[[UUID], dict]], kwargs.get("get_additional_kwargs"))
+
+        wrapped_node = get_wrapped_node(node)
+        if not wrapped_node:
+            raise NotImplementedError(
+                "Unable to serialize standalone adornment nodes. Please use adornment nodes as a decorator."
+            )
+
+        wrapped_node_display_class = get_node_display_class(BaseNodeDisplay, wrapped_node)
+        wrapped_node_display = wrapped_node_display_class()
+        additional_kwargs = get_additional_kwargs(wrapped_node_display.node_id) if get_additional_kwargs else {}
+        serialized_wrapped_node = wrapped_node_display.serialize(display_context, **kwargs, **additional_kwargs)
+
+        adornments = cast(JsonArray, serialized_wrapped_node.get("adornments")) or []
+        serialized_wrapped_node["adornments"] = adornments + [adornment] if adornment else adornments
+
+        return serialized_wrapped_node

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

@@ -3,8 +3,8 @@ from typing import Dict, Generic, List, Optional, TypeVar, cast
 
 from vellum.workflows.nodes import MapNode
 from vellum.workflows.types.core import JsonObject
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
+from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -12,7 +12,7 @@ from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class imp
 _MapNodeType = TypeVar("_MapNodeType", bound=MapNode)
 
 
-class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeType]):
+class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNodeType]):
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

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

@@ -1,10 +1,42 @@
-from typing import Generic, TypeVar
+import inspect
+from typing import Any, Generic, TypeVar, cast
 
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
-from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
+from vellum.workflows.types.core import JsonArray, JsonObject
+from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
+from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _RetryNodeType = TypeVar("_RetryNodeType", bound=RetryNode)
 
 
-class BaseRetryNodeDisplay(BaseNodeDisplay[_RetryNodeType], Generic[_RetryNodeType]):
-    pass
+class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_RetryNodeType]):
+    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+        node = self._node
+        node_id = self.node_id
+
+        attributes: JsonArray = []
+        for attribute in node:
+            if inspect.isclass(attribute.instance) and issubclass(attribute.instance, BaseWorkflow):
+                # We don't need to serialize attributes that are workflows
+                continue
+
+            id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            attributes.append(
+                {
+                    "id": id,
+                    "name": attribute.name,
+                    "value": self.serialize_value(display_context, cast(BaseDescriptor, attribute.instance)),
+                }
+            )
+
+        adornment: JsonObject = {
+            "id": str(node_id),
+            "label": node.__qualname__,
+            "base": self.get_base().dict(),
+            "attributes": attributes,
+        }
+
+        return super().serialize(display_context, adornment=adornment)

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

@@ -7,6 +7,7 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.references import LazyReference
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input_value_pointer_rules
@@ -43,6 +44,10 @@ class MyNodeADisplay(BaseNodeVellumDisplay[MyNodeA]):
 class MyNodeB(BaseNode):
     example = MyNodeA.Outputs.output
     fallback_example = MyNodeA.Outputs.output.coalesce(Inputs.example_workflow_input).coalesce("fallback")
+    constant_coalesce = Inputs.example_workflow_input.coalesce("default_value")
+    lazy_coalesce: BaseDescriptor[str] = LazyReference(
+        lambda: MyNodeA.Outputs.output.coalesce(Inputs.example_workflow_input)
+    )
 
 
 @pytest.mark.parametrize(
@@ -76,6 +81,32 @@ class MyNodeB(BaseNode):
                 ConstantValuePointer(type="CONSTANT_VALUE", data=StringVellumValue(value="fallback")),
             ],
         ),
+        (
+            MyNodeB.constant_coalesce,
+            [
+                InputVariablePointer(
+                    type="INPUT_VARIABLE",
+                    data=InputVariableData(input_variable_id="a154c29d-fac0-4cd0-ba88-bc52034f5470"),
+                ),
+                ConstantValuePointer(type="CONSTANT_VALUE", data=StringVellumValue(value="default_value")),
+            ],
+        ),
+        (
+            MyNodeB.lazy_coalesce,
+            [
+                NodeOutputPointer(
+                    type="NODE_OUTPUT",
+                    data=NodeOutputData(
+                        node_id="b48fa5e0-d7d3-4fe3-ae48-615415011cc5",
+                        output_id="4b16a629-11a1-4b3f-a965-a57b872d13b8",
+                    ),
+                ),
+                InputVariablePointer(
+                    type="INPUT_VARIABLE",
+                    data=InputVariableData(input_variable_id="a154c29d-fac0-4cd0-ba88-bc52034f5470"),
+                ),
+            ],
+        ),
     ],
 )
 def test_create_node_input_value_pointer_rules(

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

@@ -1,55 +1,69 @@
+import inspect
 from uuid import UUID
 from typing import Any, Callable, ClassVar, Generic, Optional, Tuple, Type, TypeVar, cast
 
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
-from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME, get_wrapped_node
+from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
 from vellum.workflows.references.output import OutputReference
-from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.uuids import uuid4_from_hash
+from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
-from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
+from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _TryNodeType = TypeVar("_TryNodeType", bound=TryNode)
 
 
-class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeType]):
+class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNodeType]):
     error_output_id: ClassVar[Optional[UUID]] = None
 
-    def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
+    def serialize(self, display_context: "WorkflowDisplayContext", **kwargs: Any) -> JsonObject:
         node = self._node
+        node_id = self.node_id
+
+        # We let the inner node serialize first and then append to it
+        attributes: JsonArray = []
+        for attribute in node:
+            if inspect.isclass(attribute.instance) and issubclass(attribute.instance, BaseWorkflow):
+                # We don't need to serialize attributes that are workflows
+                continue
+
+            id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            attributes.append(
+                {
+                    "id": id,
+                    "name": attribute.name,
+                    "value": self.serialize_value(display_context, cast(BaseDescriptor, attribute.instance)),
+                }
+            )
 
-        inner_node = get_wrapped_node(node)
-        if not inner_node:
-            subworkflow = raise_if_descriptor(node.subworkflow)
-            if not isinstance(subworkflow.graph, type) or not issubclass(subworkflow.graph, BaseNode):
-                raise NotImplementedError(
-                    "Unable to serialize Try Nodes that wrap subworkflows containing more than one Node."
-                )
-
-            inner_node = subworkflow.graph
-        elif inner_node.__bases__[0] is BaseNode:
-            # If the wrapped node is a generic node, we let generic node do adornment handling
-            class TryBaseNodeDisplay(BaseNodeDisplay[node]):  # type: ignore[valid-type]
-                pass
-
-            return TryBaseNodeDisplay().serialize(display_context)
-
-        # We need the node display class of the underlying node because
-        # it contains the logic for serializing the node and potential display overrides
-        node_display_class = get_node_display_class(BaseNodeDisplay, inner_node)
-        node_display = node_display_class()
-
-        serialized_node = node_display.serialize(
+        adornment: JsonObject = {
+            "id": str(node_id),
+            "label": node.__qualname__,
+            "base": self.get_base().dict(),
+            "attributes": attributes,
+        }
+
+        # We need the inner node's ID to generate the error output ID
+        # Long term we want to hoist error_output_id append from inner node displays to this display,
+        # But that's a lot of work and this allows us to punt a little longer
+        serialized_node = super().serialize(
             display_context,
-            error_output_id=self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id"),
+            adornment=adornment,
+            get_additional_kwargs=lambda node_id: {
+                "error_output_id": self.error_output_id or uuid4_from_hash(f"{node_id}|error_output_id")
+            },
         )
 
+        if serialized_node["type"] == "GENERIC":
+            return serialized_node
+
         serialized_node_definition = serialized_node.get("definition")
         if isinstance(serialized_node_definition, dict):
             serialized_node_definition_module = serialized_node_definition.get("module")

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

@@ -5,8 +5,10 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.references import NodeReference
+from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
+from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
 from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule
 from vellum_ee.workflows.display.vellum import (
     ConstantValuePointer,
@@ -56,6 +58,12 @@ def create_node_input_value_pointer_rules(
                 raise ValueError(f"Expected NodeReference {value.name} to have an instance")
             value = cast(BaseDescriptor, value.instance)
 
+        if isinstance(value, LazyReference):
+            child_descriptor = get_child_descriptor(value, display_context)
+            return create_node_input_value_pointer_rules(
+                child_descriptor, display_context, [], pointer_type=pointer_type
+            )
+
         if isinstance(value, CoalesceExpression):
             # Recursively handle the left-hand side
             lhs_rules = create_node_input_value_pointer_rules(value.lhs, display_context, [], pointer_type=pointer_type)

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

@@ -1,3 +1,4 @@
+import pytest
 from uuid import uuid4
 
 from deepdiff import DeepDiff
@@ -11,6 +12,7 @@ from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.base import WorkflowInputsDisplay
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
+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.workflows.get_vellum_workflow_display_class import get_workflow_display
 from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
@@ -32,7 +34,7 @@ class InnerRetryGenericNodeDisplay(BaseNodeDisplay[InnerRetryGenericNode.__wrapp
     pass
 
 
-class OuterRetryNodeDisplay(BaseNodeDisplay[InnerRetryGenericNode]):
+class OuterRetryNodeDisplay(BaseRetryNodeDisplay[InnerRetryGenericNode]):  # type: ignore
     pass
 
 
@@ -213,3 +215,25 @@ def test_serialize_node__try(serialize_node):
         },
         serialized_node,
     )
+
+
+@pytest.mark.skip(reason="Not implemented")
+def test_serialize_node__stacked():
+    @TryNode.wrap()
+    @RetryNode.wrap(max_attempts=5)
+    class InnerStackedGenericNode(BaseNode):
+        pass
+
+    # AND a workflow that uses the adornment node
+    class StackedWorkflow(BaseWorkflow):
+        graph = InnerStackedGenericNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=StackedWorkflow,
+    )
+    exec_config = workflow_display.serialize()
+
+    # THEN the workflow display is created successfully
+    assert exec_config is not None

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

@@ -596,6 +596,20 @@ def test_serialize_workflow__try_wrapped():
             ],
             "name": "TryNode",
         },
+        "adornments": [
+            {
+                "id": "3344083c-a32c-4a32-920b-0fb5093448fa",
+                "label": "TryNode",
+                "base": {"name": "TryNode", "module": ["vellum", "workflows", "nodes", "core", "try_node", "node"]},
+                "attributes": [
+                    {
+                        "id": "ab2fbab0-e2a0-419b-b1ef-ce11ecf11e90",
+                        "name": "on_error_code",
+                        "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                    }
+                ],
+            }
+        ],
     }
 
     final_output_nodes = workflow_raw_data["nodes"][2:]

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

@@ -1,9 +1,11 @@
+import pytest
+
 from deepdiff import DeepDiff
 
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
-from tests.workflows.basic_try_node.workflow import SimpleTryExample
+from tests.workflows.basic_try_node.workflow import SimpleTryExample, StandaloneTryExample
 
 
 def test_serialize_workflow():
@@ -69,3 +71,19 @@ def test_serialize_workflow():
 
     try_node = workflow_raw_data["nodes"][1]
     assert try_node["id"] == "1381c078-efa2-4255-89a1-7b4cb742c7fc"
+
+
+def test_serialize_workflow__standalone():
+    # GIVEN a Workflow with a standalone TryNode
+    # WHEN we serialize it
+    with pytest.raises(NotImplementedError) as exc:
+        workflow_display = get_workflow_display(
+            base_display_class=VellumWorkflowDisplay, workflow_class=StandaloneTryExample
+        )
+        workflow_display.serialize()
+
+    # THEN we should get an error
+    assert (
+        exc.value.args[0]
+        == "Unable to serialize standalone adornment nodes. Please use adornment nodes as a decorator."
+    )