vellum-ai 0.12.14__py3-none-any.whl → 0.12.16__py3-none-any.whl

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

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -15,12 +12,7 @@ def test_serialize_workflow():
     workflow_display = get_workflow_display(
         base_display_class=VellumWorkflowDisplay, workflow_class=BasicInlineSubworkflowWorkflow
     )
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -156,7 +148,23 @@ def test_serialize_workflow():
                                 "bases": [],
                             },
                         },
-                        {"type": "MOCKED"},
+                        {
+                            "id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
+                            "label": "StartNode",
+                            "type": "GENERIC",
+                            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                            "definition": {
+                                "name": "StartNode",
+                                "module": ["tests", "workflows", "basic_inline_subworkflow", "workflow"],
+                                "bases": [
+                                    {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}
+                                ],
+                            },
+                            "trigger": {"id": "a95a34f2-e894-4fb6-a2c9-15d12c1e3135", "merge_behavior": "AWAIT_ANY"},
+                            "ports": [],
+                            "adornments": None,
+                            "attributes": [],
+                        },
                         {
                             "id": "a773c3a5-78cb-4250-8d29-7282e8a579d3",
                             "type": "TERMINAL",

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

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -13,12 +10,7 @@ def test_serialize_workflow():
     # GIVEN a Workflow that uses a MapNode
     # WHEN we serialize it
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleMapExample)
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -141,7 +133,23 @@ def test_serialize_workflow():
                                 "bases": [],
                             },
                         },
-                        {"type": "MOCKED"},
+                        {
+                            "id": "baf6d316-dc75-41e8-96c0-015aede96309",
+                            "label": "Iteration",
+                            "type": "GENERIC",
+                            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+                            "definition": {
+                                "name": "Iteration",
+                                "module": ["tests", "workflows", "basic_map_node", "workflow"],
+                                "bases": [
+                                    {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}
+                                ],
+                            },
+                            "trigger": {"id": "01324747-9bc0-4ecd-a8ab-40dca5a94e2e", "merge_behavior": "AWAIT_ANY"},
+                            "ports": [],
+                            "adornments": None,
+                            "attributes": [],
+                        },
                         {
                             "id": "6f4883b2-70b1-4e1c-ae15-7d0f5aec810b",
                             "type": "TERMINAL",

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

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -15,12 +12,7 @@ def test_serialize_workflow__await_all():
     workflow_display = get_workflow_display(
         base_display_class=VellumWorkflowDisplay, workflow_class=AwaitAllPassingWorkflow
     )
-
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -76,22 +68,8 @@ def test_serialize_workflow__await_all():
         },
     }
 
-    passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "MOCKED"]
-    assert not DeepDiff(
-        [
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
-            {
-                "type": "MOCKED",
-            },
-        ],
-        passthrough_nodes,
-        ignore_order=True,
-    )
+    passthrough_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "GENERIC"]
+    assert len(passthrough_nodes) == 3
 
     merge_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "MERGE")
     assert not DeepDiff(

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

@@ -1,8 +1,5 @@
-from unittest import mock
-
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -13,10 +10,7 @@ def test_serialize_workflow():
     # GIVEN a Workflow with a TryNode
     # WHEN we serialize it
     workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=SimpleTryExample)
-
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        serialized_workflow: dict = workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -82,4 +76,4 @@ def test_serialize_workflow():
     }
 
     try_node = workflow_raw_data["nodes"][1]
-    assert try_node == {"type": "MOCKED"}
+    assert try_node["id"] == "1381c078-efa2-4255-89a1-7b4cb742c7fc"

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

@@ -1,9 +1,7 @@
 import pytest
-from unittest import mock
 
 from deepdiff import DeepDiff
 
-from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
 
@@ -13,17 +11,12 @@ from tests.workflows.complex_final_output_node.missing_workflow_output import Mi
 
 def test_serialize_workflow__missing_final_output_node():
     # GIVEN a Workflow that is missing a Terminal Node
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=MissingFinalOutputNodeWorkflow
+    )
 
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        workflow_display = get_workflow_display(
-            base_display_class=VellumWorkflowDisplay, workflow_class=MissingFinalOutputNodeWorkflow
-        )
-
-        # WHEN we serialize it
-        serialized_workflow: dict = workflow_display.serialize()
+    # WHEN we serialize it
+    serialized_workflow: dict = workflow_display.serialize()
 
     # THEN we should get a serialized representation of the Workflow
     assert serialized_workflow.keys() == {
@@ -97,10 +90,8 @@ def test_serialize_workflow__missing_final_output_node():
         },
     }
 
-    passthrough_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "MOCKED")
-    assert passthrough_node == {
-        "type": "MOCKED",
-    }
+    passthrough_node = next(node for node in workflow_raw_data["nodes"] if node["type"] == "GENERIC")
+    assert passthrough_node["id"] == "32d88cab-e9fa-4a56-9bc2-fb6e1fd0897f"
 
     final_output_nodes = [node for node in workflow_raw_data["nodes"] if node["type"] == "TERMINAL"]
     assert not DeepDiff(
@@ -218,17 +209,12 @@ def test_serialize_workflow__missing_final_output_node():
 
 def test_serialize_workflow__missing_workflow_output():
     # GIVEN a Workflow that contains a terminal node that is unreferenced by the Workflow's Outputs
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=MissingWorkflowOutputWorkflow
+    )
 
-    # TODO: Support serialization of BaseNode
-    # https://app.shortcut.com/vellum/story/4871/support-serialization-of-base-node
-    with mock.patch.object(BaseNodeVellumDisplay, "serialize") as mocked_serialize:
-        mocked_serialize.return_value = {"type": "MOCKED"}
-        workflow_display = get_workflow_display(
-            base_display_class=VellumWorkflowDisplay, workflow_class=MissingWorkflowOutputWorkflow
-        )
-
-        # WHEN we serialize it, it should throw an error
-        with pytest.raises(ValueError) as exc_info:
-            workflow_display.serialize()
+    # WHEN we serialize it, it should throw an error
+    with pytest.raises(ValueError) as exc_info:
+        workflow_display.serialize()
 
     assert exc_info.value.args[0] == "Unable to serialize terminal nodes that are not referenced by workflow outputs."

vellum_ee/workflows/display/types.py

@@ -36,8 +36,12 @@ class WorkflowDisplayContext(
     workflow_display_class: Type["BaseWorkflowDisplay"]
     workflow_display: WorkflowMetaDisplayType
     workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = field(default_factory=dict)
+    global_workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = field(
+        default_factory=dict
+    )
     node_displays: Dict[Type[BaseNode], "NodeDisplayType"] = field(default_factory=dict)
-    node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], "NodeOutputDisplay"]] = field(
+    global_node_displays: Dict[Type[BaseNode], NodeDisplayType] = field(default_factory=dict)
+    global_node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], "NodeOutputDisplay"]] = field(
         default_factory=dict
     )
     entrypoint_displays: Dict[Type[BaseNode], EntrypointDisplayType] = field(default_factory=dict)

vellum_ee/workflows/display/utils/vellum.py

@@ -58,13 +58,13 @@ def create_node_input_value_pointer_rule(
     value: Any, display_context: WorkflowDisplayContext
 ) -> NodeInputValuePointerRule:
     if isinstance(value, OutputReference):
-        upstream_node, output_display = display_context.node_output_displays[value]
-        upstream_node_display = display_context.node_displays[upstream_node]
+        upstream_node, output_display = display_context.global_node_output_displays[value]
+        upstream_node_display = display_context.global_node_displays[upstream_node]
         return NodeOutputPointer(
             data=NodeOutputData(node_id=str(upstream_node_display.node_id), output_id=str(output_display.id)),
         )
     if isinstance(value, WorkflowInputReference):
-        workflow_input_display = display_context.workflow_input_displays[value]
+        workflow_input_display = display_context.global_workflow_input_displays[value]
         return InputVariablePointer(data=InputVariableData(input_variable_id=str(workflow_input_display.id)))
     if isinstance(value, VellumSecretReference):
         # TODO: Pass through the name instead of retrieving the ID

vellum_ee/workflows/display/vellum.py

@@ -38,6 +38,10 @@ class NodeDisplayData(UniversalBaseModel):
     comment: Optional[NodeDisplayComment] = None
 
 
+class GenericNodeDisplayData(UniversalBaseModel):
+    position: NodeDisplayPosition = Field(default_factory=NodeDisplayPosition)
+
+
 class CodeResourceDefinition(UniversalBaseModel):
     name: str
     module: List[str]

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -1,9 +1,9 @@
 from abc import abstractmethod
-from copy import deepcopy
+from copy import copy
 from functools import cached_property
 import logging
 from uuid import UUID
-from typing import Any, Dict, Generic, Optional, Tuple, Type, get_args
+from typing import Any, Dict, Generic, Iterator, List, Optional, Tuple, Type, get_args
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
@@ -71,6 +71,10 @@ class BaseWorkflowDisplay(
     # Used to store the mapping between workflows and their display classes
     _workflow_display_registry: Dict[Type[WorkflowType], Type["BaseWorkflowDisplay"]] = {}
 
+    _errors: List[Exception]
+
+    _dry_run: bool
+
     def __init__(
         self,
         workflow: Type[WorkflowType],
@@ -85,12 +89,15 @@ class BaseWorkflowDisplay(
                 EdgeDisplayType,
             ]
         ] = None,
+        dry_run: bool = False,
     ):
         self._workflow = workflow
         self._parent_display_context = parent_display_context
+        self._errors: List[Exception] = []
+        self._dry_run = dry_run
 
     @abstractmethod
-    def serialize(self, raise_errors: bool = True) -> JsonObject:
+    def serialize(self) -> JsonObject:
         pass
 
     @classmethod
@@ -110,7 +117,18 @@ class BaseWorkflowDisplay(
     def node_display_base_class(self) -> Type[NodeDisplayType]:
         pass
 
-    def _enrich_node_output_displays(
+    def add_error(self, error: Exception) -> None:
+        if self._dry_run:
+            self._errors.append(error)
+            return
+
+        raise error
+
+    @property
+    def errors(self) -> Iterator[Exception]:
+        return iter(self._errors)
+
+    def _enrich_global_node_output_displays(
         self,
         node: Type[BaseNode],
         node_display: NodeDisplayType,
@@ -126,7 +144,7 @@ class BaseWorkflowDisplay(
                 inner_node = get_wrapped_node(node)
                 if inner_node._is_wrapped_node:
                     inner_node_display = self._get_node_display(inner_node)
-                    self._enrich_node_output_displays(inner_node, inner_node_display, node_output_displays)
+                    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
@@ -174,19 +192,26 @@ class BaseWorkflowDisplay(
     ]:
         workflow_display = self._generate_workflow_meta_display()
 
-        # If we're dealing with a nested workflow, then it should have access to the outputs of all nodes
-        node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]] = (
-            deepcopy(self._parent_display_context.node_output_displays) if self._parent_display_context else {}
+        global_node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]] = (
+            copy(self._parent_display_context.global_node_output_displays) if self._parent_display_context else {}
         )
 
         node_displays: Dict[Type[BaseNode], NodeDisplayType] = {}
+
+        global_node_displays: Dict[Type[BaseNode], NodeDisplayType] = (
+            copy(self._parent_display_context.global_node_displays) if self._parent_display_context else {}
+        )
+
         port_displays: Dict[Port, PortDisplay] = {}
 
         # TODO: We should still serialize nodes that are in the workflow's directory but aren't used in the graph.
         # https://app.shortcut.com/vellum/story/5394
         for node in self._workflow.get_nodes():
+            if node in global_node_displays:
+                continue
             node_display = self._get_node_display(node)
             node_displays[node] = node_display
+            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):
@@ -195,22 +220,23 @@ class BaseWorkflowDisplay(
 
                 if inner_node._is_wrapped_node:
                     node_displays[inner_node] = inner_node_display
+                    global_node_displays[inner_node] = inner_node_display
 
-            self._enrich_node_output_displays(node, node_display, node_output_displays)
+            self._enrich_global_node_output_displays(node, node_display, global_node_output_displays)
             self._enrich_node_port_displays(node, node_display, port_displays)
 
+        workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {}
         # If we're dealing with a nested workflow, then it should have access to the inputs of its parents.
-        workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = (
-            deepcopy(self._parent_display_context.workflow_input_displays) if self._parent_display_context else {}
+        global_workflow_input_displays = (
+            copy(self._parent_display_context.workflow_input_displays) if self._parent_display_context else {}
         )
         for workflow_input in self._workflow.get_inputs_class():
-            if workflow_input in workflow_input_displays:
-                continue
-
             workflow_input_display_overrides = self.inputs_display.get(workflow_input)
-            workflow_input_displays[workflow_input] = self._generate_workflow_input_display(
+            input_display = self._generate_workflow_input_display(
                 workflow_input, overrides=workflow_input_display_overrides
             )
+            workflow_input_displays[workflow_input] = input_display
+            global_workflow_input_displays[workflow_input] = input_display
 
         entrypoint_displays: Dict[Type[BaseNode], EntrypointDisplayType] = {}
         for entrypoint in self._workflow.get_entrypoints():
@@ -253,8 +279,10 @@ class BaseWorkflowDisplay(
         return WorkflowDisplayContext(
             workflow_display=workflow_display,
             workflow_input_displays=workflow_input_displays,
+            global_workflow_input_displays=global_workflow_input_displays,
             node_displays=node_displays,
-            node_output_displays=node_output_displays,
+            global_node_output_displays=global_node_output_displays,
+            global_node_displays=global_node_displays,
             entrypoint_displays=entrypoint_displays,
             workflow_output_displays=workflow_output_displays,
             edge_displays=edge_displays,

vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py

@@ -10,6 +10,7 @@ def get_workflow_display(
     workflow_class: Type[WorkflowType],
     root_workflow_class: Optional[Type[WorkflowType]] = None,
     parent_display_context: Optional[WorkflowDisplayContext] = None,
+    dry_run: bool = False,
 ) -> WorkflowDisplayType:
     try:
         workflow_display_class = base_display_class.get_from_workflow_display_registry(workflow_class)
@@ -20,6 +21,7 @@ def get_workflow_display(
                 workflow_class=workflow_class.__bases__[0],
                 root_workflow_class=workflow_class if root_workflow_class is None else root_workflow_class,
                 parent_display_context=parent_display_context,
+                dry_run=dry_run,
             )
         except IndexError:
             return base_display_class(workflow_class)
@@ -27,4 +29,5 @@ def get_workflow_display(
     return workflow_display_class(  # type: ignore[return-value]
         workflow_class,
         parent_display_context=parent_display_context,
+        dry_run=dry_run,
     )

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -53,7 +53,7 @@ class VellumWorkflowDisplay(
 ):
     node_display_base_class = BaseNodeVellumDisplay
 
-    def serialize(self, raise_errors: bool = True) -> JsonObject:
+    def serialize(self) -> JsonObject:
         input_variables: JsonArray = []
         for workflow_input, workflow_input_display in self.display_context.workflow_input_displays.items():
             default = primitive_to_vellum_value(workflow_input.instance) if workflow_input.instance else None
@@ -106,12 +106,9 @@ class VellumWorkflowDisplay(
 
             try:
                 serialized_node = node_display.serialize(self.display_context)
-            except NotImplementedError:
-                logger.warning("Unable to serialize node", extra={"node": node.__name__})
-                if raise_errors:
-                    raise
-                else:
-                    continue
+            except NotImplementedError as e:
+                self.add_error(e)
+                continue
 
             nodes.append(serialized_node)
 
@@ -205,7 +202,9 @@ class VellumWorkflowDisplay(
         # If there are terminal nodes with no workflow output reference,
         # raise a serialization error
         if len(unreferenced_final_output_node_outputs) > 0:
-            raise ValueError("Unable to serialize terminal nodes that are not referenced by workflow outputs.")
+            self.add_error(
+                ValueError("Unable to serialize terminal nodes that are not referenced by workflow outputs.")
+            )
 
         # Add an edge for each edge in the workflow
         all_edge_displays: List[EdgeVellumDisplay] = [