vellum-ai 0.13.10__py3-none-any.whl → 0.13.12__py3-none-any.whl

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

@@ -4,10 +4,12 @@ from deepdiff import DeepDiff
 
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.references.constant import ConstantValueReference
+from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 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.types import NodeOutputDisplay
-from vellum_ee.workflows.display.nodes.vellum.base_node import BaseNodeDisplay
 
 
 class Inputs(BaseInputs):
@@ -19,11 +21,7 @@ class ConstantValueGenericNode(BaseNode):
 
 
 def test_serialize_node__constant_value(serialize_node):
-    input_id = uuid4()
-    serialized_node = serialize_node(
-        node_class=ConstantValueGenericNode,
-        global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=input_id)},
-    )
+    serialized_node = serialize_node(ConstantValueGenericNode)
 
     assert not DeepDiff(
         {
@@ -67,6 +65,92 @@ def test_serialize_node__constant_value(serialize_node):
     )
 
 
+class ConstantValueReferenceGenericNode(BaseNode):
+    attr: str = ConstantValueReference("hello")
+
+
+def test_serialize_node__constant_value_reference(serialize_node):
+    serialized_node = serialize_node(ConstantValueReferenceGenericNode)
+
+    assert not DeepDiff(
+        {
+            "id": "9271e2b1-f47e-47a4-95ae-51299dedb62f",
+            "label": "ConstantValueReferenceGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "ConstantValueReferenceGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "8cc0b4c4-4ae4-4248-8fd5-bfb2e658eb51", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "fe696d84-47c2-4325-8020-34a1c586a759", "name": "default", "type": "DEFAULT"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "460aeb68-7369-43d2-9d3d-37caa425611f",
+                    "name": "attr",
+                    "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "hello"}},
+                }
+            ],
+            "outputs": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
+class LazyReferenceGenericNode(BaseNode):
+    attr: str = LazyReference(lambda: ConstantValueReference("hello"))
+
+
+def test_serialize_node__lazy_reference(serialize_node):
+    serialized_node = serialize_node(LazyReferenceGenericNode)
+
+    assert not DeepDiff(
+        {
+            "id": "29563b11-bd4d-47b0-b017-372f78aeaef5",
+            "label": "LazyReferenceGenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "LazyReferenceGenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "56e9791a-078a-4bb7-90bc-a26c3991c70f", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "acb761a0-fcc2-4d21-bc8c-d0d560912c04", "name": "default", "type": "DEFAULT"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "4370b381-9165-4fb4-881e-480507abe069",
+                    "name": "attr",
+                    "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "hello"}},
+                }
+            ],
+            "outputs": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )
+
+
 class WorkflowInputGenericNode(BaseNode):
     attr: str = Inputs.input
 
@@ -284,3 +368,100 @@ def test_serialize_node__node_execution(serialize_node):
         serialized_node,
         ignore_order=True,
     )
+
+
+class CoalesceNodeA(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output: str
+
+
+class CoalesceNodeADisplay(BaseNodeDisplay[CoalesceNodeA]):
+    pass
+
+
+class CoalesceNodeB(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        output: str
+
+
+class CoalesceNodeBDisplay(BaseNodeDisplay[CoalesceNodeB]):
+    pass
+
+
+class CoalesceNodeFinal(BaseNode):
+    attr = CoalesceNodeA.Outputs.output.coalesce(CoalesceNodeB.Outputs.output)
+
+
+class CoalesceNodeFinalDisplay(BaseNodeDisplay[CoalesceNodeFinal]):
+    pass
+
+
+def test_serialize_node__coalesce(serialize_node):
+    coalesce_node_a_output_id = uuid4()
+    coalesce_node_b_output_id = uuid4()
+    serialized_node = serialize_node(
+        node_class=CoalesceNodeFinal,
+        global_node_displays={
+            CoalesceNodeA: CoalesceNodeADisplay(),
+            CoalesceNodeB: CoalesceNodeBDisplay(),
+            CoalesceNodeFinal: CoalesceNodeFinalDisplay(),
+        },
+        global_node_output_displays={
+            CoalesceNodeA.Outputs.output: (
+                CoalesceNodeA,
+                NodeOutputDisplay(id=coalesce_node_a_output_id, name="output"),
+            ),
+            CoalesceNodeB.Outputs.output: (
+                CoalesceNodeB,
+                NodeOutputDisplay(id=coalesce_node_b_output_id, name="output"),
+            ),
+        },
+    )
+
+    assert not DeepDiff(
+        {
+            "id": "84d0ce62-afd1-4186-b0e3-5dd8e5ca8b65",
+            "label": "CoalesceNodeFinal",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "CoalesceNodeFinal",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "5165f887-153b-4ecd-9219-1beb3cf4f906", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "08ab456d-f541-4400-8305-e97f30cbe745", "name": "default", "type": "DEFAULT"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "5a22dd48-9ef3-456b-85b8-7662cf7823eb",
+                    "name": "attr",
+                    "value": {
+                        "type": "BINARY_EXPRESSION",
+                        "lhs": {
+                            "type": "NODE_OUTPUT",
+                            "node_id": "20c340f2-409c-4d31-b44b-eeffd76938d5",
+                            "node_output_id": str(coalesce_node_a_output_id),
+                        },
+                        "operator": "coalesce",
+                        "rhs": {
+                            "type": "NODE_OUTPUT",
+                            "node_id": "9fac3c54-4e75-46ca-a1f2-a80fc6d8ad3f",
+                            "node_output_id": str(coalesce_node_b_output_id),
+                        },
+                    },
+                }
+            ],
+            "outputs": [],
+        },
+        serialized_node,
+        ignore_order=True,
+    )

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

@@ -5,8 +5,8 @@ from deepdiff import DeepDiff
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 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.types import NodeOutputDisplay
-from vellum_ee.workflows.display.nodes.vellum.base_node import BaseNodeDisplay
 
 
 class Inputs(BaseInputs):

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

@@ -7,8 +7,8 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 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.types import NodeOutputDisplay
-from vellum_ee.workflows.display.nodes.vellum.base_node import BaseNodeDisplay
 
 
 class Inputs(BaseInputs):

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

@@ -88,7 +88,21 @@ def test_serialize_workflow():
         {
             "id": "5cf9c5e3-0eae-4daf-8d73-8b9536258eb9",
             "type": "ERROR",
-            "inputs": [],
+            "inputs": [
+                {
+                    "id": "690d825f-6ffd-493e-8141-c86d384e6150",
+                    "key": "error_source_input_id",
+                    "value": {
+                        "rules": [
+                            {
+                                "type": "CONSTANT_VALUE",
+                                "data": {"type": "STRING", "value": "Input threshold was too low"},
+                            }
+                        ],
+                        "combinator": "OR",
+                    },
+                }
+            ],
             "data": {
                 "name": "error-node",
                 "label": "Fail Node",

vellum_ee/workflows/display/utils/vellum.py

@@ -6,6 +6,7 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.and_ import AndExpression
 from vellum.workflows.expressions.begins_with import BeginsWithExpression
 from vellum.workflows.expressions.between import BetweenExpression
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.expressions.contains import ContainsExpression
 from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
 from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
@@ -150,5 +151,7 @@ def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperato
         return "and"
     elif isinstance(descriptor, OrExpression):
         return "or"
+    elif isinstance(descriptor, CoalesceExpression):
+        return "coalesce"
     else:
         raise ValueError(f"Unsupported descriptor type: {descriptor}")

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -14,6 +14,7 @@ from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.utils.uuids import uuid4_from_hash
+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.types import PortDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -43,7 +44,7 @@ class VellumWorkflowDisplay(
         WorkflowMetaVellumDisplayOverrides,
         WorkflowInputsVellumDisplay,
         WorkflowInputsVellumDisplayOverrides,
-        BaseNodeVellumDisplay,
+        BaseNodeDisplay,
         EntrypointVellumDisplay,
         EntrypointVellumDisplayOverrides,
         EdgeVellumDisplay,
@@ -52,7 +53,7 @@ class VellumWorkflowDisplay(
         WorkflowOutputVellumDisplayOverrides,
     ]
 ):
-    node_display_base_class = BaseNodeVellumDisplay
+    node_display_base_class = BaseNodeDisplay
 
     def serialize(self) -> JsonObject:
         input_variables: JsonArray = []
@@ -137,7 +138,7 @@ class VellumWorkflowDisplay(
                     workflow_output_display.node_input_id,
                 )
 
-                source_node_display: Optional[BaseNodeVellumDisplay]
+                source_node_display: Optional[BaseNodeDisplay]
                 first_rule = node_input.value.rules[0]
                 if first_rule.type == "NODE_OUTPUT":
                     source_node_id = UUID(first_rule.data.node_id)
@@ -170,15 +171,20 @@ class VellumWorkflowDisplay(
                 )
 
                 if source_node_display:
+                    if isinstance(source_node_display, BaseNodeVellumDisplay):
+                        source_handle_id = source_node_display.get_source_handle_id(
+                            port_displays=self.display_context.port_displays
+                        )
+                    else:
+                        source_handle_id = source_node_display.get_node_port_display(
+                            source_node_display._node.Ports.default
+                        ).id
+
                     synthetic_output_edges.append(
                         {
                             "id": str(workflow_output_display.edge_id),
                             "source_node_id": str(source_node_display.node_id),
-                            "source_handle_id": str(
-                                source_node_display.get_source_handle_id(
-                                    port_displays=self.display_context.port_displays
-                                )
-                            ),
+                            "source_handle_id": str(source_handle_id),
                             "target_node_id": str(workflow_output_display.node_id),
                             "target_handle_id": str(workflow_output_display.target_handle_id),
                             "type": "DEFAULT",
@@ -279,7 +285,7 @@ class VellumWorkflowDisplay(
         self,
         entrypoint: Type[BaseNode],
         workflow_display: WorkflowMetaVellumDisplay,
-        node_displays: Dict[Type[BaseNode], BaseNodeVellumDisplay],
+        node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
         overrides: Optional[EntrypointVellumDisplayOverrides] = None,
     ) -> EntrypointVellumDisplay:
         entrypoint_node_id = workflow_display.entrypoint_node_id
@@ -293,8 +299,12 @@ class VellumWorkflowDisplay(
         )
 
         entrypoint_target = get_unadorned_node(entrypoint)
-        target_node_id = node_displays[entrypoint_target].node_id
-        target_handle_id = node_displays[entrypoint_target].get_target_handle_id()
+        target_node_display = node_displays[entrypoint_target]
+        target_node_id = target_node_display.node_id
+        if isinstance(target_node_display, BaseNodeVellumDisplay):
+            target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(entrypoint_node_id)
+        else:
+            target_handle_id = target_node_display.get_trigger_id()
 
         edge_display = self._generate_edge_display_from_source(
             entrypoint_node_id, source_handle_id, target_node_id, target_handle_id, overrides=edge_display_overrides
@@ -339,7 +349,7 @@ class VellumWorkflowDisplay(
     def _generate_edge_display(
         self,
         edge: Edge,
-        node_displays: Dict[Type[BaseNode], BaseNodeVellumDisplay],
+        node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
         port_displays: Dict[Port, PortDisplay],
         overrides: Optional[EdgeVellumDisplayOverrides] = None,
     ) -> EdgeVellumDisplay:
@@ -352,7 +362,11 @@ class VellumWorkflowDisplay(
 
         target_node_display = node_displays[target_node]
         target_node_id = target_node_display.node_id
-        target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
+
+        if isinstance(target_node_display, BaseNodeVellumDisplay):
+            target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
+        else:
+            target_handle_id = target_node_display.get_trigger_id()
 
         return self._generate_edge_display_from_source(
             source_node_id, source_handle_id, target_node_id, target_handle_id, overrides

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

@@ -1,121 +0,0 @@
-import inspect
-from uuid import UUID
-from typing import Any, Generic, Optional, TypeVar, cast
-
-from vellum.workflows.constants import UNDEF
-from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.nodes.bases.base import BaseNode
-from vellum.workflows.nodes.utils import get_wrapped_node
-from vellum.workflows.types.core import JsonArray, JsonObject
-from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
-from vellum.workflows.workflows.base import BaseWorkflow
-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
-from vellum_ee.workflows.display.vellum import GenericNodeDisplayData
-
-_BaseNodeType = TypeVar("_BaseNodeType", bound=BaseNode)
-
-
-class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeType]):
-    def serialize(
-        self, display_context: WorkflowDisplayContext, adornments: Optional[JsonArray] = None, **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 generic node attributes containing a subworkflow
-                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)),
-                }
-            )
-
-        wrapped_node = get_wrapped_node(node)
-        if wrapped_node is not None:
-            display_class = get_node_display_class(BaseNodeDisplay, wrapped_node)
-
-            adornment: JsonObject = {
-                "id": str(node_id),
-                "label": node.__qualname__,
-                "base": self.get_base().dict(),
-                "attributes": attributes,
-            }
-
-            existing_adornments = adornments if adornments is not None else []
-            return display_class().serialize(display_context, adornments=existing_adornments + [adornment])
-
-        ports: JsonArray = []
-        for port in node.Ports:
-            id = str(self.get_node_port_display(port).id)
-
-            if port._condition_type:
-                ports.append(
-                    {
-                        "id": id,
-                        "name": port.name,
-                        "type": port._condition_type.value,
-                        "expression": (
-                            self.serialize_condition(display_context, port._condition) if port._condition else None
-                        ),
-                    }
-                )
-            else:
-                ports.append(
-                    {
-                        "id": id,
-                        "name": port.name,
-                        "type": "DEFAULT",
-                    }
-                )
-
-        outputs: JsonArray = []
-        for output in node.Outputs:
-            type = primitive_type_to_vellum_variable_type(output)
-            value = (
-                self.serialize_value(display_context, output.instance)
-                if output.instance is not None and output.instance != UNDEF
-                else None
-            )
-
-            outputs.append(
-                {
-                    "id": str(uuid4_from_hash(f"{node_id}|{output.name}")),
-                    "name": output.name,
-                    "type": type,
-                    "value": value,
-                }
-            )
-
-        return {
-            "id": str(node_id),
-            "label": node.__qualname__,
-            "type": "GENERIC",
-            "display_data": self.get_generic_node_display_data().dict(),
-            "base": self.get_base().dict(),
-            "definition": self.get_definition().dict(),
-            "trigger": {
-                "id": str(self.get_trigger_id()),
-                "merge_behavior": node.Trigger.merge_behavior.value,
-            },
-            "ports": ports,
-            "adornments": adornments,
-            "attributes": attributes,
-            "outputs": outputs,
-        }
-
-    def get_target_handle_id(self) -> UUID:
-        return self.get_trigger_id()
-
-    def get_generic_node_display_data(self) -> GenericNodeDisplayData:
-        explicit_value = self._get_explicit_node_display_attr("display_data", GenericNodeDisplayData)
-        return explicit_value if explicit_value else GenericNodeDisplayData()