vellum-ai 0.14.64__py3-none-any.whl → 0.14.66__py3-none-any.whl

vellum_cli/pull.py

@@ -20,7 +20,7 @@ from vellum_cli.config import (
     WorkflowDeploymentConfig,
     load_vellum_cli_config,
 )
-from vellum_cli.logger import load_cli_logger
+from vellum_cli.logger import handle_cli_error, load_cli_logger
 
 ERROR_LOG_FILE_NAME = "error.log"
 METADATA_FILE_NAME = "metadata.json"
@@ -201,16 +201,42 @@ def pull_command(
         zip_bytes = b"".join(response)
     except ApiError as e:
         if e.status_code == 401 or e.status_code == 403:
-            raise Exception("Please make sure your `VELLUM_API_KEY` environment variable is set correctly.")
+            handle_cli_error(
+                logger,
+                title="Authentication failed",
+                message="Unable to authenticate with the Vellum API.",
+                suggestion="Please make sure your `VELLUM_API_KEY` environment variable is set correctly and that you have access to this workflow.",  # noqa: E501
+            )
+
+        if e.status_code == 404:
+            handle_cli_error(
+                logger,
+                title="Workflow not found",
+                message=f"The workflow with ID '{pk}' could not be found.",
+                suggestion="Please verify the workflow ID is correct and that you have access to it in your workspace.",
+            )
 
         if e.status_code == 500:
-            raise Exception(
-                "The Pull API failed with an unexpected error. Please try again later and contact support if the problem persists."  # noqa: E501
+            handle_cli_error(
+                logger,
+                title="Server error occurred",
+                message="The Vellum API encountered an internal server error while processing your request.",
+                suggestion="Please try again in a few moments. If the problem persists, contact Vellum support with the workflow ID and timestamp.",  # noqa: E501
+            )
+
+        if e.status_code == 502 or e.status_code == 503 or e.status_code == 504:
+            handle_cli_error(
+                logger,
+                title="Service temporarily unavailable",
+                message="The Vellum API is temporarily unavailable or experiencing high load.",
+                suggestion="Please wait a moment and try again. If the issue continues, check the Vellum status page or contact support.",  # noqa: E501
             )
 
-        # TODO: We should return an Origin header in to validate this case
-        raise Exception(
-            "The API we tried to pull is invalid. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        handle_cli_error(
+            logger,
+            title="API request failed",
+            message=f"The API request failed with status code {e.status_code}.",
+            suggestion="Please verify your `VELLUM_API_URL` environment variable is set correctly and try again.",
         )
 
     zip_buffer = io.BytesIO(zip_bytes)
@@ -299,8 +325,11 @@ def pull_command(
                         logger.info(f"Writing to {target_file}...")
                         target.write(content)
     except zipfile.BadZipFile:
-        raise Exception(
-            "The API we tried to pull from returned an invalid zip file. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        handle_cli_error(
+            logger,
+            title="Invalid response format",
+            message="The API returned an invalid zip file format.",
+            suggestion="Please verify your `VELLUM_API_URL` environment variable is set correctly and try again.",
         )
 
     if include_json:

vellum_cli/tests/test_image_push_error_handling.py

@@ -0,0 +1,184 @@
+import subprocess
+from unittest.mock import patch
+from uuid import uuid4
+
+from click.testing import CliRunner
+
+from vellum.client.core.api_error import ApiError
+from vellum_cli import main as cli_main
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_401_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(status_code=401, body="Unauthorized")
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Authentication failed" in result.output
+    assert "VELLUM_API_KEY" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_500_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(
+        status_code=500, body="Internal Server Error"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Server error" in result.output
+    assert "try again later" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_docker_service_token_other_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+    ]
+
+    vellum_client.container_images.docker_service_token.side_effect = ApiError(
+        status_code=429, body="Too Many Requests"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "API request failed" in result.output
+    assert "HTTP 429" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_401_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(status_code=401, body="Unauthorized")
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Authentication failed" in result.output
+    assert "VELLUM_API_KEY" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_500_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(
+        status_code=500, body="Internal Server Error"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "Server error" in result.output
+    assert "try again later" in result.output
+
+
+@patch("subprocess.run")
+@patch("docker.from_env")
+def test_image_push_container_image_other_error(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    mock_docker_client = mock_docker_from_env.return_value
+    mock_docker_client.images.get.return_value.id = "test-image-id"
+    mock_docker_client.images.push.return_value = ["pushed"]
+
+    mock_run.side_effect = [
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b'[{"RepoDigests": ["test-repo@sha256:abcd1234"]}]'),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = type(
+        "obj", (object,), {"access_token": "345678mnopqr", "organization_id": str(uuid4()), "repository": "myrepo.net"}
+    )()
+
+    vellum_client.container_images.push_container_image.side_effect = ApiError(
+        status_code=429, body="Too Many Requests"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest"])
+
+    assert result.exit_code == 1
+    assert "API request failed" in result.output
+    assert "HTTP 429" in result.output

vellum_cli/tests/test_pull.py

@@ -890,6 +890,7 @@ def test_pull__module_name_from_deployment_name(vellum_client):
         assert f.read() == "print('hello')"
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__invalid_zip_file(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -903,11 +904,12 @@ def test_pull__invalid_zip_file(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == "The API we tried to pull from returned an invalid zip file. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        "Invalid response format" in result.output
+        or "Please verify your `VELLUM_API_URL` environment variable is set correctly" in result.output
     )
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__json_decode_error(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -927,11 +929,12 @@ def test_pull__json_decode_error(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == "The API we tried to pull is invalid. Please make sure your `VELLUM_API_URL` environment variable is set correctly."  # noqa: E501
+        "API request failed" in result.output
+        or "Please verify your `VELLUM_API_URL` environment variable is set correctly" in result.output
     )
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__unauthorized_error_path(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -948,9 +951,10 @@ def test_pull__unauthorized_error_path(vellum_client):
 
     # THEN the command returns an error
     assert result.exit_code == 1
-    assert str(result.exception) == "Please make sure your `VELLUM_API_KEY` environment variable is set correctly."
+    assert "Please make sure your `VELLUM_API_KEY` environment variable is set correctly" in result.output
 
 
+@pytest.mark.usefixtures("mock_module")
 def test_pull__unexpected_error_path(vellum_client):
     workflow_deployment = "test-workflow-deployment-id"
 
@@ -968,9 +972,8 @@ def test_pull__unexpected_error_path(vellum_client):
     # THEN the command returns an error
     assert result.exit_code == 1
     assert (
-        str(result.exception)
-        == """The Pull API failed with an unexpected error. \
-Please try again later and contact support if the problem persists."""
+        "Server error occurred" in result.output
+        or "Please try again in a few moments. If the problem persists, contact Vellum support" in result.output
     )