diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 4404bf6898f..869bab3d827 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -480,6 +480,7 @@ router_settings: | DEFAULT_FAILURE_THRESHOLD_PERCENT | Threshold percentage of failures to cool down a deployment. Default is 0.5 (50%) | DEFAULT_FLUSH_INTERVAL_SECONDS | Default interval in seconds for flushing operations. Default is 5 | DEFAULT_HEALTH_CHECK_INTERVAL | Default interval in seconds for health checks. Default is 300 (5 minutes) +| DEFAULT_HEALTH_CHECK_PROMPT | Default prompt used during health checks for non-image models. Default is "test from litellm" | DEFAULT_IMAGE_HEIGHT | Default height for images. Default is 300 | DEFAULT_IMAGE_TOKEN_COUNT | Default token count for images. Default is 250 | DEFAULT_IMAGE_WIDTH | Default width for images. Default is 300 diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 7df7685f335..6f98265e40a 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -106,6 +106,13 @@ model_list: mode: image_generation # 👈 ADD THIS ``` +#### Custom Health Check Prompt + +By default, health checks use the prompt `"test from litellm"`. You can customize this prompt globally by setting an environment variable, or per-model via config: + +```bash +DEFAULT_HEALTH_CHECK_PROMPT="this is a test prompt" +``` ### Text Completion Models diff --git a/litellm/constants.py b/litellm/constants.py index 7a39596a7ef..17de0e0bd19 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1,6 +1,7 @@ import os from typing import List, Literal +DEFAULT_HEALTH_CHECK_PROMPT = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION = str( os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview") ) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 2f4bb081819..427a16a9801 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -8,7 +8,7 @@ from typing import List, Optional import litellm logger = logging.getLogger(__name__) -from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS +from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS, DEFAULT_HEALTH_CHECK_PROMPT ILLEGAL_DISPLAY_PARAMS = [ "messages", @@ -99,7 +99,7 @@ async def _perform_health_check(model_list: list, details: Optional[bool] = True litellm.ahealth_check( model["litellm_params"], mode=mode, - prompt="test from litellm", + prompt=DEFAULT_HEALTH_CHECK_PROMPT, input=["test from litellm"], ), timeout, diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index 4d55a25f334..8af19f9c18a 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -518,4 +518,67 @@ async def test_ahealth_check_ocr(): mode="ocr", ) print(response) - return response \ No newline at end of file + return response + +@pytest.mark.asyncio +async def test_image_generation_health_check_prompt(monkeypatch): + """Health checks should respect default and environment-configured prompts.""" + + import importlib + import litellm.constants as litellm_constants + import litellm.proxy.health_check as health_check + + def reload_modules(): + reloaded_constants = importlib.reload(litellm_constants) + reloaded_health_check = importlib.reload(health_check) + return reloaded_constants, reloaded_health_check + + async def run_health_check(health_check_module): + health_check_calls = [] + + async def mock_health_check(litellm_params, mode=None, prompt=None, input=None): + health_check_calls.append( + { + "mode": mode, + "prompt": prompt, + "model": litellm_params.get("model"), + } + ) + return {"status": "healthy"} + + model_list = [ + { + "litellm_params": {"model": "dall-e-3", "api_key": "fake-key"}, + "model_info": { + "mode": "image_generation", + }, + } + ] + + with patch( + "litellm.proxy.health_check.litellm.ahealth_check", + side_effect=mock_health_check, + ): + await health_check_module._perform_health_check(model_list) + + return health_check_calls + + # Default prompt is used when env var is unset + monkeypatch.delenv("DEFAULT_HEALTH_CHECK_PROMPT", raising=False) + litellm_constants, health_check = reload_modules() + health_check_calls = await run_health_check(health_check) + + assert len(health_check_calls) == 1 + assert ( + health_check_calls[0]["prompt"] + == litellm_constants.DEFAULT_HEALTH_CHECK_PROMPT + ) + + # Environment override should change the prompt without code changes + override_prompt = "environment override prompt" + monkeypatch.setenv("DEFAULT_HEALTH_CHECK_PROMPT", override_prompt) + litellm_constants, health_check = reload_modules() + health_check_calls = await run_health_check(health_check) + + assert len(health_check_calls) == 1 + assert health_check_calls[0]["prompt"] == override_prompt \ No newline at end of file