vellum-ai 0.13.14__py3-none-any.whl → 0.13.15__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.13",
+            "X-Fern-SDK-Version": "0.13.15",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

{vellum_ai-0.13.14.dist-info → vellum_ai-0.13.15.dist-info}/METADATA

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

{vellum_ai-0.13.14.dist-info → vellum_ai-0.13.15.dist-info}/RECORD

@@ -1,20 +1,20 @@
 vellum_cli/CONTRIBUTING.md,sha256=FtDC7BGxSeMnwCXAUssFsAIElXtmJE-O5Z7BpolcgvI,2935
 vellum_cli/README.md,sha256=2NudRoLzWxNKqnuVy1JuQ7DerIaxWGYkrH8kMd-asIE,90
-vellum_cli/__init__.py,sha256=tRcOi6oLzT3EXORlHwGzo0Sn1-npf8Jf7mMi90NDUh0,10035
+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/image_push.py,sha256=SJwhwWJsLjwGNezNVd_oCVpFMfPsAB3dfLWmriZZUtw,4419
 vellum_cli/logger.py,sha256=PuRFa0WCh4sAGFS5aqWB0QIYpS6nBWwPJrIXpWxugV4,1022
 vellum_cli/ping.py,sha256=lWyJw6sziXjyTopTYRdFF5hV-sYPVDdX0yVbG5fzcY4,585
 vellum_cli/pull.py,sha256=zf0y22XptUYI_hMP_4Q1CEo9s2wALsTJcCXNd-_ibd8,7551
-vellum_cli/push.py,sha256=0Q3Of25KzqDq0KTgkXQ22nAhSH4_JI424hKwKkj3RWs,8627
+vellum_cli/push.py,sha256=VuAQ26stoweihwiY8ms8RpU6i200fO5SjU7MiMIDvAo,9217
 vellum_cli/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_cli/tests/conftest.py,sha256=AFYZryKA2qnUuCPBxBKmHLFoPiE0WhBFFej9tNwSHdc,1526
 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=KRvQ_fwsFH6tBU49O-KOOAzcHxdB-DPbUkdIdahjuTs,19841
-vellum_cli/tests/test_push.py,sha256=3to7bOOVoWB888ee7TIAVl2VgXCsJhP_aPcrXwF9c_c,14023
+vellum_cli/tests/test_push.py,sha256=AgJVaDeIDa--iVvDrUMxjuvOoC-ylwD_A5KumUJTtJU,18810
 vellum_ee/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -112,7 +112,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=WBZxe9jwGagoDy-B3qFt1xxFJYUu4kED2-27248SBsQ,1869
+vellum/client/core/client_wrapper.py,sha256=_D3P4oNIluyIpVSoNUIGx4jK18hJGMdrTL62pI-ApYg,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
@@ -1441,8 +1441,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.14.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.13.14.dist-info/METADATA,sha256=-XRE0gS6VYSqoJt7wrYK1LcXz99LGaGwgY-9L7THqsg,5335
-vellum_ai-0.13.14.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.13.14.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.13.14.dist-info/RECORD,,
+vellum_ai-0.13.15.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.13.15.dist-info/METADATA,sha256=H6f0l4ZqdCd9ANG3dLIdY-vX9NNebMI9n_hBLBQ49AI,5335
+vellum_ai-0.13.15.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.13.15.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.13.15.dist-info/RECORD,,

vellum_cli/__init__.py

@@ -56,6 +56,12 @@ def workflows():
 
 @workflows.command(name="push")
 @click.argument("module", required=False)
+@click.option(
+    "--workflow-sandbox-id",
+    type=str,
+    help="""The specific Workflow Sandbox ID to use when pushing. Must either be already associated \
+with the provided module or be available for use. The Workflow Sandbox must also exist in Vellum.""",
+)
 @click.option("--deploy", is_flag=True, help="Deploy the Workflow after pushing it to Vellum")
 @click.option("--deployment-label", type=str, help="Label to use for the Deployment")
 @click.option("--deployment-name", type=str, help="Unique name for the Deployment")
@@ -74,6 +80,7 @@ def workflows():
 @click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
 def workflows_push(
     module: Optional[str],
+    workflow_sandbox_id: Optional[str],
     deploy: Optional[bool],
     deployment_label: Optional[str],
     deployment_name: Optional[str],
@@ -90,6 +97,7 @@ def workflows_push(
 
     push_command(
         module=module,
+        workflow_sandbox_id=workflow_sandbox_id,
         deploy=deploy,
         deployment_label=deployment_label,
         deployment_name=deployment_name,
@@ -103,6 +111,7 @@ def workflows_push(
 
 @push.command(name="*", hidden=True)
 @click.pass_context
+@click.option("--workflow-sandbox-id", type=str, help="The specific Workflow Sandbox ID to use when pushing")
 @click.option("--deploy", is_flag=True, help="Deploy the Resource after pushing it to Vellum")
 @click.option("--deployment-label", type=str, help="Label to use for the Deployment")
 @click.option("--deployment-name", type=str, help="Unique name for the Deployment")
@@ -121,6 +130,7 @@ def workflows_push(
 @click.option("--workspace", type=str, help="The specific Workspace config to use when pushing")
 def push_module(
     ctx: click.Context,
+    workflow_sandbox_id: Optional[str],
     deploy: Optional[bool],
     deployment_label: Optional[str],
     deployment_name: Optional[str],
@@ -135,6 +145,7 @@ def push_module(
     if ctx.parent:
         push_command(
             module=ctx.parent.invoked_subcommand,
+            workflow_sandbox_id=workflow_sandbox_id,
             deploy=deploy,
             deployment_label=deployment_label,
             deployment_name=deployment_name,

vellum_cli/push.py

@@ -23,6 +23,7 @@ from vellum_ee.workflows.display.workflows.vellum_workflow_display import Vellum
 
 def push_command(
     module: Optional[str] = None,
+    workflow_sandbox_id: Optional[str] = None,
     deploy: Optional[bool] = None,
     deployment_label: Optional[str] = None,
     deployment_name: Optional[str] = None,
@@ -39,13 +40,24 @@ def push_command(
     if not config.workflows:
         raise ValueError("No Workflows found in project to push.")
 
-    if len(config.workflows) > 1 and not module:
-        raise ValueError("Multiple workflows found in project to push. Pushing only a single workflow is supported.")
+    workflow_configs = (
+        [
+            w
+            for w in config.workflows
+            if (module and w.module == module) or (workflow_sandbox_id and w.workflow_sandbox_id == workflow_sandbox_id)
+        ]
+        if module or workflow_sandbox_id
+        else config.workflows
+    )
 
-    workflow_config = next((w for w in config.workflows if w.module == module), None) if module else config.workflows[0]
-    if workflow_config is None:
+    if len(workflow_configs) == 0:
         raise ValueError(f"No workflow config for '{module}' found in project to push.")
 
+    if len(workflow_configs) > 1:
+        raise ValueError("Multiple workflows found in project to push. Pushing only a single workflow is supported.")
+
+    workflow_config = workflow_configs[0]
+
     logger.info(f"Loading workflow from {workflow_config.module}")
     resolved_workspace = workspace or workflow_config.workspace or DEFAULT_WORKSPACE_CONFIG.name
     workspace_config = (
@@ -65,6 +77,15 @@ def push_command(
         )
         config.workflows.append(workflow_config)
 
+    if (
+        workflow_sandbox_id
+        and workflow_config.workflow_sandbox_id
+        and workflow_config.workflow_sandbox_id != workflow_sandbox_id
+    ):
+        raise ValueError(
+            f"Workflow sandbox id '{workflow_sandbox_id}' is already associated with '{workflow_config.module}'."
+        )
+
     client = create_vellum_client(
         api_key=api_key,
     )
@@ -138,7 +159,7 @@ def push_command(
             # https://app.shortcut.com/vellum/story/5585
             exec_config=json.dumps(exec_config),
             label=label,
-            workflow_sandbox_id=workflow_config.workflow_sandbox_id,
+            workflow_sandbox_id=workflow_config.workflow_sandbox_id or workflow_sandbox_id,
             artifact=artifact,
             # We should check with fern if we could auto-serialize typed object fields for us
             # https://app.shortcut.com/vellum/story/5568

vellum_cli/tests/test_push.py

@@ -29,6 +29,21 @@ def _extract_tar_gz(tar_gz_bytes: bytes) -> dict[str, str]:
     return files
 
 
+def _ensure_workflow_py(temp_dir: str, module: str) -> str:
+    base_dir = os.path.join(temp_dir, *module.split("."))
+    os.makedirs(base_dir, exist_ok=True)
+    workflow_py_file_content = """\
+from vellum.workflows import BaseWorkflow
+
+class ExampleWorkflow(BaseWorkflow):
+    pass
+"""
+    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
+        f.write(workflow_py_file_content)
+
+    return workflow_py_file_content
+
+
 def test_push__no_config(mock_module):
     # GIVEN no config file set
     mock_module.set_pyproject_toml({"workflows": []})
@@ -89,16 +104,7 @@ def test_push__happy_path(mock_module, vellum_client, base_command):
     module = mock_module.module
 
     # AND a workflow exists in the module successfully
-    base_dir = os.path.join(temp_dir, *module.split("."))
-    os.makedirs(base_dir, exist_ok=True)
-    workflow_py_file_content = """\
-from vellum.workflows import BaseWorkflow
-
-class ExampleWorkflow(BaseWorkflow):
-    pass
-"""
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
-        f.write(workflow_py_file_content)
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, module)
 
     # AND the push API call returns successfully
     vellum_client.workflows.push.return_value = WorkflowPushResponse(
@@ -137,22 +143,146 @@ class ExampleWorkflow(BaseWorkflow):
     ],
     ids=["push", "workflows_push"],
 )
-def test_push__deployment(mock_module, vellum_client, base_command):
+def test_push__workflow_sandbox_option__existing_id(mock_module, vellum_client, base_command):
     # GIVEN a single workflow configured
     temp_dir = mock_module.temp_dir
     module = mock_module.module
+    existing_workflow_sandbox_id = mock_module.workflow_sandbox_id
 
     # AND a workflow exists in the module successfully
-    base_dir = os.path.join(temp_dir, *module.split("."))
-    os.makedirs(base_dir, exist_ok=True)
-    workflow_py_file_content = """\
-from vellum.workflows import BaseWorkflow
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, module)
 
-class ExampleWorkflow(BaseWorkflow):
-    pass
-"""
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
-        f.write(workflow_py_file_content)
+    # AND the push API call would return successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=existing_workflow_sandbox_id,
+    )
+
+    # WHEN calling `vellum push` with the workflow sandbox option on an existing config
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [module, "--workflow-sandbox-id", existing_workflow_sandbox_id])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # Get the last part of the module path and format it
+    expected_label = mock_module.module.split(".")[-1].replace("_", " ").title()
+    expected_artifact_name = f"{mock_module.module.replace('.', '__')}.tar.gz"
+
+    # AND we should have called the push API with the correct args
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
+    assert call_args["label"] == expected_label
+    assert call_args["workflow_sandbox_id"] == existing_workflow_sandbox_id
+    assert call_args["artifact"].name == expected_artifact_name
+    assert "deplyment_config" not in call_args
+
+    extracted_files = _extract_tar_gz(call_args["artifact"].read())
+    assert extracted_files["workflow.py"] == workflow_py_file_content
+
+
+def test_push__workflow_sandbox_option__existing_no_module(mock_module, vellum_client):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    first_module = mock_module.module
+    second_module = f"{first_module}2"
+    first_workflow_sandbox_id = mock_module.workflow_sandbox_id
+    second_workflow_sandbox_id = str(uuid4())
+
+    # AND the pyproject.toml has two workflow sandboxes configured
+    mock_module.set_pyproject_toml(
+        {
+            "workflows": [
+                {"module": first_module, "workflow_sandbox_id": first_workflow_sandbox_id},
+                {"module": second_module, "workflow_sandbox_id": second_workflow_sandbox_id},
+            ]
+        }
+    )
+
+    # AND a workflow exists for both modules
+    _ensure_workflow_py(temp_dir, first_module)
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, second_module)
+
+    # AND the push API call would return successfully for the second module
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=second_workflow_sandbox_id,
+    )
+
+    # WHEN calling `vellum push` with the workflow sandbox option on the second module
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", "--workflow-sandbox-id", second_workflow_sandbox_id])
+
+    # THEN it should succeed
+    assert result.exit_code == 0
+
+    # Get the last part of the module path and format it
+    expected_label = second_module.split(".")[-1].replace("_", " ").title()
+    expected_artifact_name = f"{second_module.replace('.', '__')}.tar.gz"
+
+    # AND we should have called the push API with the correct args
+    vellum_client.workflows.push.assert_called_once()
+    call_args = vellum_client.workflows.push.call_args.kwargs
+    assert json.loads(call_args["exec_config"])["workflow_raw_data"]["definition"]["name"] == "ExampleWorkflow"
+    assert call_args["label"] == expected_label
+    assert call_args["workflow_sandbox_id"] == second_workflow_sandbox_id
+    assert call_args["artifact"].name == expected_artifact_name
+    assert "deplyment_config" not in call_args
+
+    extracted_files = _extract_tar_gz(call_args["artifact"].read())
+    assert extracted_files["workflow.py"] == workflow_py_file_content
+
+
+def test_push__workflow_sandbox_option__existing_id_different_module(mock_module):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    second_module = f"{module}2"
+    first_workflow_sandbox_id = mock_module.workflow_sandbox_id
+    second_workflow_sandbox_id = str(uuid4())
+    set_pyproject_toml = mock_module.set_pyproject_toml
+
+    # AND the pyproject.toml has two workflow sandboxes configured
+    set_pyproject_toml(
+        {
+            "workflows": [
+                {"module": module, "workflow_sandbox_id": first_workflow_sandbox_id},
+                {"module": second_module, "workflow_sandbox_id": second_workflow_sandbox_id},
+            ]
+        }
+    )
+
+    # AND a workflow exists in both modules successfully
+    _ensure_workflow_py(temp_dir, module)
+    _ensure_workflow_py(temp_dir, second_module)
+
+    # WHEN calling `vellum push` with the first module and the second workflow sandbox id
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", module, "--workflow-sandbox-id", second_workflow_sandbox_id])
+
+    # THEN it should fail
+    assert result.exit_code == 1
+    assert result.exception
+    assert (
+        str(result.exception)
+        == "Multiple workflows found in project to push. Pushing only a single workflow is supported."
+    )
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["push"],
+        ["workflows", "push"],
+    ],
+    ids=["push", "workflows_push"],
+)
+def test_push__deployment(mock_module, vellum_client, base_command):
+    # GIVEN a single workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+
+    # AND a workflow exists in the module successfully
+    workflow_py_file_content = _ensure_workflow_py(temp_dir, module)
 
     # AND the push API call returns successfully
     vellum_client.workflows.push.return_value = WorkflowPushResponse(
@@ -252,16 +382,7 @@ def test_push__strict_option_returns_diffs(mock_module, vellum_client):
     module = mock_module.module
 
     # AND a workflow exists in the module successfully
-    base_dir = os.path.join(temp_dir, *module.split("."))
-    os.makedirs(base_dir, exist_ok=True)
-    workflow_py_file_content = """\
-from vellum.workflows import BaseWorkflow
-
-class ExampleWorkflow(BaseWorkflow):
-    pass
-"""
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
-        f.write(workflow_py_file_content)
+    _ensure_workflow_py(temp_dir, module)
 
     # AND the push API call returns a 4xx response with diffs
     vellum_client.workflows.push.side_effect = ApiError(
@@ -356,16 +477,7 @@ MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
         )
 
     # AND a workflow exists in the module successfully
-    base_dir = os.path.join(temp_dir, *module.split("."))
-    os.makedirs(base_dir, exist_ok=True)
-    workflow_py_file_content = """\
-from vellum.workflows import BaseWorkflow
-
-class ExampleWorkflow(BaseWorkflow):
-    pass
-"""
-    with open(os.path.join(temp_dir, *module.split("."), "workflow.py"), "w") as f:
-        f.write(workflow_py_file_content)
+    _ensure_workflow_py(temp_dir, module)
 
     # AND the push API call returns a new workflow sandbox id
     new_workflow_sandbox_id = str(uuid4())