vellum-ai 0.14.16__py3-none-any.whl → 0.14.18__py3-none-any.whl

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

@@ -54,7 +54,7 @@ def create_node_input_value_pointer_rules(
 
     if isinstance(value, BaseDescriptor):
         if isinstance(value, NodeReference):
-            if not value.instance:
+            if value.instance is None:
                 raise ValueError(f"Expected NodeReference {value.name} to have an instance")
             value = cast(BaseDescriptor, value.instance)
 

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

@@ -85,7 +85,7 @@ def test_serialize_workflow():
                     },
                 },
                 {
-                    "id": "bb712499-4265-4d71-bc14-0c3d8ca6a7de",
+                    "id": "e6c9d062-1a4f-4d4d-a258-59b75c6062fc",
                     "key": "info",
                     "value": {
                         "rules": [

vellum_ee/workflows/display/vellum.py

@@ -1,11 +1,10 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from uuid import UUID
 from typing import List, Literal, Optional, Union
 
 from pydantic import Field
 
-from vellum import PromptParameters, VellumVariable, VellumVariableType
+from vellum import VellumVariableType
 from vellum.client.types.array_vellum_value import ArrayVellumValue
 from vellum.client.types.vellum_value import VellumValue
 from vellum.core import UniversalBaseModel
@@ -138,24 +137,6 @@ class WorkflowOutputVellumDisplay(WorkflowOutputVellumDisplayOverrides):
     pass
 
 
-class WorkflowNodeType(str, Enum):
-    PROMPT = "PROMPT"
-    TEMPLATING = "TEMPLATING"
-    NOTE = "NOTE"
-    CODE_EXECUTION = "CODE_EXECUTION"
-    METRIC = "METRIC"
-    SEARCH = "SEARCH"
-    WEBHOOK = "WEBHOOK"
-    MERGE = "MERGE"
-    CONDITIONAL = "CONDITIONAL"
-    API = "API"
-    ENTRYPOINT = "ENTRYPOINT"
-    TERMINAL = "TERMINAL"
-    SUBWORKFLOW = "SUBWORKFLOW"
-    MAP = "MAP"
-    ERROR = "ERROR"
-
-
 class ConstantValuePointer(UniversalBaseModel):
     type: Literal["CONSTANT_VALUE"] = "CONSTANT_VALUE"
     data: VellumValue
@@ -220,131 +201,3 @@ class NodeInput(UniversalBaseModel):
     id: str
     key: str
     value: NodeInputValuePointer
-
-
-class BaseWorkflowNode(UniversalBaseModel):
-    id: str
-    inputs: List[NodeInput]
-    type: str
-    display_data: Optional[NodeDisplayData] = None
-    base: CodeResourceDefinition
-    definition: CodeResourceDefinition
-
-
-class EntrypointNodeData(UniversalBaseModel):
-    source_handle_id: str
-
-
-class EntrypointNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.ENTRYPOINT] = WorkflowNodeType.ENTRYPOINT
-    data: EntrypointNodeData
-
-
-class PromptTemplateBlockData(UniversalBaseModel):
-    version: Literal[1] = 1
-    # blocks: List[PromptBlockRequest]
-
-
-class PromptVersionExecConfig(UniversalBaseModel):
-    parameters: PromptParameters
-    input_variables: List[VellumVariable]
-    prompt_template_block_data: PromptTemplateBlockData
-
-
-class BasePromptNodeData(UniversalBaseModel):
-    label: str
-    output_id: str
-    error_output_id: Optional[str] = None
-    array_output_id: str
-    source_handle_id: str
-    target_handle_id: str
-
-
-class InlinePromptNodeData(BasePromptNodeData):
-    variant: Literal["INLINE"] = "INLINE"
-    exec_config: PromptVersionExecConfig
-    ml_model_name: str
-
-
-class DeploymentPromptNodeData(BasePromptNodeData):
-    variant: Literal["DEPLOYMENT"] = "DEPLOYMENT"
-    deployment_id: str
-    release_tag: str
-
-
-PromptNodeData = Union[
-    InlinePromptNodeData,
-    DeploymentPromptNodeData,
-]
-
-
-class PromptNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.PROMPT] = WorkflowNodeType.PROMPT
-    data: PromptNodeData
-
-
-class SearchNodeData(UniversalBaseModel):
-    label: str
-
-    results_output_id: str
-    text_output_id: str
-    error_output_id: Optional[str] = None
-
-    source_handle_id: str
-    target_handle_id: str
-
-    query_node_input_id: str
-    document_index_node_input_id: str
-    weights_node_input_id: str
-    limit_node_input_id: str
-    separator_node_input_id: str
-    result_merging_enabled_node_input_id: str
-    external_id_filters_node_input_id: str
-    metadata_filters_node_input_id: str
-
-
-class SearchNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.SEARCH] = WorkflowNodeType.SEARCH
-    data: SearchNodeData
-
-
-class FinalOutputNodeData(UniversalBaseModel):
-    label: str
-    name: str
-    target_handle_id: str
-    output_id: str
-    output_type: VellumVariableType
-    node_input_id: str
-
-
-class FinalOutputNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.TERMINAL] = WorkflowNodeType.TERMINAL
-    data: FinalOutputNodeData
-
-
-WorkflowNode = Union[
-    EntrypointNode,
-    PromptNode,
-    SearchNode,
-    FinalOutputNode,
-]
-
-
-class WorkflowEdge(UniversalBaseModel):
-    id: str
-    source_node_id: str
-    source_handle_id: str
-    target_node_id: str
-    target_handle_id: str
-
-
-class WorkflowRawData(UniversalBaseModel):
-    nodes: List[WorkflowNode]
-    edges: List[WorkflowEdge]
-    display_data: Optional[WorkflowDisplayData] = None
-
-
-class WorkflowVersionExecConfig(UniversalBaseModel):
-    workflow_raw_data: WorkflowRawData
-    input_variables: List[VellumVariable]
-    output_variables: List[VellumVariable]

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -429,7 +429,7 @@ class BaseWorkflowDisplay(
                     )
                     subworkflow_display_context = subworkflow_display.get_event_display_context()
 
-            node_event_displays[str(node_id)] = NodeEventDisplayContext(
+            node_event_displays[node_id] = NodeEventDisplayContext(
                 input_display=input_display,
                 output_display=output_display,
                 port_display=port_display_meta,

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -5,9 +5,11 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.inline_subworkflow_node.node import InlineSubworkflowNode
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
+from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.retry_node import BaseRetryNodeDisplay
+from vellum_ee.workflows.display.nodes.vellum.try_node import BaseTryNodeDisplay
 from vellum_ee.workflows.display.vellum import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -112,13 +114,36 @@ def test_get_event_display_context__node_display_filled_without_base_display():
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the node display should be included
-    assert str(StartNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(StartNode.__id__)]
+    assert StartNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[StartNode.__id__]
 
     # AND so should their output ids
     assert StartNode.__output_ids__ == node_event_display.output_display
 
 
+def test_get_event_display_context__node_display_filled_without_output_display():
+    # GIVEN a simple workflow
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+    class MyWorkflow(BaseWorkflow):
+        graph = StartNode
+
+    class StartNodeDisplay(BaseNodeDisplay[StartNode]):
+        pass
+
+    # WHEN we gather the event display context
+    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+
+    # THEN the node display should be included
+    assert StartNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[StartNode.__id__]
+
+    # AND so should their output ids
+    assert node_event_display.output_display.keys() == {"foo"}
+
+
 def test_get_event_display_context__node_display_to_include_subworkflow_display():
     # GIVEN a simple workflow
     class InnerNode(BaseNode):
@@ -138,26 +163,40 @@ def test_get_event_display_context__node_display_to_include_subworkflow_display(
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(SubworkflowNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(SubworkflowNode.__id__)]
+    assert SubworkflowNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[SubworkflowNode.__id__]
 
     assert node_event_display.subworkflow_display is not None
-    assert str(InnerNode.__id__) in node_event_display.subworkflow_display.node_displays
-
-
-def test_get_event_display_context__node_display_for_adornment_nodes():
-    # GIVEN a simple workflow with a retry node adornment
-    @RetryNode.wrap(max_attempts=4)
+    assert InnerNode.__id__ in node_event_display.subworkflow_display.node_displays
+
+
+@pytest.mark.parametrize(
+    ["AdornmentNode", "AdornmentNodeDisplay", "expected_adornment_output_names"],
+    [
+        [RetryNode, BaseRetryNodeDisplay, {"foo"}],
+        [TryNode, BaseTryNodeDisplay, {"foo", "error"}],
+    ],
+    ids=["retry_node", "try_node"],
+)
+def test_get_event_display_context__node_display_for_adornment_nodes(
+    AdornmentNode,
+    AdornmentNodeDisplay,
+    expected_adornment_output_names,
+):
+    # GIVEN a simple workflow with an adornment
+    @AdornmentNode.wrap()
     class MyNode(BaseNode):
-        pass
+        class Outputs(BaseNode.Outputs):
+            foo: str
 
     class MyWorkflow(BaseWorkflow):
         graph = MyNode
 
     # AND a display class for the node
+    adornment_node_id = uuid4()
     inner_node_id = uuid4()
 
-    @BaseRetryNodeDisplay.wrap()
+    @AdornmentNodeDisplay.wrap(node_id=adornment_node_id)
     class MyNodeDisplay(BaseNodeDisplay[MyNode]):
         node_id = inner_node_id
 
@@ -165,10 +204,15 @@ def test_get_event_display_context__node_display_for_adornment_nodes():
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(MyNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(MyNode.__id__)]
+    assert adornment_node_id in display_context.node_displays
+    node_event_display = display_context.node_displays[adornment_node_id]
     assert node_event_display.subworkflow_display is not None
-    assert str(inner_node_id) in node_event_display.subworkflow_display.node_displays
+    assert inner_node_id in node_event_display.subworkflow_display.node_displays
+
+    # AND the inner node should have the correct outputs
+    inner_node_display = node_event_display.subworkflow_display.node_displays[inner_node_id]
+    assert inner_node_display.output_display.keys() == {"foo"}
+    assert node_event_display.output_display.keys() == expected_adornment_output_names
 
 
 def test_get_event_display_context__templating_node_input_display():
@@ -187,7 +231,7 @@ def test_get_event_display_context__templating_node_input_display():
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(MyNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(MyNode.__id__)]
+    assert MyNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[MyNode.__id__]
 
     assert node_event_display.input_display.keys() == {"inputs.foo"}

vellum_ee/workflows/tests/test_display_meta.py

@@ -1,7 +1,7 @@
 import pytest
 import os
 import sys
-from uuid import UUID, uuid4
+from uuid import uuid4
 
 from vellum.workflows import BaseWorkflow
 from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
@@ -35,22 +35,22 @@ def test_base_class_dynamic_import(files):
         "node_displays": {
             "533c6bd8-6088-4abc-a168-8c1758abcd33": {
                 "input_display": {
-                    "example_var_1": UUID("a0d1d7cf-242a-4bd9-a437-d308a7ced9b3"),
-                    "template": UUID("f97d721a-e685-498e-90c3-9c3d9358fdad"),
+                    "example_var_1": "a0d1d7cf-242a-4bd9-a437-d308a7ced9b3",
+                    "template": "f97d721a-e685-498e-90c3-9c3d9358fdad",
                 },
-                "output_display": {"result": UUID("423bc529-1a1a-4f72-af4d-cbdb5f0a5929")},
-                "port_display": {"default": UUID("afda9a19-0618-42e1-9b63-5d0db2a88f62")},
+                "output_display": {"result": "423bc529-1a1a-4f72-af4d-cbdb5f0a5929"},
+                "port_display": {"default": "afda9a19-0618-42e1-9b63-5d0db2a88f62"},
                 "subworkflow_display": None,
             },
             "f3ef4b2b-fec9-4026-9cc6-e5eac295307f": {
-                "input_display": {"node_input": UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")},
-                "output_display": {"value": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
+                "input_display": {"node_input": "fe6cba85-2423-4b5e-8f85-06311a8be5fb"},
+                "output_display": {"value": "5469b810-6ea6-4362-9e79-e360d44a1405"},
                 "port_display": {},
                 "subworkflow_display": None,
             },
         },
-        "workflow_inputs": {"input_value": UUID("2268a996-bd17-4832-b3ff-f5662d54b306")},
-        "workflow_outputs": {"final_output": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
+        "workflow_inputs": {"input_value": "2268a996-bd17-4832-b3ff-f5662d54b306"},
+        "workflow_outputs": {"final_output": "5469b810-6ea6-4362-9e79-e360d44a1405"},
     }
     assert display_meta
-    assert display_meta.dict() == expected_result
+    assert display_meta.model_dump(mode="json") == expected_result