Merge pull request #38291 from BerriAI/devin_ai_lit6160_health_check_image_edit_mode

fix(health): support `mode: image_edit` in health checks

litellm/litellm_core_utils/health_check_helpers.py

@@ -2,6 +2,7 @@
 Helper functions for health check calls.
 """
 
+import base64
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Final, Literal
 
@@ -13,6 +14,14 @@ if TYPE_CHECKING:
 # Minimal PDF for health checks - base64 encoded 1-page PDF with just "test"
 TEST_PDF_URL = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y="
 
+# Minimal image for health checks - base64 encoded 512x512 solid-gray PNG
+TEST_IMAGE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFlklEQVR42u3VMQEAAAzCMKQjHQ97l0jo0xSAlyIBgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAUgAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAADcDrctaAb6XeXAAAAAASUVORK5CYII="
+
+
+def get_image_file_for_health_check() -> bytes:
+    """Return the image used for health checks."""
+    return base64.b64decode(TEST_IMAGE_BASE64)
+
 
 class HealthCheckHelpers:
     @staticmethod
@@ -127,6 +136,7 @@ class HealthCheckHelpers:
             "audio_speech",
             "audio_transcription",
             "image_generation",
+            "image_edit",
             "video_generation",
             "rerank",
             "realtime",
@@ -185,6 +195,11 @@ class HealthCheckHelpers:
                 **_filter_model_params(model_params=model_params),
                 prompt=prompt,
             ),
+            "image_edit": lambda: litellm.aimage_edit(
+                **_filter_model_params(model_params=model_params),
+                image=get_image_file_for_health_check(),
+                prompt=prompt or "test",
+            ),
             "video_generation": lambda: litellm.avideo_generation(
                 **_filter_model_params(model_params=model_params),
                 prompt=prompt or "test video generation",

litellm/proxy/health_endpoints/_health_endpoints.py

@@ -1796,6 +1796,7 @@ async def test_model_connection(
         "audio_speech",
         "audio_transcription",
         "image_generation",
+        "image_edit",
         "video_generation",
         "batch",
         "rerank",

tests/test_litellm/litellm_core_utils/test_health_check_helpers.py

@@ -12,6 +12,55 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS
 
 
+@pytest.mark.asyncio
+async def test_image_edit_health_check_handler_uses_png_and_prompt():
+    model_params = {"model": "openai/gpt-image-1", "api_key": "sk-test"}
+    mode_handlers = HealthCheckHelpers.get_mode_handlers(
+        model="gpt-image-1",
+        custom_llm_provider="openai",
+        model_params=model_params,
+    )
+
+    assert "image_edit" in mode_handlers
+
+    with patch(  # test-quality-ok: the public health-check path has no dependency injection seam
+        "litellm.aimage_edit", new_callable=AsyncMock, return_value={}
+    ) as mock_aimage_edit:
+        await mode_handlers["image_edit"]()
+        await HealthCheckHelpers.get_mode_handlers(
+            model="gpt-image-1",
+            custom_llm_provider="openai",
+            model_params=model_params,
+            prompt="edit this image",
+        )["image_edit"]()
+
+    assert mock_aimage_edit.call_count == 2
+    default_call = mock_aimage_edit.call_args_list[0].kwargs
+    explicit_call = mock_aimage_edit.call_args_list[1].kwargs
+    assert default_call["model"] == "openai/gpt-image-1"
+    assert default_call["prompt"] == "test"
+    assert explicit_call["prompt"] == "edit this image"
+    image = default_call["image"]
+    assert isinstance(image, bytes)
+    assert image.startswith(b"\x89PNG")
+    assert int.from_bytes(image[16:20], "big") == 512
+    assert int.from_bytes(image[20:24], "big") == 512
+
+
+@pytest.mark.asyncio
+async def test_ahealth_check_supports_image_edit_mode():
+    with patch(  # test-quality-ok: the public health-check path has no dependency injection seam
+        "litellm.aimage_edit", new_callable=AsyncMock, return_value={}
+    ):
+        result = await ahealth_check(
+            {"model": "gpt-image-1", "api_key": "sk-test"},
+            mode="image_edit",
+        )
+
+    assert "error" not in result
+    assert "Mode image_edit not supported" not in str(result)
+
+
 def test_update_model_params_with_health_check_tracking_information():
     """Test _update_model_params_with_health_check_tracking_information adds required tracking info."""
     initial_model_params = {"model": "gpt-3.5-turbo", "api_key": "test_key"}

tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py

@@ -6,11 +6,14 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
+import respx
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, PrismaError
 
+import litellm
 import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module
+from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64
 
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@@ -2674,3 +2677,38 @@ class TestNoRedisWarning:
         ):
             details = await _health_endpoints_module._get_health_readiness_details()
         assert details["show_no_redis_warning"] is False
+
+
+def test_test_model_connection_accepts_image_edit_mode(monkeypatch):
+    """
+    Regression: /health/test_connection rejected mode=image_edit with a 422
+    before image_edit was added to its mode Literal, breaking the UI Test
+    Connection button for image edit deployments.
+    """
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    litellm.in_memory_llm_clients_cache.flush_cache()
+
+    app = FastAPI()
+    app.include_router(_health_endpoints_module.router)
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    client = TestClient(app)
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),  # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam
+        respx.mock(assert_all_called=True) as respx_mock,
+    ):
+        respx_mock.post(host="api.openai.com", path="/v1/images/edits").respond(
+            json={"created": 1700000000, "data": [{"b64_json": TEST_IMAGE_BASE64}]}
+        )
+        response = client.post(
+            "/health/test_connection",
+            json={
+                "mode": "image_edit",
+                "litellm_params": {"model": "openai/gpt-image-2", "api_key": "sk-test"},
+            },
+        )
+
+    assert response.status_code == 200, response.text
+    assert response.json()["status"] == "success"

ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx

@@ -6,6 +6,7 @@ export const TEST_MODES = [
   { value: "audio_speech", label: "Audio Speech - /audio/speech" },
   { value: "audio_transcription", label: "Audio Transcription - /audio/transcriptions" },
   { value: "image_generation", label: "Image Generation - /images/generations" },
+  { value: "image_edit", label: "Image Edit - /images/edits" },
   { value: "video_generation", label: "Video Generation - /videos" },
   { value: "rerank", label: "Rerank - /rerank" },
   { value: "realtime", label: "Realtime - /realtime" },

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -22703,7 +22703,7 @@ export interface components {
              * Mode
              * @description The mode to test the model with. If not provided, auto-detected from model capabilities.
              */
-            mode?: ("chat" | "completion" | "embedding" | "audio_speech" | "audio_transcription" | "image_generation" | "video_generation" | "batch" | "rerank" | "realtime" | "responses" | "ocr") | null;
+            mode?: ("chat" | "completion" | "embedding" | "audio_speech" | "audio_transcription" | "image_generation" | "image_edit" | "video_generation" | "batch" | "rerank" | "realtime" | "responses" | "ocr") | null;
             /**
              * Model Info
              * @description Model info for the health check