vellum-ai 0.13.18__py3-none-any.whl → 0.13.19__py3-none-any.whl

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.18",
+            "X-Fern-SDK-Version": "0.13.19",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/workflows/nodes/core/retry_node/node.py

@@ -1,3 +1,4 @@
+import time
 from typing import Callable, Generic, Optional, Type
 
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -17,11 +18,13 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
     Used to retry a Subworkflow a specified number of times.
 
     max_attempts: int - The maximum number of attempts to retry the Subworkflow
+    delay: float - The number of seconds to wait between retries
     retry_on_error_code: Optional[VellumErrorCode] = None - The error code to retry on
     subworkflow: Type["BaseWorkflow[SubworkflowInputs, BaseState]"] - The Subworkflow to execute
     """
 
     max_attempts: int
+    delay: Optional[float] = None
     retry_on_error_code: Optional[WorkflowErrorCode] = None
     retry_on_condition: Optional[BaseDescriptor] = None
 
@@ -29,7 +32,9 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
         attempt_number: int
 
     def run(self) -> BaseNode.Outputs:
-        last_exception = Exception("max_attempts must be greater than 0")
+        if self.max_attempts <= 0:
+            raise Exception("max_attempts must be greater than 0")
+
         for index in range(self.max_attempts):
             attempt_number = index + 1
             context = WorkflowContext(vellum_client=self._context.vellum_client)
@@ -49,30 +54,29 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
 
                 return node_outputs
             elif terminal_event.name == "workflow.execution.paused":
-                last_exception = NodeException(
+                raise NodeException(
                     code=WorkflowErrorCode.INVALID_OUTPUTS,
                     message=f"Subworkflow unexpectedly paused on attempt {attempt_number}",
                 )
-                break
             elif self.retry_on_error_code and self.retry_on_error_code != terminal_event.error.code:
-                last_exception = NodeException(
+                raise NodeException(
                     code=WorkflowErrorCode.INVALID_OUTPUTS,
                     message=f"""Unexpected rejection on attempt {attempt_number}: {terminal_event.error.code.value}.
 Message: {terminal_event.error.message}""",
                 )
-                break
             elif self.retry_on_condition and not resolve_value(self.retry_on_condition, self.state):
-                last_exception = NodeException(
+                raise NodeException(
                     code=WorkflowErrorCode.INVALID_OUTPUTS,
                     message=f"""Rejection failed on attempt {attempt_number}: {terminal_event.error.code.value}.
 Message: {terminal_event.error.message}""",
                 )
-                break
             else:
                 last_exception = NodeException(
                     terminal_event.error.message,
                     code=terminal_event.error.code,
                 )
+                if self.delay:
+                    time.sleep(self.delay)
 
         raise last_exception
 
@@ -80,6 +84,7 @@ Message: {terminal_event.error.message}""",
     def wrap(
         cls,
         max_attempts: int,
+        delay: Optional[float] = None,
         retry_on_error_code: Optional[WorkflowErrorCode] = None,
         retry_on_condition: Optional[BaseDescriptor] = None,
     ) -> Callable[..., Type["RetryNode"]]:
@@ -87,6 +92,7 @@ Message: {terminal_event.error.message}""",
             cls,
             attributes={
                 "max_attempts": max_attempts,
+                "delay": delay,
                 "retry_on_error_code": retry_on_error_code,
                 "retry_on_condition": retry_on_condition,
             },

vellum/workflows/outputs/base.py

@@ -6,6 +6,8 @@ from pydantic_core import core_schema
 
 from vellum.workflows.constants import UNDEF
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.utils import get_class_attr_names, infer_types
 
@@ -183,6 +185,15 @@ class _BaseOutputsMeta(type):
 
 class BaseOutputs(metaclass=_BaseOutputsMeta):
     def __init__(self, **kwargs: Any) -> None:
+        declared_fields = {descriptor.name for descriptor in self.__class__}
+        provided_fields = set(kwargs.keys())
+
+        if not provided_fields.issubset(declared_fields):
+            raise NodeException(
+                message=f"Unexpected outputs: {provided_fields - declared_fields}",
+                code=WorkflowErrorCode.INVALID_OUTPUTS,
+            )
+
         for name, value in kwargs.items():
             setattr(self, name, value)
 

{vellum_ai-0.13.18.dist-info → vellum_ai-0.13.19.dist-info}/METADATA

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

{vellum_ai-0.13.18.dist-info → vellum_ai-0.13.19.dist-info}/RECORD

@@ -2,7 +2,7 @@ vellum_cli/CONTRIBUTING.md,sha256=FtDC7BGxSeMnwCXAUssFsAIElXtmJE-O5Z7BpolcgvI,29
 vellum_cli/README.md,sha256=2NudRoLzWxNKqnuVy1JuQ7DerIaxWGYkrH8kMd-asIE,90
 vellum_cli/__init__.py,sha256=uEn2Nlt2Z0kBc79NcO4-rgOIE7H9nsMEEjWF6MLDlPo,10591
 vellum_cli/aliased_group.py,sha256=ugW498j0yv4ALJ8vS9MsO7ctDW7Jlir9j6nE_uHAP8c,3363
-vellum_cli/config.py,sha256=LVRB-SEJcpQYfg2QGcjKHmRSAijdSFADbS90gDY4AI8,6829
+vellum_cli/config.py,sha256=Bsb3mnvKvv3oOTcCuxpgC7lWPMqt6eJhgRA6VEE-vL4,9266
 vellum_cli/image_push.py,sha256=SJwhwWJsLjwGNezNVd_oCVpFMfPsAB3dfLWmriZZUtw,4419
 vellum_cli/logger.py,sha256=PuRFa0WCh4sAGFS5aqWB0QIYpS6nBWwPJrIXpWxugV4,1022
 vellum_cli/ping.py,sha256=lWyJw6sziXjyTopTYRdFF5hV-sYPVDdX0yVbG5fzcY4,585
@@ -13,7 +13,7 @@ vellum_cli/tests/conftest.py,sha256=AFYZryKA2qnUuCPBxBKmHLFoPiE0WhBFFej9tNwSHdc,
 vellum_cli/tests/test_config.py,sha256=uvKGDc8BoVyT9_H0Z-g8469zVxomn6Oi3Zj-vK7O_wU,2631
 vellum_cli/tests/test_main.py,sha256=qDZG-aQauPwBwM6A2DIu1494n47v3pL28XakTbLGZ-k,272
 vellum_cli/tests/test_ping.py,sha256=QtbhYKMYn1DFnDyBij2mkQO32j9KOpZ5Pf0yek7k_Ao,1284
-vellum_cli/tests/test_pull.py,sha256=CEnVpb9i-_I0wOSabjgJxleRuvvAHFfDB0_f4CvpwMI,23742
+vellum_cli/tests/test_pull.py,sha256=JURmgGs5lSnpzefSx4K13eF2swv7O8OF86-4df81Zjo,25241
 vellum_cli/tests/test_push.py,sha256=kcM8jp6vTRBwfk6cAC_tHa4Eliuy_zZ-X1DfIpSEIW0,18299
 vellum_ee/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -28,7 +28,7 @@ vellum_ee/workflows/display/nodes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5J
 vellum_ee/workflows/display/nodes/tests/test_base_node_display.py,sha256=QqR3Ly0RNrXwOeLdW5nERDFt0gRPf76n1bPES6o5UN4,1093
 vellum_ee/workflows/display/nodes/types.py,sha256=St1BB6no528OyELGiyRabWao0GGw6mLhstQAvEACbGk,247
 vellum_ee/workflows/display/nodes/utils.py,sha256=sloya5TpXsnot1HURc9L51INwflRqUzHxRVnCS9Cd-4,973
-vellum_ee/workflows/display/nodes/vellum/__init__.py,sha256=Moeuvm0u7c_6wtxUWOJqTJYbrnG4Rp_5mKRoYpfewvQ,1578
+vellum_ee/workflows/display/nodes/vellum/__init__.py,sha256=VHx6wSs9wuuiFlZpSQSd3mhECz4SUy2wEBuTSv--_As,1578
 vellum_ee/workflows/display/nodes/vellum/api_node.py,sha256=hoV-cUtS6H9kmRQXHd2py95GRWI_dAnnaPwvlNBkDOQ,8571
 vellum_ee/workflows/display/nodes/vellum/code_execution_node.py,sha256=z00Z3L0d4PsUQo4S8FRDTtOFLtjdi17TJbatNVF4nM8,4288
 vellum_ee/workflows/display/nodes/vellum/conditional_node.py,sha256=ybLIa4uclqVIy3VAQvI1ivg2tnK5Ug_1R5a69DFqL7E,11104
@@ -56,7 +56,7 @@ vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=h4bE187
 vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py,sha256=EenmEdBtHUFQ0OS-kE7Vboax3JnDdj-K4Qixt5oR0Po,2253
-vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=6wgLNd-Ae-33SPIYhOvQ9KJ8-5ach98v_Z3g7nrN_xs,8026
+vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py,sha256=KYdohS5pRgHM0DcUaK0tHa40f0gSvDKi2K5On0zNEU8,8305
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py,sha256=yPXhdZxEDunNl5LL5Obb0jeO34djt7F-GiaTJhp_fmo,16742
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py,sha256=-12ZkZb3f5gyoNASV2yeQtMo5HmNsVEo8nXwL6IC-I8,6261
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=6th6kCwzql6lddjkTQx4Jbvvs4ChqtJwctW-B4QuBhI,37352
@@ -113,7 +113,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=8nZt88C9SVwWanjLbIQMU3rzb32h5UZfFMBx3VPHB50,111887
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=JapYtCCqy-ryd-fk1i1krQS6K630bgNpUw4zE0Lct8s,1869
+vellum/client/core/client_wrapper.py,sha256=rnVChhyob3hO3iXrfL4gsCAYGAOdglXQdXoLL07FKLg,1869
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -1323,7 +1323,7 @@ vellum/workflows/nodes/core/map_node/node.py,sha256=bjCVMAzkqJUvaLWVBObjskcutwLG
 vellum/workflows/nodes/core/map_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/map_node/tests/test_node.py,sha256=P7ImQyeLcK-aEJUYCX1gPQyuOxdxnSPpI3mNPpKQ62Y,1919
 vellum/workflows/nodes/core/retry_node/__init__.py,sha256=lN2bIy5a3Uzhs_FYCrooADyYU6ZGShtvLKFWpelwPvo,60
-vellum/workflows/nodes/core/retry_node/node.py,sha256=QEpxhKOyxDkRoAn2b0PToZWtAGQetSQYVTpb9yCOLlw,4028
+vellum/workflows/nodes/core/retry_node/node.py,sha256=WD96o-eOj3dwEEe2nqxwBbmLTIyPRYB3Lk4T6XHRX74,4214
 vellum/workflows/nodes/core/retry_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/core/retry_node/tests/test_node.py,sha256=RM_OHwxrHwyxvlQQBJPqVBxpedFuWQ9h2-Xa3kP75sc,4399
 vellum/workflows/nodes/core/templating_node/__init__.py,sha256=GmyuYo81_A1_Bz6id69ozVFS6FKiuDsZTiA3I6MaL2U,70
@@ -1389,7 +1389,7 @@ vellum/workflows/nodes/experimental/openai_chat_completion_node/__init__.py,sha2
 vellum/workflows/nodes/experimental/openai_chat_completion_node/node.py,sha256=1EGeiaT-Zoo6pttQFKKBcdf3dmhAbjKGaErYD5FFwlc,10185
 vellum/workflows/nodes/utils.py,sha256=T7krLipDSI007JkkH2zsfyQ-tfOBH4hyFzQ0YJKYpuw,4025
 vellum/workflows/outputs/__init__.py,sha256=AyZ4pRh_ACQIGvkf0byJO46EDnSix1ZCAXfvh-ms1QE,94
-vellum/workflows/outputs/base.py,sha256=a7W6rNSDSawwGAXYjNTF2iHb9lnZu7WFSOagZIyy__k,7976
+vellum/workflows/outputs/base.py,sha256=Wu48tqyQoxpflBUcwzKeZjSVp1LPKrBwuIvnx__9H90,8459
 vellum/workflows/ports/__init__.py,sha256=bZuMt-R7z5bKwpu4uPW7LlJeePOQWmCcDSXe5frUY5g,101
 vellum/workflows/ports/node_ports.py,sha256=g4A-8iUAvEJSkaWppbvzAR8XU02R9U-qLN4rP2Kq4Aw,2743
 vellum/workflows/ports/port.py,sha256=rc3GB7dDQCUs0IbY08a92-31YzJHQgBeww13brSJ2Js,3172
@@ -1442,8 +1442,8 @@ vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
 vellum/workflows/workflows/base.py,sha256=k0kUWWko4fHyCqLSU_1cBK_pXZpl9MXekWiG-bdOAo0,18353
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
-vellum_ai-0.13.18.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.13.18.dist-info/METADATA,sha256=aLUy-xu3bfn69UAWushKWy3YWWoMMruLabea-UEOO1o,5335
-vellum_ai-0.13.18.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.13.18.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.13.18.dist-info/RECORD,,
+vellum_ai-0.13.19.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.13.19.dist-info/METADATA,sha256=NzphP-OKn-7lL7hjSD1e1eT709mTax682Cs85AT0O_4,5335
+vellum_ai-0.13.19.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.13.19.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.13.19.dist-info/RECORD,,

vellum_cli/config.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from dataclasses import field
 import json
 import os
@@ -88,6 +89,50 @@ class WorkflowConfig(UniversalBaseModel):
         )
 
 
+def merge_workflows_by_sandbox_id(
+    workflows: List[WorkflowConfig], other_workflows: List[WorkflowConfig]
+) -> List[WorkflowConfig]:
+    merged_workflows: List[WorkflowConfig] = []
+    for self_workflow in workflows:
+        if self_workflow.workflow_sandbox_id is None:
+            # If the user defines a workflow in the pyproject.toml without a sandbox_id,
+            # we merge the workflow with one of the ones in the lockfile.
+            other_workflow = next(
+                (
+                    other_workflow
+                    for other_workflow in other_workflows
+                    if self_workflow.workspace == other_workflow.workspace
+                ),
+                None,
+            )
+            if other_workflow is not None:
+                merged_workflows.append(self_workflow.merge(other_workflow))
+            else:
+                merged_workflows.append(self_workflow)
+        else:
+            # If the user defines a workflow in the pyproject.toml with a sandbox_id,
+            # we merge the workflow with one of the ones in the lockfile with the same sandbox_id.
+            other_workflow = next(
+                (
+                    other_workflow
+                    for other_workflow in other_workflows
+                    if self_workflow.workflow_sandbox_id == other_workflow.workflow_sandbox_id
+                ),
+                None,
+            )
+            if other_workflow is not None:
+                merged_workflows.append(self_workflow.merge(other_workflow))
+            else:
+                merged_workflows.append(self_workflow)
+
+    workflow_sandbox_ids_so_far = {workflow.workflow_sandbox_id for workflow in merged_workflows}
+    for other_workflow in other_workflows:
+        if other_workflow.workflow_sandbox_id not in workflow_sandbox_ids_so_far:
+            merged_workflows.append(other_workflow)
+
+    return merged_workflows
+
+
 class VellumCliConfig(UniversalBaseModel):
     version: Literal["1.0"] = "1.0"
     workflows: List[WorkflowConfig] = field(default_factory=list)
@@ -97,24 +142,31 @@ class VellumCliConfig(UniversalBaseModel):
         lockfile_path = os.path.join(os.getcwd(), LOCKFILE_PATH)
         with open(lockfile_path, "w") as f:
             json.dump(self.model_dump(), f, indent=2, cls=DefaultStateEncoder)
+            # Makes sure the file ends with a newline, consistent with most autoformatters
+            f.write("\n")
 
     def merge(self, other: "VellumCliConfig") -> "VellumCliConfig":
         if other.version != self.version:
             raise ValueError("Lockfile version mismatch")
 
-        self_workflow_by_module = {workflow.module: workflow for workflow in self.workflows}
-        other_workflow_by_module = {workflow.module: workflow for workflow in other.workflows}
-        all_modules = sorted(set(self_workflow_by_module.keys()).union(set(other_workflow_by_module.keys())))
+        self_workflows_by_module = self.get_workflows_by_module_mapping()
+        other_workflows_by_module = other.get_workflows_by_module_mapping()
+        all_modules = sorted(set(self_workflows_by_module.keys()).union(set(other_workflows_by_module.keys())))
         merged_workflows = []
         for module in all_modules:
-            self_workflow = self_workflow_by_module.get(module)
-            other_workflow = other_workflow_by_module.get(module)
-            if self_workflow and other_workflow:
-                merged_workflows.append(self_workflow.merge(other_workflow))
-            elif self_workflow:
-                merged_workflows.append(self_workflow)
-            elif other_workflow:
-                merged_workflows.append(other_workflow)
+            self_workflows = self_workflows_by_module.get(module)
+            other_workflows = other_workflows_by_module.get(module)
+            if self_workflows and other_workflows:
+                merged_workflows.extend(
+                    merge_workflows_by_sandbox_id(
+                        self_workflows,
+                        other_workflows,
+                    )
+                )
+            elif self_workflows:
+                merged_workflows.extend(self_workflows)
+            elif other_workflows:
+                merged_workflows.extend(other_workflows)
 
         self_workspace_by_name = {workspace.name: workspace for workspace in self.workspaces}
         other_workspace_by_name = {workspace.name: workspace for workspace in other.workspaces}
@@ -132,6 +184,12 @@ class VellumCliConfig(UniversalBaseModel):
 
         return VellumCliConfig(workflows=merged_workflows, workspaces=merged_workspaces, version=self.version)
 
+    def get_workflows_by_module_mapping(self) -> Dict[str, List[WorkflowConfig]]:
+        workflows_by_module = defaultdict(list)
+        for workflow in self.workflows:
+            workflows_by_module[workflow.module].append(workflow)
+        return workflows_by_module
+
 
 def load_vellum_cli_config(root_dir: Optional[str] = None) -> VellumCliConfig:
     if root_dir is None:

vellum_cli/tests/test_pull.py

@@ -428,7 +428,7 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "workflows": [
                     {
                         "module": module,
-                        "workflow_sandbox_id": "0edc07cd-45b9-43e8-99bc-1f181972a857",
+                        "workflow_sandbox_id": workflow_sandbox_id,
                         "ignore": "tests/*",
                         "deployments": [
                             {
@@ -646,3 +646,44 @@ def test_pull__module_not_in_config(vellum_client, mock_module):
                 "workspace": "default",
             }
         ]
+
+
+def test_pull__multiple_instances_of_same_module__keep_when_pulling_another_module(vellum_client, mock_module):
+    # GIVEN a module on the user's filesystem
+    module = mock_module.module
+    temp_dir = mock_module.temp_dir
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+
+    # AND the vellum lock file has two instances of some other module
+    lock_data = {
+        "workflows": [
+            {
+                "module": "some_other_module",
+                "workflow_sandbox_id": str(uuid4()),
+                "workspace": "default",
+            },
+            {
+                "module": "some_other_module",
+                "workflow_sandbox_id": str(uuid4()),
+                "workspace": "other",
+            },
+        ]
+    }
+    lock_json = os.path.join(temp_dir, "vellum.lock.json")
+    with open(lock_json, "w") as f:
+        json.dump(lock_data, f)
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+
+    # WHEN the user runs the pull command on the new module
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", module, "--workflow-sandbox-id", workflow_sandbox_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0, (result.output, result.exception)
+
+    # AND the lockfile should have all three entries
+    with open(lock_json) as f:
+        lock_data = json.load(f)
+        assert len(lock_data["workflows"]) == 3

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

@@ -18,7 +18,6 @@ from .try_node import BaseTryNodeDisplay
 
 # All node display classes must be imported here to be registered in BaseNodeDisplay's node display registry
 __all__ = [
-    "BaseRetryNodeDisplay",
     "BaseAPINodeDisplay",
     "BaseCodeExecutionNodeDisplay",
     "BaseConditionalNodeDisplay",
@@ -31,6 +30,7 @@ __all__ = [
     "BaseMergeNodeDisplay",
     "BaseNoteNodeDisplay",
     "BasePromptDeploymentNodeDisplay",
+    "BaseRetryNodeDisplay",
     "BaseSearchNodeDisplay",
     "BaseSubworkflowDeploymentNodeDisplay",
     "BaseTemplatingNodeDisplay",

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

@@ -80,6 +80,11 @@ def test_serialize_node__retry(serialize_node):
                         "module": ["vellum", "workflows", "nodes", "core", "retry_node", "node"],
                     },
                     "attributes": [
+                        {
+                            "id": "8a07dc58-3fed-41d4-8ca6-31ee0bb86c61",
+                            "name": "delay",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
                         {
                             "id": "f388e93b-8c68-4f54-8577-bbd0c9091557",
                             "name": "max_attempts",