diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 3a79eb78b17..43b2eca2fd5 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/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", diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index c23cde6052e..cc49ae574cc 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/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", diff --git a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py index 8f4799e3e7d..c2e2f92ad8a 100644 --- a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py +++ b/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"} diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 72c6f77a4f7..b75d8809c31 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/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" diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx index 30deb79469c..81d89cc41fd 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_modes.tsx +++ b/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" }, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3d657ea53a7..9af332abd2f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/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