vellum-ai 0.14.65__tar.gz → 0.14.67__tar.gz

{vellum_ai-0.14.65 → vellum_ai-0.14.67}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.65
+Version: 0.14.67
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -21,7 +21,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Typing :: Typed
 Requires-Dist: Jinja2 (>=3.1.0,<4.0.0)
-Requires-Dist: click (>=8.1.0,<9.0.0)
+Requires-Dist: click (>=8.1.7,<9.0.0)
 Requires-Dist: docker (>=7.1.0,<8.0.0)
 Requires-Dist: httpx (>=0.21.2)
 Requires-Dist: openai (>=1.0.0,<2.0.0)

{vellum_ai-0.14.65 → vellum_ai-0.14.67}/ee/vellum_cli/image_push.py

@@ -9,9 +9,10 @@ import docker
 from docker import DockerClient
 from dotenv import load_dotenv
 
+from vellum.client.core.api_error import ApiError
 from vellum.workflows.vellum_client import create_vellum_client, create_vellum_environment
 from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, load_vellum_cli_config
-from vellum_cli.logger import load_cli_logger
+from vellum_cli.logger import handle_cli_error, load_cli_logger
 
 _SUPPORTED_ARCHITECTURE = "amd64"
 
@@ -82,7 +83,34 @@ def image_push_command(image: str, tags: Optional[List[str]] = None, workspace:
             exit(1)
     else:
         logger.info("Authenticating...")
-        auth = vellum_client.container_images.docker_service_token()
+        try:
+            auth = vellum_client.container_images.docker_service_token()
+        except ApiError as e:
+            if e.status_code == 401 or e.status_code == 403:
+                handle_cli_error(
+                    logger,
+                    title="Authentication failed",
+                    message="Unable to authenticate with Vellum API. Please check your API key.",
+                    suggestion="Make sure your VELLUM_API_KEY environment variable is set correctly.",
+                )
+                return
+            elif e.status_code == 500:
+                handle_cli_error(
+                    logger,
+                    title="Server error",
+                    message="The Vellum API failed with an unexpected server error.",
+                    suggestion="Please try again later and contact support if the problem persists.",
+                )
+                return
+            else:
+                handle_cli_error(
+                    logger,
+                    title="API request failed",
+                    message=f"Failed to get Docker service token from Vellum API (HTTP {e.status_code}).",
+                    suggestion="Please check your configuration and try again. If the problem persists, "
+                    "contact support.",
+                )
+                return
 
         docker_client.login(
             username="oauth2accesstoken",
@@ -143,11 +171,38 @@ def image_push_command(image: str, tags: Optional[List[str]] = None, workspace:
 
     logger.info(f"Updating Vellum metadata and validating image works in our system with image digest: {sha}...")
 
-    vellum_client.container_images.push_container_image(
-        name=image_name,
-        sha=sha,
-        tags=all_tags,
-    )
+    try:
+        vellum_client.container_images.push_container_image(
+            name=image_name,
+            sha=sha,
+            tags=all_tags,
+        )
+    except ApiError as e:
+        if e.status_code == 401 or e.status_code == 403:
+            handle_cli_error(
+                logger,
+                title="Authentication failed",
+                message="Unable to push container image metadata to Vellum API. Please check your API key.",
+                suggestion="Make sure your VELLUM_API_KEY environment variable is set correctly.",
+            )
+            return
+        elif e.status_code == 500:
+            handle_cli_error(
+                logger,
+                title="Server error",
+                message="The Vellum API failed with an unexpected server error while pushing container image metadata.",
+                suggestion="Please try again later and contact support if the problem persists.",
+            )
+            return
+        else:
+            handle_cli_error(
+                logger,
+                title="API request failed",
+                message=f"Failed to push container image metadata to Vellum API (HTTP {e.status_code}).",
+                suggestion="Please check your configuration and try again. If the problem persists, contact support.",
+            )
+            return
+
     logger.info(f"Image successfully pushed as {image_name} to vellum with tags: {all_tags}.")
 
 

{vellum_ai-0.14.65 → vellum_ai-0.14.67}/ee/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_ai-0.14.67/ee/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_ai-0.14.65 → vellum_ai-0.14.67}/ee/vellum_cli/tests/test_pull.py

@@ -745,7 +745,7 @@ def test_pull__same_pull_twice__one_entry_in_lockfile(vellum_client, mock_module
     zip_contents = _zip_file_map({"workflow.py": "print('hello')"})
     responses = iter([zip_contents, zip_contents])
 
-    def workflows_pull_side_effect(*args, **kwargs):
+    def workflows_pull_side_effect(*_args, **_kwargs):
         return iter([next(responses)])
 
     vellum_client.workflows.pull.side_effect = workflows_pull_side_effect
@@ -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
     )