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

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -16,15 +16,28 @@ from typing import (
     get_origin,
 )
 
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.expressions.between import BetweenExpression
+from vellum.workflows.expressions.is_nil import IsNilExpression
+from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
+from vellum.workflows.expressions.is_not_null import IsNotNullExpression
+from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
+from vellum.workflows.expressions.is_null import IsNullExpression
+from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
+from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference
+from vellum.workflows.references.execution_count import ExecutionCountReference
+from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.references.workflow_input import WorkflowInputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.names import pascal_to_title_case
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator, primitive_to_vellum_value
 from vellum_ee.workflows.display.vellum import CodeResourceDefinition
 
 if TYPE_CHECKING:
@@ -76,6 +89,9 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         )
         return node_definition
 
+    def get_trigger_id(self) -> UUID:
+        return uuid4_from_hash(f"{self.node_id}|trigger")
+
     def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
         explicit_display = self.output_display.get(output)
         if explicit_display:
@@ -174,3 +190,87 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
         node_class = cls.infer_node_class()
         cls._node_display_registry[node_class] = cls
+
+    def serialize_condition(self, display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
+        if isinstance(
+            condition,
+            (
+                IsNullExpression,
+                IsNotNullExpression,
+                IsNilExpression,
+                IsNotNilExpression,
+                IsUndefinedExpression,
+                IsNotUndefinedExpression,
+            ),
+        ):
+            lhs = self.serialize_value(display_context, condition._expression)
+            return {
+                "type": "UNARY_EXPRESSION",
+                "lhs": lhs,
+                "operator": convert_descriptor_to_operator(condition),
+            }
+        elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
+            base = self.serialize_value(display_context, condition._value)
+            lhs = self.serialize_value(display_context, condition._start)
+            rhs = self.serialize_value(display_context, condition._end)
+
+            return {
+                "type": "TERNARY_EXPRESSION",
+                "base": base,
+                "operator": convert_descriptor_to_operator(condition),
+                "lhs": lhs,
+                "rhs": rhs,
+            }
+        else:
+            lhs = self.serialize_value(display_context, condition._lhs)  # type: ignore[attr-defined]
+            rhs = self.serialize_value(display_context, condition._rhs)  # type: ignore[attr-defined]
+
+            return {
+                "type": "BINARY_EXPRESSION",
+                "lhs": lhs,
+                "operator": convert_descriptor_to_operator(condition),
+                "rhs": rhs,
+            }
+
+    def serialize_value(self, display_context: "WorkflowDisplayContext", value: BaseDescriptor) -> JsonObject:
+        if isinstance(value, WorkflowInputReference):
+            workflow_input_display = display_context.global_workflow_input_displays[value]
+            return {
+                "type": "WORKFLOW_INPUT",
+                "input_variable_id": str(workflow_input_display.id),
+            }
+
+        if isinstance(value, OutputReference):
+            upstream_node, output_display = display_context.global_node_output_displays[value]
+            upstream_node_display = display_context.global_node_displays[upstream_node]
+
+            return {
+                "type": "NODE_OUTPUT",
+                "node_id": str(upstream_node_display.node_id),
+                "node_output_id": str(output_display.id),
+            }
+
+        if isinstance(value, VellumSecretReference):
+            return {
+                "type": "VELLUM_SECRET",
+                "vellum_secret_name": value.name,
+            }
+
+        if isinstance(value, ExecutionCountReference):
+            node_class_display = display_context.global_node_displays[value.node_class]
+
+            return {
+                "type": "EXECUTION_COUNTER",
+                "node_id": str(node_class_display.node_id),
+            }
+
+        if not isinstance(value, BaseDescriptor):
+            vellum_value = primitive_to_vellum_value(value)
+            return {
+                "type": "CONSTANT_VALUE",
+                "value": vellum_value.dict(),
+            }
+
+        # If it's not any of the references we know about,
+        # then try to serialize it as a nested value
+        return self.serialize_condition(display_context, value)

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

@@ -1,41 +1,37 @@
-from typing import Any, Generic, TypeVar, cast
+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.expressions.between import BetweenExpression
-from vellum.workflows.expressions.is_nil import IsNilExpression
-from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
-from vellum.workflows.expressions.is_not_null import IsNotNullExpression
-from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
-from vellum.workflows.expressions.is_null import IsNullExpression
-from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
-from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.nodes.bases.base import BaseNode
-from vellum.workflows.references.execution_count import ExecutionCountReference
-from vellum.workflows.references.output import OutputReference
-from vellum.workflows.references.vellum_secret import VellumSecretReference
-from vellum.workflows.references.workflow_input import WorkflowInputReference
+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.vellum.utils import convert_descriptor_to_operator
+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.utils.vellum import primitive_to_vellum_value
 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, **kwargs: Any) -> JsonObject:
+    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:
-            id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
+            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,
@@ -44,9 +40,23 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
                 }
             )
 
+        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 idx, port in enumerate(node.Ports):
-            id = str(uuid4_from_hash(f"{node_id}|{idx}"))
+        for port in node.Ports:
+            id = str(self.get_node_port_display(port).id)
 
             if port._condition_type:
                 ports.append(
@@ -94,99 +104,18 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
             "base": self.get_base().dict(),
             "definition": self.get_definition().dict(),
             "trigger": {
-                "id": str(uuid4_from_hash(f"{node_id}|trigger")),
+                "id": str(self.get_trigger_id()),
                 "merge_behavior": node.Trigger.merge_behavior.value,
             },
             "ports": ports,
-            "adornments": None,
+            "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()
-
-    def serialize_condition(self, display_context: WorkflowDisplayContext, condition: BaseDescriptor) -> JsonObject:
-        if isinstance(
-            condition,
-            (
-                IsNullExpression,
-                IsNotNullExpression,
-                IsNilExpression,
-                IsNotNilExpression,
-                IsUndefinedExpression,
-                IsNotUndefinedExpression,
-            ),
-        ):
-            lhs = self.serialize_value(display_context, condition._expression)
-            return {
-                "type": "UNARY_EXPRESSION",
-                "lhs": lhs,
-                "operator": convert_descriptor_to_operator(condition),
-            }
-        elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
-            base = self.serialize_value(display_context, condition._value)
-            lhs = self.serialize_value(display_context, condition._start)
-            rhs = self.serialize_value(display_context, condition._end)
-
-            return {
-                "type": "TERNARY_EXPRESSION",
-                "base": base,
-                "operator": convert_descriptor_to_operator(condition),
-                "lhs": lhs,
-                "rhs": rhs,
-            }
-        else:
-            lhs = self.serialize_value(display_context, condition._lhs)  # type: ignore[attr-defined]
-            rhs = self.serialize_value(display_context, condition._rhs)  # type: ignore[attr-defined]
-
-            return {
-                "type": "BINARY_EXPRESSION",
-                "lhs": lhs,
-                "operator": convert_descriptor_to_operator(condition),
-                "rhs": rhs,
-            }
-
-    def serialize_value(self, display_context: WorkflowDisplayContext, value: BaseDescriptor) -> JsonObject:
-        if isinstance(value, WorkflowInputReference):
-            workflow_input_display = display_context.global_workflow_input_displays[value]
-            return {
-                "type": "WORKFLOW_INPUT",
-                "input_variable_id": str(workflow_input_display.id),
-            }
-
-        if isinstance(value, OutputReference):
-            upstream_node, output_display = display_context.global_node_output_displays[value]
-            upstream_node_display = display_context.global_node_displays[upstream_node]
-
-            return {
-                "type": "NODE_OUTPUT",
-                "node_id": str(upstream_node_display.node_id),
-                "node_output_id": str(output_display.id),
-            }
-
-        if isinstance(value, VellumSecretReference):
-            return {
-                "type": "VELLUM_SECRET",
-                "vellum_secret_name": value.name,
-            }
-
-        if isinstance(value, ExecutionCountReference):
-            node_class_display = display_context.global_node_displays[value.node_class]
-
-            return {
-                "type": "EXECUTION_COUNTER",
-                "node_id": str(node_class_display.node_id),
-            }
-
-        if not isinstance(value, BaseDescriptor):
-            vellum_value = primitive_to_vellum_value(value)
-            return {
-                "type": "CONSTANT_VALUE",
-                "value": vellum_value.dict(),
-            }
-
-        # If it's not any of the references we know about,
-        # then try to serialize it as a nested value
-        return self.serialize_condition(display_context, value)

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

@@ -17,8 +17,9 @@ from vellum.workflows.nodes.displayable import ConditionalNode
 from vellum.workflows.types.core import ConditionType, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
-from vellum_ee.workflows.display.nodes.vellum.utils import convert_descriptor_to_operator, create_node_input
+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.utils.vellum import convert_descriptor_to_operator
 from vellum_ee.workflows.display.vellum import NodeInput
 
 _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)

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

@@ -13,6 +13,7 @@ from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeV
 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_node import BaseNodeDisplay as GenericBaseNodeDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _TryNodeType = TypeVar("_TryNodeType", bound=TryNode)
@@ -33,6 +34,12 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
                 )
 
             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(GenericBaseNodeDisplay[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

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

@@ -2,31 +2,7 @@ from uuid import UUID
 from typing import Any, List, Optional, Type, Union, cast
 
 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
-from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
-from vellum.workflows.expressions.does_not_equal import DoesNotEqualExpression
-from vellum.workflows.expressions.ends_with import EndsWithExpression
-from vellum.workflows.expressions.equals import EqualsExpression
-from vellum.workflows.expressions.greater_than import GreaterThanExpression
-from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
-from vellum.workflows.expressions.in_ import InExpression
-from vellum.workflows.expressions.is_nil import IsNilExpression
-from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
-from vellum.workflows.expressions.is_not_null import IsNotNullExpression
-from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
-from vellum.workflows.expressions.is_null import IsNullExpression
-from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
-from vellum.workflows.expressions.less_than import LessThanExpression
-from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
-from vellum.workflows.expressions.not_between import NotBetweenExpression
-from vellum.workflows.expressions.not_in import NotInExpression
-from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.references import NodeReference
 from vellum.workflows.utils.uuids import uuid4_from_hash
@@ -124,48 +100,3 @@ def create_pointer(
         return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_variable_value)
     else:
         raise ValueError(f"Pointer type {pointer_type} not supported")
-
-
-def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> str:
-    if isinstance(descriptor, EqualsExpression):
-        return "="
-    elif isinstance(descriptor, DoesNotEqualExpression):
-        return "!="
-    elif isinstance(descriptor, LessThanExpression):
-        return "<"
-    elif isinstance(descriptor, GreaterThanExpression):
-        return ">"
-    elif isinstance(descriptor, LessThanOrEqualToExpression):
-        return "<="
-    elif isinstance(descriptor, GreaterThanOrEqualToExpression):
-        return ">="
-    elif isinstance(descriptor, ContainsExpression):
-        return "contains"
-    elif isinstance(descriptor, BeginsWithExpression):
-        return "beginsWith"
-    elif isinstance(descriptor, EndsWithExpression):
-        return "endsWith"
-    elif isinstance(descriptor, DoesNotContainExpression):
-        return "doesNotContain"
-    elif isinstance(descriptor, DoesNotBeginWithExpression):
-        return "doesNotBeginWith"
-    elif isinstance(descriptor, DoesNotEndWithExpression):
-        return "doesNotEndWith"
-    elif isinstance(descriptor, (IsNullExpression, IsNilExpression, IsUndefinedExpression)):
-        return "null"
-    elif isinstance(descriptor, (IsNotNullExpression, IsNotNilExpression, IsNotUndefinedExpression)):
-        return "notNull"
-    elif isinstance(descriptor, InExpression):
-        return "in"
-    elif isinstance(descriptor, NotInExpression):
-        return "notIn"
-    elif isinstance(descriptor, BetweenExpression):
-        return "between"
-    elif isinstance(descriptor, NotBetweenExpression):
-        return "notBetween"
-    elif isinstance(descriptor, AndExpression):
-        return "and"
-    elif isinstance(descriptor, OrExpression):
-        return "or"
-    else:
-        raise ValueError(f"Unsupported descriptor type: {descriptor}")

vellum_ee/workflows/display/tests/test_vellum_workflow_display.py

@@ -88,3 +88,59 @@ def test_vellum_workflow_display__serialize_input_variables_with_capitalized_var
             "extensions": {"color": None},
         }
     ]
+
+
+def test_vellum_workflow_display_serialize_valid_handle_ids_for_base_nodes():
+    # GIVEN a workflow between two base nodes
+    class StartNode(BaseNode):
+        pass
+
+    class EndNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            hello = "world"
+
+    class Workflow(BaseWorkflow):
+        graph = StartNode >> EndNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_value = EndNode.Outputs.hello
+
+    # AND a display class for this workflow
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay,
+        workflow_class=Workflow,
+    )
+
+    # WHEN we serialize the workflow
+    exec_config = workflow_display.serialize()
+
+    # THEN the serialized workflow handle ids are valid
+    raw_data = exec_config.get("workflow_raw_data")
+    assert isinstance(raw_data, dict)
+    nodes = raw_data.get("nodes")
+    edges = raw_data.get("edges")
+
+    assert isinstance(nodes, list)
+    assert isinstance(edges, list)
+
+    edge_source_handle_ids = {edge.get("source_handle_id") for edge in edges if isinstance(edge, dict)}
+    edge_target_handle_ids = {edge.get("target_handle_id") for edge in edges if isinstance(edge, dict)}
+
+    for node in nodes:
+        assert isinstance(node, dict)
+
+        if node["type"] in {"ENTRYPOINT", "TERMINAL"}:
+            continue
+
+        ports = node.get("ports")
+        assert isinstance(ports, list)
+        for port in ports:
+            assert isinstance(port, dict)
+            assert (
+                port["id"] in edge_source_handle_ids
+            ), f"Port {port['id']} from node {node['label']} not found in edge source handle ids"
+
+        assert isinstance(node["trigger"], dict)
+        assert (
+            node["trigger"]["id"] in edge_target_handle_ids
+        ), f"Trigger {node['trigger']['id']} from node {node['label']} not found in edge target handle ids"

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

@@ -1,6 +1,6 @@
 import pytest
 from uuid import uuid4
-from typing import Dict, Tuple, Type
+from typing import Any, Dict, Tuple, Type
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.references.output import OutputReference
@@ -20,11 +20,12 @@ from vellum_ee.workflows.display.workflows.vellum_workflow_display import Vellum
 def serialize_node():
     def _serialize_node(
         node_class: Type[NodeType],
+        base_class: type[BaseNodeDisplay[Any]] = BaseNodeDisplay,
         global_workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {},
         global_node_displays: Dict[Type[BaseNode], NodeDisplayType] = {},
         global_node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]] = {},
     ) -> JsonObject:
-        node_display_class = get_node_display_class(BaseNodeDisplay, node_class)
+        node_display_class = get_node_display_class(base_class, node_class)
         node_display = node_display_class()
 
         context: WorkflowDisplayContext = WorkflowDisplayContext(