vellum-ai 1.5.0__py3-none-any.whl → 1.5.1__py3-none-any.whl

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -323,12 +323,12 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         )
         return node_definition
 
-    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
         explicit_display = self.output_display.get(output)
         if explicit_display:
-            return self._node, explicit_display
+            return explicit_display
 
-        return (self._node, NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name))
+        return 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)
@@ -347,6 +347,9 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         if default_port in port_displays:
             return port_displays[default_port].id
 
+        if default_port:
+            return self.get_node_port_display(default_port).id
+
         first_port = next((port for port in unadorned_node.Ports), None)
         if not first_port:
             raise ValueError(f"Node {self._node.__name__} must have at least one port.")

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

@@ -1,8 +1,7 @@
 from uuid import UUID
-from typing import ClassVar, Dict, Generic, Optional, TypeVar, cast
+from typing import ClassVar, Dict, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable import APINode
-from vellum.workflows.references.output import OutputReference
 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.utils import raise_if_descriptor
@@ -178,11 +177,9 @@ class BaseAPINodeDisplay(BaseNodeDisplay[_APINodeType], Generic[_APINodeType]):
         ]
         inputs.extend(additional_header_inputs)
 
-        _, text_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
-        _, json_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.json)]
-        _, status_code_output_display = display_context.global_node_output_displays[
-            cast(OutputReference, node.Outputs.status_code)
-        ]
+        text_output_display = self.get_node_output_display(node.Outputs.text)
+        json_output_display = self.get_node_output_display(node.Outputs.json)
+        status_code_output_display = self.get_node_output_display(node.Outputs.status_code)
 
         serialized_node: JsonObject = {
             "id": str(node_id),

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

@@ -7,6 +7,8 @@ from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes import InlineSubworkflowNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.types.core import JsonObject
+from vellum.workflows.workflows.base import BaseWorkflow
+from vellum_ee.workflows.display.exceptions import NodeValidationError
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -31,14 +33,25 @@ class BaseInlineSubworkflowNodeDisplay(
         node = self._node
         node_id = self.node_id
 
-        node_inputs, workflow_inputs = self._generate_node_and_workflow_inputs(node_id, node, display_context)
+        subworkflow_class = raise_if_descriptor(node.subworkflow)
+        if subworkflow_class is None:
+            display_context.add_error(
+                NodeValidationError(
+                    "InlineSubworkflowNode requires a subworkflow to be defined",
+                    node_class_name=node.__class__.__name__,
+                )
+            )
+            subworkflow_class = BaseWorkflow
 
+        node_inputs, workflow_inputs = self._generate_node_and_workflow_inputs(
+            node_id, node, display_context, subworkflow_class
+        )
         subworkflow_display = get_workflow_display(
             base_display_class=display_context.workflow_display_class,
-            workflow_class=raise_if_descriptor(node.subworkflow),
+            workflow_class=subworkflow_class,
             parent_display_context=display_context,
         )
-        workflow_outputs = self._generate_workflow_outputs(node, subworkflow_display.display_context)
+        workflow_outputs = self._generate_workflow_outputs(node, subworkflow_display.display_context, subworkflow_class)
         serialized_subworkflow = subworkflow_display.serialize()
 
         return {
@@ -63,8 +76,8 @@ class BaseInlineSubworkflowNodeDisplay(
         node_id: UUID,
         node: Type[InlineSubworkflowNode],
         display_context: WorkflowDisplayContext,
+        subworkflow: Type[BaseWorkflow],
     ) -> Tuple[List[NodeInput], List[VellumVariable]]:
-        subworkflow = raise_if_descriptor(node.subworkflow)
         subworkflow_inputs_class = subworkflow.get_inputs_class()
         subworkflow_inputs = raise_if_descriptor(node.subworkflow_inputs)
 
@@ -115,9 +128,10 @@ class BaseInlineSubworkflowNodeDisplay(
         self,
         node: Type[InlineSubworkflowNode],
         display_context: WorkflowDisplayContext,
+        subworkflow: Type[BaseWorkflow],
     ) -> List[VellumVariable]:
         workflow_outputs: List[VellumVariable] = []
-        for output_descriptor in raise_if_descriptor(node.subworkflow).Outputs:  # type: ignore[union-attr]
+        for output_descriptor in subworkflow.Outputs:  # type: ignore[union-attr]
             workflow_output_display = display_context.workflow_output_displays[output_descriptor]
             output_type = infer_vellum_variable_type(output_descriptor)
             workflow_outputs.append(

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

@@ -1,7 +1,6 @@
 import inspect
-from typing import Any, Generic, Tuple, Type, TypeVar
+from typing import Any, Generic, TypeVar
 
-from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
 from vellum.workflows.references.output import OutputReference
@@ -66,7 +65,7 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
 
         return serialized_node
 
-    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
         inner_node = self._node.__wrapped_node__
         if not inner_node:
             return super().get_node_output_display(output)

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

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from uuid import UUID
-from typing import Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
 
 from vellum import (
     MetadataFilterConfigRequest,
@@ -9,7 +9,6 @@ from vellum import (
 )
 from vellum.workflows.nodes.displayable.bases.types import MetadataLogicalCondition, MetadataLogicalConditionGroup
 from vellum.workflows.nodes.displayable.search_node import SearchNode
-from vellum.workflows.references import OutputReference
 from vellum.workflows.types.core import JsonArray, JsonObject
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
@@ -49,10 +48,8 @@ class BaseSearchNodeDisplay(BaseNodeDisplay[_SearchNodeType], Generic[_SearchNod
         node_id = self.node_id
         node_inputs = self._generate_search_node_inputs(node_id, node, display_context)
 
-        _, results_output_display = display_context.global_node_output_displays[
-            cast(OutputReference, node.Outputs.results)
-        ]
-        _, text_output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
+        results_output_display = display_context.global_node_output_displays[node.Outputs.results]
+        text_output_display = display_context.global_node_output_displays[node.Outputs.text]
 
         return {
             "id": str(node_id),

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

@@ -50,7 +50,7 @@ class BaseTemplatingNodeDisplay(BaseNodeDisplay[_TemplatingNodeType], Generic[_T
         # Misc type ignore is due to `node.Outputs` being generic
         # https://app.shortcut.com/vellum/story/4784
         output_descriptor = node.Outputs.result  # type: ignore [misc]
-        _, output_display = display_context.global_node_output_displays[output_descriptor]
+        output_display = display_context.global_node_output_displays[output_descriptor]
         inferred_output_type = primitive_type_to_vellum_variable_type(output_descriptor)
 
         return {

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

@@ -121,9 +121,8 @@ def test_create_node_input_value_pointer_rules(
                 ),
             },
             global_node_output_displays={
-                MyNodeA.Outputs.output: (
-                    MyNodeA,
-                    NodeOutputDisplay(id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"),
+                MyNodeA.Outputs.output: NodeOutputDisplay(
+                    id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"
                 ),
             },
             global_node_displays={

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

@@ -1,8 +1,7 @@
 import inspect
 from uuid import UUID
-from typing import Any, ClassVar, Generic, Optional, Tuple, Type, TypeVar
+from typing import Any, ClassVar, Generic, Optional, TypeVar
 
-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
 from vellum.workflows.references.output import OutputReference
@@ -76,7 +75,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
 
         return serialized_node
 
-    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
         inner_node = self._node.__wrapped_node__
         if not inner_node:
             return super().get_node_output_display(output)
@@ -84,7 +83,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
         node_display_class = get_node_display_class(inner_node)
         node_display = node_display_class()
         if output.name == "error":
-            return inner_node, NodeOutputDisplay(
+            return NodeOutputDisplay(
                 id=self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id"),
                 name="error",
             )

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

@@ -329,7 +329,7 @@ def test_serialize_node__node_output(serialize_node):
         global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=workflow_input_id)},
         global_node_displays={NodeWithOutput: NodeWithOutputDisplay()},
         global_node_output_displays={
-            NodeWithOutput.Outputs.output: (NodeWithOutput, NodeOutputDisplay(id=node_output_id, name="output"))
+            NodeWithOutput.Outputs.output: NodeOutputDisplay(id=node_output_id, name="output")
         },
     )
 
@@ -521,14 +521,8 @@ def test_serialize_node__coalesce(serialize_node):
             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"),
-            ),
+            CoalesceNodeA.Outputs.output: NodeOutputDisplay(id=coalesce_node_a_output_id, name="output"),
+            CoalesceNodeB.Outputs.output: NodeOutputDisplay(id=coalesce_node_b_output_id, name="output"),
         },
     )
 
@@ -602,7 +596,7 @@ def test_serialize_node__dataclass_with_node_output_reference(serialize_node):
         node_class=GenericNodeWithDataclass,
         global_node_displays={NodeWithOutput: NodeWithOutputDisplay()},
         global_node_output_displays={
-            NodeWithOutput.Outputs.result: (NodeWithOutput, NodeOutputDisplay(id=node_output_id, name="result"))
+            NodeWithOutput.Outputs.result: NodeOutputDisplay(id=node_output_id, name="result")
         },
     )
 
@@ -634,7 +628,7 @@ def test_serialize_node__pydantic_with_node_output_reference(serialize_node):
         node_class=GenericNodeWithPydantic,
         global_node_displays={NodeWithOutput: NodeWithOutputDisplay()},
         global_node_output_displays={
-            NodeWithOutput.Outputs.result: (NodeWithOutput, NodeOutputDisplay(id=node_output_id, name="result"))
+            NodeWithOutput.Outputs.result: NodeOutputDisplay(id=node_output_id, name="result")
         },
     )
 

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

@@ -127,7 +127,7 @@ def test_serialize_node__node_output_reference(serialize_node):
         global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=workflow_input_id)},
         global_node_displays={NodeWithOutput: NodeWithOutputDisplay()},
         global_node_output_displays={
-            NodeWithOutput.Outputs.output: (NodeWithOutput, NodeOutputDisplay(id=node_output_id, name="output"))
+            NodeWithOutput.Outputs.output: NodeOutputDisplay(id=node_output_id, name="output")
         },
     )
 

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

@@ -292,7 +292,7 @@ def test_serialize_node__node_output_reference(serialize_node):
         global_workflow_input_displays={Inputs.input: WorkflowInputsDisplay(id=workflow_input_id)},
         global_node_displays={NodeWithOutput: NodeWithOutputDisplay()},
         global_node_output_displays={
-            NodeWithOutput.Outputs.output: (NodeWithOutput, NodeOutputDisplay(id=node_output_id, name="output"))
+            NodeWithOutput.Outputs.output: NodeOutputDisplay(id=node_output_id, name="output")
         },
     )
 

vellum_ee/workflows/display/types.py

@@ -28,7 +28,7 @@ if TYPE_CHECKING:
 WorkflowInputsDisplays = Dict[WorkflowInputReference, WorkflowInputsDisplay]
 StateValueDisplays = Dict[StateValueReference, StateValueDisplay]
 NodeDisplays = Dict[Type[BaseNode], BaseNodeDisplay]
-NodeOutputDisplays = Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]]
+NodeOutputDisplays = Dict[OutputReference, NodeOutputDisplay]
 EntrypointDisplays = Dict[Type[BaseNode], EntrypointDisplay]
 WorkflowOutputDisplays = Dict[BaseDescriptor, WorkflowOutputDisplay]
 EdgeDisplays = Dict[Tuple[Port, Type[BaseNode]], EdgeDisplay]
@@ -51,12 +51,12 @@ class WorkflowDisplayContext:
     workflow_output_displays: WorkflowOutputDisplays = field(default_factory=dict)
     edge_displays: EdgeDisplays = field(default_factory=dict)
     port_displays: PortDisplays = field(default_factory=dict)
+    dry_run: bool = False
     _errors: List[Exception] = field(default_factory=list)
     _invalid_nodes: List[Type[BaseNode]] = field(default_factory=list)
-    _dry_run: bool = False
 
     def add_error(self, error: Exception, node: Optional[Type[BaseNode]] = None) -> None:
-        if self._dry_run:
+        if self.dry_run:
             self._errors.append(error)
             return
 

vellum_ee/workflows/display/utils/expressions.py

@@ -41,7 +41,9 @@ 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.expressions.parse_json import ParseJsonExpression
+from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
+from vellum.workflows.nodes.utils import get_unadorned_node
 from vellum.workflows.references.constant import ConstantValueReference
 from vellum.workflows.references.environment_variable import EnvironmentVariableReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
@@ -296,12 +298,17 @@ def serialize_value(executable_id: UUID, display_context: "WorkflowDisplayContex
         }
 
     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]
+        output_display = display_context.global_node_output_displays[value]
+
+        upstream_node_class = value.outputs_class.__parent_class__
+        if not issubclass(upstream_node_class, BaseNode):
+            raise ValueError(f"Output references must be to a node, not {upstream_node_class}")
+        unadorned_upstream_node_class = get_unadorned_node(upstream_node_class)
+        upstream_node = display_context.global_node_displays[unadorned_upstream_node_class]
 
         return {
             "type": "NODE_OUTPUT",
-            "node_id": str(upstream_node_display.node_id),
+            "node_id": str(upstream_node.node_id),
             "node_output_id": str(output_display.id),
         }
 

vellum_ee/workflows/display/utils/vellum.py

@@ -8,6 +8,7 @@ from vellum.client.types.vellum_variable_type import VellumVariableType
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
+from vellum.workflows.nodes.utils import get_unadorned_node
 from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.lazy import LazyReference
@@ -108,8 +109,14 @@ def create_node_input_value_pointer_rule(
 
             raise ValueError(f"Reference to outputs '{value.outputs_class.__qualname__}' is invalid.")
 
-        upstream_node, output_display = display_context.global_node_output_displays[value]
-        upstream_node_display = display_context.global_node_displays[upstream_node]
+        output_display = display_context.global_node_output_displays[value]
+
+        upstream_node_class = value.outputs_class.__parent_class__
+        if not issubclass(upstream_node_class, BaseNode):
+            raise ValueError(f"Output references must be to a node, not {upstream_node_class}")
+        unadorned_upstream_node_class = get_unadorned_node(upstream_node_class)
+        upstream_node_display = display_context.global_node_displays[unadorned_upstream_node_class]
+
         return NodeOutputPointer(
             data=NodeOutputData(node_id=str(upstream_node_display.node_id), output_id=str(output_display.id)),
         )

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -479,7 +479,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         self,
         node: Type[BaseNode],
         node_display: BaseNodeDisplay,
-        node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]],
+        node_output_displays: Dict[OutputReference, NodeOutputDisplay],
     ):
         """This method recursively adds nodes wrapped in decorators to the node_output_displays dictionary."""
 
@@ -624,7 +624,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             edge_displays=edge_displays,
             port_displays=port_displays,
             workflow_display_class=self.__class__,
-            _dry_run=self._dry_run,
+            dry_run=self._dry_run,
         )
 
     def _generate_workflow_meta_display(self) -> WorkflowMetaDisplay: