vellum-ai 0.13.8__tar.gz → 0.13.9__tar.gz

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.13.8
+Version: 0.13.9
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -76,12 +76,12 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         )
         return node_definition
 
-    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
+    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
         explicit_display = self.output_display.get(output)
         if explicit_display:
-            return explicit_display
+            return self._node, explicit_display
 
-        return NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name)
+        return (self._node, NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name))
 
     def get_node_port_display(self, port: Port) -> PortDisplay:
         overrides = self.port_displays.get(port)

vellum_ai-0.13.9/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py

@@ -0,0 +1,63 @@
+from typing import Any, Dict, cast
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
+from vellum.workflows.nodes.core.try_node.node import TryNode
+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
+
+
+def test_try_node_display__serialize_with_error_output() -> None:
+    # GIVEN a Base Node wrapped with a TryNode
+    @TryNode.wrap()
+    class MyNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            hello = "world"
+
+    # AND a displayable node referencing
+    class OtherNode(TemplatingNode):
+        template = "{{ hello }} {{ error }}"
+        inputs = {
+            "hello": MyNode.Outputs.hello,
+            "error": MyNode.Outputs.error,
+        }
+
+    # AND a workflow referencing the two nodes
+    class MyWorkflow(BaseWorkflow):
+        graph = MyNode >> OtherNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
+
+    # THEN the correct inputs should be serialized on the node
+    serialized_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(OtherNode.__id__)
+    )
+    hello_input_value = next(input["value"] for input in serialized_node["inputs"] if input["key"] == "hello")
+    error_input_value = next(input["value"] for input in serialized_node["inputs"] if input["key"] == "error")
+    assert hello_input_value == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "data": {
+                    "node_id": str(MyNode.__wrapped_node__.__id__),
+                    "output_id": "c8fbe459-c9ee-4639-a82d-961180cf9411",
+                },
+                "type": "NODE_OUTPUT",
+            }
+        ],
+    }
+    assert error_input_value == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "data": {
+                    "node_id": str(MyNode.__wrapped_node__.__id__),
+                    "output_id": "efe6e307-3ea4-4862-a26f-4c4416bb4537",
+                },
+                "type": "NODE_OUTPUT",
+            }
+        ],
+    }

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -1,13 +1,12 @@
 import pytest
 from uuid import UUID, uuid4
-from typing import List, cast
+from typing import List
 
 from vellum.client.types.string_vellum_value import StringVellumValue
 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 OutputReference, WorkflowInputReference
 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
@@ -92,12 +91,12 @@ def test_create_node_input_value_pointer_rules(
                 entrypoint_node_display=NodeDisplayData(),
             ),
             global_workflow_input_displays={
-                cast(WorkflowInputReference, Inputs.example_workflow_input): WorkflowInputsVellumDisplayOverrides(
+                Inputs.example_workflow_input: WorkflowInputsVellumDisplayOverrides(
                     id=UUID("a154c29d-fac0-4cd0-ba88-bc52034f5470"),
                 ),
             },
             global_node_output_displays={
-                cast(OutputReference, MyNodeA.Outputs.output): (
+                MyNodeA.Outputs.output: (
                     MyNodeA,
                     NodeOutputDisplay(id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"),
                 ),

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/ee/vellum_ee/workflows/display/nodes/vellum/try_node.py

@@ -1,15 +1,17 @@
 from uuid import UUID
-from typing import Any, Callable, ClassVar, Generic, Optional, Type, TypeVar, cast
+from typing import Any, Callable, ClassVar, Generic, Optional, Tuple, Type, TypeVar, cast
 
 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.references.output import OutputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.utils import get_original_base
 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.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
@@ -22,9 +24,8 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
     def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
         node = self._node
 
-        try:
-            inner_node = get_wrapped_node(node)
-        except AttributeError:
+        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(
@@ -57,6 +58,22 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
 
         return serialized_node
 
+    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+        inner_node = self._node.__wrapped_node__
+        if not inner_node:
+            return super().get_node_output_display(output)
+
+        node_display_class = get_node_display_class(BaseNodeVellumDisplay, inner_node)
+        node_display = node_display_class()
+        if output.name == "error":
+            return inner_node, NodeOutputDisplay(
+                id=self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id"),
+                name="error",
+            )
+
+        inner_output = getattr(inner_node.Outputs, output.name)
+        return node_display.get_node_output_display(inner_output)
+
     @classmethod
     def wrap(cls, error_output_id: Optional[UUID] = None) -> Callable[..., Type["BaseTryNodeDisplay"]]:
         _error_output_id = error_output_id

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/ee/vellum_ee/workflows/display/nodes/vellum/utils.py

@@ -28,8 +28,7 @@ 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.nodes.utils import get_wrapped_node, has_wrapped_node
-from vellum.workflows.references import NodeReference, OutputReference
+from vellum.workflows.references import NodeReference
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule
@@ -56,15 +55,6 @@ def create_node_input(
     pointer_type: Optional[Type[NodeInputValuePointerRule]] = ConstantValuePointer,
 ) -> NodeInput:
     input_id = str(input_id) if input_id else str(uuid4_from_hash(f"{node_id}|{input_name}"))
-    if (
-        isinstance(value, OutputReference)
-        and value.outputs_class._node_class
-        and has_wrapped_node(value.outputs_class._node_class)
-    ):
-        wrapped_node = get_wrapped_node(value.outputs_class._node_class)
-        if wrapped_node._is_wrapped_node:
-            value = getattr(wrapped_node.Outputs, value.name)
-
     rules = create_node_input_value_pointer_rules(value, display_context, pointer_type=pointer_type)
     return NodeInput(
         id=input_id,

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py

@@ -685,8 +685,7 @@ def test_serialize_workflow__try_wrapped():
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
             },
         ],
-        final_output_nodes,
-        ignore_order=True,
+        sorted(final_output_nodes, key=lambda x: x["id"], reverse=True),
     )
 
     # AND each edge should be serialized correctly

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -11,7 +11,7 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.nodes.utils import get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum.workflows.types.core import JsonObject
@@ -138,19 +138,18 @@ class BaseWorkflowDisplay(
     ):
         """This method recursively adds nodes wrapped in decorators to the node_output_displays dictionary."""
 
+        inner_node = get_wrapped_node(node)
+        if inner_node:
+            inner_node_display = self._get_node_display(inner_node)
+            self._enrich_global_node_output_displays(inner_node, inner_node_display, node_output_displays)
+
         for node_output in node.Outputs:
             if node_output in node_output_displays:
                 continue
 
-            if has_wrapped_node(node):
-                inner_node = get_wrapped_node(node)
-                if inner_node._is_wrapped_node:
-                    inner_node_display = self._get_node_display(inner_node)
-                    self._enrich_global_node_output_displays(inner_node, inner_node_display, node_output_displays)
-
             # 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, node_display.get_node_output_display(node_output)
+            node_output_displays[node_output] = node_display.get_node_output_display(node_output)
 
     def _enrich_node_port_displays(
         self,
@@ -160,16 +159,15 @@ class BaseWorkflowDisplay(
     ):
         """This method recursively adds nodes wrapped in decorators to the port_displays dictionary."""
 
+        inner_node = get_wrapped_node(node)
+        if inner_node:
+            inner_node_display = self._get_node_display(inner_node)
+            self._enrich_node_port_displays(inner_node, inner_node_display, port_displays)
+
         for port in node.Ports:
             if port in port_displays:
                 continue
 
-            if has_wrapped_node(node):
-                inner_node = get_wrapped_node(node)
-                if inner_node._is_wrapped_node:
-                    inner_node_display = self._get_node_display(inner_node)
-                    self._enrich_node_port_displays(inner_node, inner_node_display, port_displays)
-
             port_displays[port] = node_display.get_node_port_display(port)
 
     def _get_node_display(self, node: Type[BaseNode]) -> NodeDisplayType:
@@ -216,13 +214,11 @@ class BaseWorkflowDisplay(
             global_node_displays[node] = node_display
 
             # Nodes wrapped in a decorator need to be in our node display dictionary for later retrieval
-            if has_wrapped_node(node):
-                inner_node = get_wrapped_node(node)
+            inner_node = get_wrapped_node(node)
+            if inner_node:
                 inner_node_display = self._get_node_display(inner_node)
-
-                if inner_node._is_wrapped_node:
-                    node_displays[inner_node] = inner_node_display
-                    global_node_displays[inner_node] = inner_node_display
+                node_displays[inner_node] = inner_node_display
+                global_node_displays[inner_node] = inner_node_display
 
             self._enrich_global_node_output_displays(node, node_display, global_node_output_displays)
             self._enrich_node_port_displays(node, node_display, port_displays)

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/ee/vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -7,7 +7,7 @@ from vellum.workflows.edges import Edge
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
-from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node, has_wrapped_node
+from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port
 from vellum.workflows.ports import Port
 from vellum.workflows.references import WorkflowInputReference
 from vellum.workflows.references.output import OutputReference
@@ -95,11 +95,8 @@ class VellumWorkflowDisplay(
         )
 
         # Add all the nodes in the workflow
-        for node, node_display in self.display_context.node_displays.items():
-            if getattr(node, "_is_wrapped_node") is True:
-                # Nodes that are wrapped or decorated by other nodes are not serialized here
-                # They are instead serialized by the wrapper node
-                continue
+        for node in self._workflow.get_nodes():
+            node_display = self.display_context.node_displays[node]
 
             try:
                 serialized_node = node_display.serialize(self.display_context)
@@ -295,13 +292,9 @@ class VellumWorkflowDisplay(
             else uuid4_from_hash(f"{self.workflow_id}|id|{entrypoint_node_id}")
         )
 
-        if has_wrapped_node(entrypoint):
-            wrapped_node = get_wrapped_node(entrypoint)
-            if wrapped_node._is_wrapped_node:
-                entrypoint = wrapped_node
-
-        target_node_id = node_displays[entrypoint].node_id
-        target_handle_id = node_displays[entrypoint].get_target_handle_id()
+        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()
 
         edge_display = self._generate_edge_display_from_source(
             entrypoint_node_id, source_handle_id, target_node_id, target_handle_id, overrides=edge_display_overrides

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vellum-ai"
-version = "0.13.8"
+version = "0.13.9"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/src/vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.13.8",
+            "X-Fern-SDK-Version": "0.13.9",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/src/vellum/plugins/vellum_mypy.py

@@ -19,6 +19,7 @@ from mypy.plugin import AttributeContext, ClassDefContext, FunctionSigContext, M
 from mypy.types import AnyType, CallableType, FunctionLike, Instance, Type as MypyType, TypeAliasType, UnionType
 
 TypeResolver = Callable[[str, List[MypyType]], MypyType]
+NODE_ATTRIBUTE_REGEX = r"^[a-z].*$"
 
 DESCRIPTOR_PATHS: list[tuple[str, str, str]] = [
     (
@@ -49,7 +50,7 @@ DESCRIPTOR_PATHS: list[tuple[str, str, str]] = [
     (
         "vellum.workflows.nodes.bases.base.BaseNode",
         "vellum.workflows.references.node.NodeReference",
-        r"^[a-z].*$",
+        NODE_ATTRIBUTE_REGEX,
     ),
 ]
 
@@ -239,10 +240,13 @@ class VellumMypyPlugin(Plugin):
         the original defined type, and the descriptor version of the type.
         """
 
-        for sym in ctx.cls.info.names.values():
+        for key, sym in ctx.cls.info.names.items():
             if not isinstance(sym.node, Var):
                 continue
 
+            if not re.match(NODE_ATTRIBUTE_REGEX, key):
+                continue
+
             type_ = sym.node.type
             if not type_:
                 continue

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/src/vellum/workflows/nodes/bases/base.py

@@ -222,7 +222,6 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
     state: StateType
     _context: WorkflowContext
     _inputs: MappingProxyType[NodeReference, Any]
-    _is_wrapped_node: bool = False
 
     class ExternalInputs(BaseInputs):
         __descriptor_class__ = ExternalInputReference

{vellum_ai-0.13.8 → vellum_ai-0.13.9}/src/vellum/workflows/nodes/utils.py

@@ -4,6 +4,7 @@ from types import ModuleType
 from typing import Any, Callable, Optional, Type, TypeVar
 
 from vellum.workflows.nodes import BaseNode
+from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.ports.port import Port
 from vellum.workflows.types.generics import NodeType
 
@@ -12,7 +13,7 @@ ADORNMENT_MODULE_NAME = "<adornment>"
 
 @cache
 def get_unadorned_node(node: Type[BaseNode]) -> Type[BaseNode]:
-    wrapped_node = getattr(node, "__wrapped_node__", None)
+    wrapped_node = get_wrapped_node(node)
     if wrapped_node is not None:
         return get_unadorned_node(wrapped_node)
 
@@ -28,22 +29,11 @@ def get_unadorned_port(port: Port) -> Port:
     return getattr(unadorned_node.Ports, port.name)
 
 
-@cache
-def get_wrapped_node(node: Type[NodeType]) -> Type[BaseNode]:
-    wrapped_node = getattr(node, "__wrapped_node__", None)
-    if wrapped_node is None:
-        raise AttributeError("Wrapped node not found")
-
-    return wrapped_node
-
+def get_wrapped_node(node: Type[NodeType]) -> Optional[Type[BaseNode]]:
+    if not issubclass(node, BaseAdornmentNode):
+        return None
 
-def has_wrapped_node(node: Type[NodeType]) -> bool:
-    try:
-        get_wrapped_node(node)
-    except AttributeError:
-        return False
-
-    return True
+    return node.__wrapped_node__
 
 
 AdornableNode = TypeVar("AdornableNode", bound=BaseNode)
@@ -57,8 +47,6 @@ def create_adornment(
         # https://app.shortcut.com/vellum/story/4116
         from vellum.workflows import BaseWorkflow
 
-        inner_cls._is_wrapped_node = True
-
         class Subworkflow(BaseWorkflow):
             graph = inner_cls