vellum-ai 0.13.14__tar.gz → 0.13.15__tar.gz

{vellum_ai-0.13.14 → vellum_ai-0.13.15}/PKG-INFO

@@ -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 → vellum_ai-0.13.15}/ee/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_ai-0.13.14 → vellum_ai-0.13.15}/ee/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_ai-0.13.14 → vellum_ai-0.13.15}/ee/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())

{vellum_ai-0.13.14 → vellum_ai-0.13.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vellum-ai"
-version = "0.13.14"
+version = "0.13.15"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-0.13.14 → vellum_ai-0.13.15}/src/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