vellum-ai 0.14.15__py3-none-any.whl → 0.14.17__py3-none-any.whl

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -59,7 +59,16 @@ _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)
+        cls = cast(Type["BaseNodeDisplay"], super().__new__(mcs, name, bases, dct))
+
+        if not dct.get("output_display"):
+            node_class = cls.infer_node_class()
+            cls.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__
+            }
+
         return cls.__annotate_node__()
 
     def __annotate_node__(cls):
@@ -73,22 +82,36 @@ class BaseNodeDisplayMeta(type):
             # Display classes are able to override the id of the node class it's parameterized by
             node_class.__id__ = display_node_id
 
-        output_display = getattr(cls, "output_display", None)
-        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
+        for reference, node_output_display in base_node_display_class.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
+            node_class.__output_ids__[reference.name] = node_output_display.id
 
         return cls
 
+    def infer_node_class(cls) -> Type[BaseNode]:
+        original_base = get_original_base(cls)
+        node_class = get_args(original_base)[0]
+        if isinstance(node_class, TypeVar):
+            bounded_class = node_class.__bound__
+            if inspect.isclass(bounded_class) and issubclass(bounded_class, BaseNode):
+                return bounded_class
+
+            if isinstance(bounded_class, ForwardRef) and bounded_class.__forward_arg__ == BaseNode.__name__:
+                return BaseNode
+
+        if issubclass(node_class, BaseNode):
+            return node_class
+
+        raise ValueError(f"Node {cls.__name__} must be a subclass of {BaseNode.__name__}")
+
 
 class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
-    output_display: Dict[OutputReference, NodeOutputDisplay] = {}
+    # Default values set by the metaclass
+    output_display: Dict[OutputReference, NodeOutputDisplay]
     port_displays: Dict[Port, PortDisplayOverrides] = {}
     node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
@@ -235,23 +258,6 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
     def get_from_node_display_registry(cls, node_class: Type[NodeType]) -> Optional[Type["BaseNodeDisplay"]]:
         return cls._node_display_registry.get(node_class)
 
-    @classmethod
-    def infer_node_class(cls) -> Type[NodeType]:
-        original_base = get_original_base(cls)
-        node_class = get_args(original_base)[0]
-        if isinstance(node_class, TypeVar):
-            bounded_class = node_class.__bound__
-            if inspect.isclass(bounded_class) and issubclass(bounded_class, BaseNode):
-                return cast(Type[NodeType], bounded_class)
-
-            if isinstance(bounded_class, ForwardRef) and bounded_class.__forward_arg__ == BaseNode.__name__:
-                return cast(Type[NodeType], BaseNode)
-
-        if issubclass(node_class, BaseNode):
-            return node_class
-
-        raise ValueError(f"Node {cls.__name__} must be a subclass of {BaseNode.__name__}")
-
     @cached_property
     def node_id(self) -> UUID:
         """Can be overridden as a class attribute to specify a custom node id."""
@@ -264,7 +270,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
 
     @property
     def _node(self) -> Type[NodeType]:
-        return self.infer_node_class()
+        return cast(Type[NodeType], self.__class__.infer_node_class())
 
     @classmethod
     def _get_explicit_node_display_attr(

vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import NodeDisplayType
@@ -43,14 +42,6 @@ def get_node_display_class(
         return {}
 
     def exec_body(ns: Dict):
-        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:
-            ns["output_display"] = output_display
-
         node_input_ids_by_name: Dict[str, UUID] = {}
         for ref in node_class:
             node_input_ids_by_name.update(_get_node_input_ids_by_ref(ref.name, ref.instance))

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

@@ -3,7 +3,6 @@ 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
@@ -12,6 +11,7 @@ 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.get_node_display_class import get_node_display_class
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 _BaseAdornmentNodeType = TypeVar("_BaseAdornmentNodeType", bound=BaseAdornmentNode)
@@ -44,12 +44,17 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
         return serialized_wrapped_node
 
     @classmethod
-    def wrap(cls, **kwargs: Any) -> Callable[..., Type[BaseNodeDisplay]]:
+    def wrap(cls, node_id: Optional[UUID] = None, **kwargs: Any) -> Callable[..., Type[BaseNodeDisplay]]:
         NodeDisplayType = TypeVar("NodeDisplayType", bound=BaseNodeDisplay)
 
-        def decorator(inner_cls: Type[NodeDisplayType]) -> Type[NodeDisplayType]:
-            node_class = inner_cls.infer_node_class()
-            wrapped_node_class = cast(Type[BaseNode], node_class.__wrapped_node__)
+        def decorator(wrapped_node_display_class: Type[NodeDisplayType]) -> Type[NodeDisplayType]:
+            node_class = wrapped_node_display_class.infer_node_class()
+            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]
@@ -58,36 +63,51 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
                 for key, kwarg in kwargs.items():
                     ns[key] = kwarg
 
-                if "node_id" not in kwargs:
-                    ns["node_id"] = uuid4_from_hash(node_class.__qualname__)
+                ns["node_id"] = node_id or uuid4_from_hash(node_class.__qualname__)
 
             AdornmentDisplay = types.new_class(
                 re.sub(r"^Base", "", cls.__name__), bases=(BaseAdornmentDisplay,), exec_body=exec_body
             )
 
-            setattr(inner_cls, "__adorned_by__", AdornmentDisplay)
+            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(inner_cls)
+            original_base_node_display = get_original_base(wrapped_node_display_class)
             original_base_node_display.__args__ = (wrapped_node_class,)
-            inner_cls._node_display_registry[wrapped_node_class] = inner_cls
-            inner_cls.__annotate_node__()
+            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
-            old_outputs = list(inner_cls.output_display.keys())
-            for old_output in old_outputs:
-                new_output = getattr(wrapped_node_class.Outputs, old_output.name)
-                inner_cls.output_display[new_output] = inner_cls.output_display.pop(old_output)
+            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(inner_cls.port_displays.keys())
+            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)
-                inner_cls.port_displays[new_port] = inner_cls.port_displays.pop(old_port)
+                wrapped_node_display_class.port_displays[new_port] = wrapped_node_display_class.port_displays.pop(
+                    old_port
+                )
 
-            return inner_cls
+            return wrapped_node_display_class
 
         return decorator

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

@@ -19,6 +19,7 @@ _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
 class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType], Generic[_InlinePromptNodeType]):
     output_id: ClassVar[Optional[UUID]] = None
     array_output_id: ClassVar[Optional[UUID]] = None
+    json_output_id: ClassVar[Optional[UUID]] = None
     prompt_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(

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

@@ -21,20 +21,18 @@ class Inputs(BaseInputs):
     input: str
 
 
-@RetryNode.wrap(max_attempts=3)
-class InnerRetryGenericNode(BaseNode):
-    input = Inputs.input
-
-    class Outputs(BaseOutputs):
-        output: str
-
+def test_serialize_node__retry(serialize_node):
+    @RetryNode.wrap(max_attempts=3)
+    class InnerRetryGenericNode(BaseNode):
+        input = Inputs.input
 
-@BaseRetryNodeDisplay.wrap(max_attempts=3)
-class InnerRetryGenericNodeDisplay(BaseNodeDisplay[InnerRetryGenericNode]):
-    pass
+        class Outputs(BaseOutputs):
+            output: str
 
+    @BaseRetryNodeDisplay.wrap(max_attempts=3)
+    class InnerRetryGenericNodeDisplay(BaseNodeDisplay[InnerRetryGenericNode]):
+        pass
 
-def test_serialize_node__retry(serialize_node):
     input_id = uuid4()
     serialized_node = serialize_node(
         node_class=InnerRetryGenericNode,
@@ -49,8 +47,8 @@ def test_serialize_node__retry(serialize_node):
     )
     assert not DeepDiff(
         {
-            "id": "f2a95e79-7d4b-47ad-b986-4f648297ec65",
-            "label": "InnerRetryGenericNode",
+            "id": "188b50aa-e518-4b7b-a5e0-e2585fb1d7b5",
+            "label": "test_serialize_node__retry.<locals>.InnerRetryGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -66,8 +64,8 @@ def test_serialize_node__retry(serialize_node):
                     "test_adornments_serialization",
                 ],
             },
-            "trigger": {"id": "af9ba01c-4cde-4632-9aa1-7673b42e7bd8", "merge_behavior": "AWAIT_ATTRIBUTES"},
-            "ports": [{"id": "c2ecc6c0-f353-4495-9b93-a61a47248556", "name": "default", "type": "DEFAULT"}],
+            "trigger": {"id": "d38a83bf-23d1-4f9d-a875-a08dc27cf397", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "078650c9-f775-4cd0-a08c-23af9983a361", "name": "default", "type": "DEFAULT"}],
             "adornments": [
                 {
                     "id": "5be7d260-74f7-4734-b31b-a46a94539586",
@@ -102,13 +100,13 @@ def test_serialize_node__retry(serialize_node):
             ],
             "attributes": [
                 {
-                    "id": "c363daa7-9482-4c0e-aee8-faa080602ee3",
+                    "id": "278df25e-58b5-43c3-b346-cf6444d893a5",
                     "name": "input",
                     "value": {"type": "WORKFLOW_INPUT", "input_variable_id": str(input_id)},
                 }
             ],
             "outputs": [
-                {"id": "8aaf6cd8-3fa5-4f17-a60f-ec7da5ec6498", "name": "output", "type": "STRING", "value": None}
+                {"id": "dc89dc0d-c0bd-47fd-88aa-ec7b262aa2f1", "name": "output", "type": "STRING", "value": None}
             ],
         },
         serialized_node,
@@ -136,20 +134,18 @@ def test_serialize_node__retry__no_display():
     assert exec_config is not None
 
 
-@TryNode.wrap()
-class InnerTryGenericNode(BaseNode):
-    input = Inputs.input
-
-    class Outputs(BaseOutputs):
-        output: str
-
+def test_serialize_node__try(serialize_node):
+    @TryNode.wrap()
+    class InnerTryGenericNode(BaseNode):
+        input = Inputs.input
 
-@BaseTryNodeDisplay.wrap()
-class InnerTryGenericNodeDisplay(BaseNodeDisplay[InnerTryGenericNode]):
-    pass
+        class Outputs(BaseOutputs):
+            output: str
 
+    @BaseTryNodeDisplay.wrap()
+    class InnerTryGenericNodeDisplay(BaseNodeDisplay[InnerTryGenericNode]):
+        pass
 
-def test_serialize_node__try(serialize_node):
     input_id = uuid4()
     serialized_node = serialize_node(
         base_class=BaseNodeVellumDisplay,
@@ -163,7 +159,7 @@ def test_serialize_node__try(serialize_node):
     assert not DeepDiff(
         {
             "id": str(InnerTryGenericNode.__wrapped_node__.__id__),
-            "label": "InnerTryGenericNode",
+            "label": "test_serialize_node__try.<locals>.InnerTryGenericNode",
             "type": "GENERIC",
             "display_data": {"position": {"x": 0.0, "y": 0.0}},
             "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
@@ -179,8 +175,8 @@ def test_serialize_node__try(serialize_node):
                     "test_adornments_serialization",
                 ],
             },
-            "trigger": {"id": "741f7f75-e921-47a9-8c05-9e66640d0866", "merge_behavior": "AWAIT_ATTRIBUTES"},
-            "ports": [{"id": "1b8f8ab5-a656-4015-926c-80655bbd9cb8", "name": "default", "type": "DEFAULT"}],
+            "trigger": {"id": "16bc1522-c408-47ad-9a22-0ef136384abf", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "8d25f244-4b12-4f8b-b202-8948698679a0", "name": "default", "type": "DEFAULT"}],
             "adornments": [
                 {
                     "id": "3344083c-a32c-4a32-920b-0fb5093448fa",
@@ -200,13 +196,13 @@ def test_serialize_node__try(serialize_node):
             ],
             "attributes": [
                 {
-                    "id": "4d8b4c2c-4f92-4c7a-abf0-b9c88a15a790",
+                    "id": "51aa0077-4060-496b-8e2e-e79d56ee6a32",
                     "name": "input",
                     "value": {"type": "WORKFLOW_INPUT", "input_variable_id": str(input_id)},
                 }
             ],
             "outputs": [
-                {"id": "63ba929b-bf79-44ee-bd1f-d259dbe8d48e", "name": "output", "type": "STRING", "value": None}
+                {"id": "ce9f8b86-6d26-4c03-8bfa-a31aa2cd97f1", "name": "output", "type": "STRING", "value": None}
             ],
         },
         serialized_node,