diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 2c7d00552b4..03a10806f91 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -9,6 +9,8 @@ from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal +from pydantic import TypeAdapter + from litellm.llms.base_llm.ocr.transformation import DocumentType from litellm.rust_bridge import runtime from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR_HEALTH_CHECK_DOCUMENT @@ -190,7 +192,10 @@ class HealthCheckHelpers: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _metadata_variable_name: Final = "litellm_metadata" - litellm_metadata: Final = HealthCheckHelpers._get_metadata_for_health_check_call() + litellm_metadata: Final[Mapping[str, object]] = { + **HealthCheckHelpers._get_metadata_for_health_check_call(), + "model_info": TypeAdapter(dict[str, object]).validate_python(model_params).get("model_info") or {}, + } model_params[_metadata_variable_name] = litellm_metadata model_params = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( data=model_params, diff --git a/litellm/main.py b/litellm/main.py index 239a85cef4c..e4110455e24 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -9512,6 +9512,7 @@ async def ahealth_check( from litellm.litellm_core_utils.cached_imports import get_litellm_logging_class from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers, default_health_check_mode from litellm.litellm_core_utils.health_check_utils import OPTIONAL_STR + from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS # Use cached import helper to lazy-load Logging class (only loads when function is called) Logging: Final = get_litellm_logging_class() @@ -9556,6 +9557,24 @@ async def ahealth_check( mode = mode or default_health_check_mode( requested_model=requested_model, model=model, custom_llm_provider=custom_llm_provider ) + litellm_logging_obj.model = requested_model + litellm_logging_obj.call_type = { + "completion": "atext_completion", + "embedding": "aembedding", + "audio_speech": "aspeech", + "audio_transcription": "atranscription", + "image_generation": "aimage_generation", + "image_edit": "aimage_edit", + "video_generation": "avideo_generation", + "rerank": "arerank", + "realtime": "arealtime", + "responses": "aresponses", + "anthropic_messages": "anthropic_messages", + "ocr": "aocr", + "decisions": "asystemone", + "extraction": "aextraction", + "batch": ("alist_batches" if custom_llm_provider in LIST_BATCHES_SUPPORTED_PROVIDERS else "acompletion"), + }.get("decisions" if is_decisions_model_mode(mode) else mode, "acompletion") if "*" in model: return await HealthCheckHelpers.ahealth_check_wildcard_models( model=model, diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index de2ed4fd4ba..550a3f66f3b 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -26,6 +26,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS, ) +from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata from litellm.litellm_core_utils.health_check_helpers import native_health_check_mode from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin, is_model_capability_field from litellm.router_strategy.auto_router.auto_router_model_naming import ( @@ -831,6 +832,24 @@ _PROBE_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object]) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +def _same_health_check_model(deployment_model: str | None, probe_params: Mapping[str, object]) -> bool: + model: Final = probe_params.get("model") + if model == deployment_model: + return True + if not isinstance(model, str) or deployment_model is None: + return False + provider_param: Final = probe_params.get("custom_llm_provider") + provider: Final = provider_param if isinstance(provider_param, str) else None + try: + deployment_name, deployment_provider = resolve_model_and_provider_for_metadata( + model=deployment_model, custom_llm_provider=provider + ) + probe_name, probe_provider = resolve_model_and_provider_for_metadata(model=model, custom_llm_provider=provider) + except litellm.exceptions.BadRequestError: + return False + return (deployment_name, deployment_provider) == (probe_name, probe_provider) + + def update_litellm_params_for_health_check( model_info: dict, litellm_params: dict, deployment_model: str | None = None ) -> dict: @@ -895,6 +914,15 @@ def update_litellm_params_for_health_check( if mode == "audio_speech": litellm_params["voice"] = model_info.get("health_check_voice", "alloy") + logging_model_info: Final = _MODEL_INFO_ADAPTER.validate_python( + _PROBE_PARAMS_ADAPTER.validate_python(litellm_params).get("model_info") or {} + ) + probe_logging_model_info: Final = ( + logging_model_info + if _same_health_check_model(base_model_owner, _PROBE_PARAMS_ADAPTER.validate_python(litellm_params)) + else {key: value for key, value in logging_model_info.items() if key == "id"} + ) + # Handle Bedrock region routing format: bedrock/region/model # This is needed because health checks bypass get_llm_provider() for the model param # Issue #15807: Without this, health checks send "region/model" as the model ID to AWS @@ -941,7 +969,12 @@ def update_litellm_params_for_health_check( model_id_params: Final = {} if model_id_pin is None else {"model_id": model_id_pin} return { key: value - for key, value in {**audio_output_params, **model_id_params, **probe_params}.items() + for key, value in { + **audio_output_params, + **model_id_params, + **probe_params, + **({"model_info": probe_logging_model_info} if "model_info" in probe_params else {}), + }.items() if not is_model_capability_field(key) } diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1211c56975b..7c7a1240c82 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -2695,7 +2695,10 @@ async def test_model_connection( stored_model: Final = _OBJECT_MAPPING.validate_python(config_litellm_params).get("model") litellm_params = update_litellm_params_for_health_check( model_info=dict(probe_model_info), - litellm_params=litellm_params, + litellm_params={ + **litellm_params, + "model_info": _OBJECT_MAPPING.validate_python(resolved_model_info or {}), + }, deployment_model=stored_model if isinstance(stored_model, str) else None, ) diff --git a/tests/unit/proxy/test_health_check_functions.py b/tests/unit/proxy/test_health_check_functions.py index 79eec44958f..fb26b9706f6 100644 --- a/tests/unit/proxy/test_health_check_functions.py +++ b/tests/unit/proxy/test_health_check_functions.py @@ -4,10 +4,19 @@ from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx +from pydantic import TypeAdapter +from starlette.requests import Request -from litellm.proxy._types import UserAPIKeyAuth +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.chatgpt import authenticator as chatgpt_authenticator +from litellm.proxy import proxy_server +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.db.health_check_latest import LatestHealthCheckRow +from litellm.proxy.health_check import _same_health_check_model, update_litellm_params_for_health_check from litellm.proxy.health_endpoints._health_endpoints import ( _aggregate_health_check_results, _build_model_param_to_info_mapping, @@ -16,8 +25,10 @@ from litellm.proxy.health_endpoints._health_endpoints import ( _save_health_check_results_if_changed, _save_health_check_to_db, latest_health_checks_endpoint, + test_model_connection as probe_model_connection, ) from litellm.proxy.utils import PrismaClient +from litellm.types.utils import StandardLoggingPayload @pytest.fixture @@ -946,5 +957,213 @@ def test_filter_deployments_to_model_groups(): assert filter_deployments_to_model_groups(model_list, frozenset()) == () +@pytest.mark.asyncio +@pytest.mark.parametrize( + "pricing", [{"base_model": "gpt-5.4"}, {"input_cost_per_token": 0.02, "output_cost_per_token": 0.03}] +) +@pytest.mark.parametrize("deployment_id", [None, "alternate-probe-deployment"]) +@pytest.mark.parametrize("override", ["health_check_model", "health_check_params", "same_model"]) +async def test_alternate_health_check_model_uses_its_own_price( + pricing: dict[str, object], deployment_id: str | None, override: str, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final[asyncio.Queue[StandardLoggingPayload]] = asyncio.Queue() + + class SpendCapture(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + events.put_nowait(TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"])) + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "_async_success_callback", [SpendCapture()]) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr( + proxy_server, + "prisma_client", + PrismaClient( + database_url="postgresql://test:test@localhost/test", proxy_logging_obj=proxy_server.proxy_logging_obj + ), + ) + monkeypatch.setitem( + litellm.model_cost, + "health-probe", + { + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00002, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + monkeypatch.setitem( + litellm.model_cost, + "gpt-5.4", + {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "litellm_provider": "openai", "mode": "chat"}, + ) + if deployment_id is not None: + monkeypatch.setitem( + litellm.model_cost, deployment_id, {**pricing, "litellm_provider": "openai", "mode": "chat"} + ) + params: Final = {"api_base": "https://health-test.example/v1", "api_key": "test-key"} + model_info: Final = { + **pricing, + **( + {"health_check_params": {"model": "openai/health-probe"}} + if override == "health_check_params" + else {"health_check_model": "openai/health-probe"} + ), + **({"id": deployment_id} if deployment_id is not None else {}), + } + payload: Final = { + "id": "chatcmpl-health-probe", + "object": "chat.completion", + "created": 1, + "model": "health-probe", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + with respx.mock: + respx.post("https://health-test.example/v1/chat/completions").mock( + return_value=httpx.Response(200, json=payload) + ) + await litellm.acompletion( + model="openai/health-probe", + messages=[{"role": "user", "content": "test"}], + litellm_metadata={"model_info": model_info} if override == "same_model" else {}, + **params, + ) + normal: Final = await asyncio.wait_for(events.get(), timeout=10) + result: Final = await probe_model_connection( + request=Request({"type": "http", "headers": []}), + mode="chat", + litellm_params={"model": "openai/health-probe" if override == "same_model" else "openai/gpt-5.4", **params}, + model_info=model_info, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["status"] == "success", result + health: Final = await asyncio.wait_for(events.get(), timeout=10) + + assert health["call_type"] == normal["call_type"] + assert health["model"] == normal["model"] + assert health["response_cost"] == normal["response_cost"] + assert health["response_cost"] > 0 + assert health["model_id"] == (deployment_id or "") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stored_model", "requested_model"), + [("openai/health-probe", "health-probe"), ("health-probe", "openai/health-probe")], +) +async def test_equivalent_health_check_model_names_keep_deployment_prices( + stored_model: str, requested_model: str, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final[asyncio.Queue[StandardLoggingPayload]] = asyncio.Queue() + + class SpendCapture(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + events.put_nowait(TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"])) + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "_async_success_callback", [SpendCapture()]) + monkeypatch.setitem( + litellm.model_cost, + "health-probe", + { + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00002, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + params: Final = { + "api_base": "https://equivalent-health.example/v1", + "api_key": "test-key", + "custom_llm_provider": "openai", + } + deployment_id: Final = "priced-health-deployment" + model_info: Final = { + "id": deployment_id, + "mode": "chat", + "input_cost_per_token": 0.02, + "output_cost_per_token": 0.03, + } + router: Final = litellm.Router( + model_list=[ + { + "model_name": "priced-health", + "litellm_params": {"model": stored_model, **params}, + "model_info": model_info, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr( + proxy_server, + "prisma_client", + PrismaClient( + database_url="postgresql://test:test@localhost/test", proxy_logging_obj=proxy_server.proxy_logging_obj + ), + ) + payload: Final = { + "id": "chatcmpl-equivalent-health", + "object": "chat.completion", + "created": 1, + "model": "health-probe", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + with respx.mock: + respx.post("https://equivalent-health.example/v1/chat/completions").mock( + return_value=httpx.Response(200, json=payload) + ) + result: Final = await probe_model_connection( + request=Request({"type": "http", "headers": []}), + mode="chat", + litellm_params={"model": requested_model, "custom_llm_provider": "openai"}, + model_info={"id": deployment_id}, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert result["status"] == "success", result + health: Final = await asyncio.wait_for(events.get(), timeout=10) + + assert health["response_cost"] == 10 * model_info["input_cost_per_token"] + 5 * model_info["output_cost_per_token"] + assert health["model"] == "health-probe" + assert health["model_id"] == deployment_id + + +def test_health_check_without_deployment_model_keeps_probe_but_drops_deployment_prices() -> None: + deployment_info: Final = { + "id": "probe-deployment", + "base_model": "openai/deployment-model", + "input_cost_per_token": 0.02, + "output_cost_per_token": 0.03, + } + probe_model: Final = "openai/health-probe" + + probe_params: Final = update_litellm_params_for_health_check( + model_info={"mode": "chat", "health_check_model": probe_model}, + litellm_params={"model_info": deployment_info}, + ) + + assert probe_params["model"] == probe_model + assert probe_params["model_info"] == {"id": deployment_info["id"]} + + +@pytest.mark.parametrize(("model", "expected"), [("health-probe", True), ("other-health-probe", False)]) +def test_chatgpt_health_check_model_comparison_does_not_access_credentials( + model: str, expected: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + def forbidden_auth_file_access(*args: object, **kwargs: object) -> None: + pytest.fail("Health-check model comparison accessed the ChatGPT auth file") + + monkeypatch.setattr(chatgpt_authenticator, "open", forbidden_auth_file_access, raising=False) + + assert ( + _same_health_check_model("chatgpt/health-probe", {"model": model, "custom_llm_provider": "chatgpt"}) is expected + ) + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 97df9cdbfbf..ddecf43de91 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -52,6 +52,7 @@ from litellm.types.utils import ( Delta, ModelResponseStream, StandardCallbackDynamicParams, + StandardLoggingPayload, StreamingChoices, TextCompletionResponse, Usage, @@ -103,6 +104,63 @@ def preserve_litellm_completion_state(monkeypatch: pytest.MonkeyPatch) -> None: WHITE_PNG: Final = (Path(__file__).parents[1] / "white_100x100.png").read_bytes() +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["image_generation", "embedding"]) +@pytest.mark.parametrize("deployment_id", [None, "health-deployment"]) +async def test_health_check_logs_the_same_cost_and_deployment_as_normal_calls( + mode: str, deployment_id: str | None, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final[asyncio.Queue[StandardLoggingPayload]] = asyncio.Queue() + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem( + litellm.model_cost, + "standard/1024-x-1024/health-image", + {"input_cost_per_image": 0.25, "litellm_provider": "openai", "mode": "image_generation"}, + ) + + class SpendCapture(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + events.put_nowait(TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"])) + + monkeypatch.setattr(litellm, "_async_success_callback", [SpendCapture()]) + model: Final = "openai/health-image" if mode == "image_generation" else "openai/text-embedding-3-small" + payload: Final = ( + {"created": 1, "data": [{"url": "https://example.com/image.png"}]} + if mode == "image_generation" + else { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1], "index": 0}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 10, "total_tokens": 10}, + } + ) + params: Final = { + "model": model, + "api_base": "https://health-test.example/v1", + "api_key": "test-key", + **({"model_info": {"id": deployment_id}} if deployment_id else {}), + **({"size": "1024x1024", "quality": "standard", "n": 1} if mode == "image_generation" else {}), + } + with respx.mock: + respx.post(url__regex=r"https://health-test.example/v1/.*").mock(return_value=httpx.Response(200, json=payload)) + if mode == "image_generation": + await litellm.aimage_generation(**params, prompt="test") + else: + await litellm.aembedding(**params, input=["test"]) + normal: Final = await asyncio.wait_for(events.get(), timeout=10) + result: Final = await litellm.ahealth_check(model_params=params.copy(), mode=mode, prompt="test") + assert "error" not in result, result + health: Final = await asyncio.wait_for(events.get(), timeout=10) + + assert health["call_type"] == normal["call_type"] + assert health["model"] == normal["model"] + assert health["response_cost"] == normal["response_cost"] + assert health["response_cost"] > 0 + assert health["model_id"] == (deployment_id or "") + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize( diff --git a/tests/unit/test_router/test_router_max_parallel_requests.py b/tests/unit/test_router/test_router_max_parallel_requests.py index 1b350f31f23..ef73e4983a9 100644 --- a/tests/unit/test_router/test_router_max_parallel_requests.py +++ b/tests/unit/test_router/test_router_max_parallel_requests.py @@ -1,11 +1,43 @@ from __future__ import annotations +import datetime from typing import Final +from unittest.mock import patch + +import pytest import litellm -import pytest from litellm.utils import calculate_max_parallel_requests +_ROUTER_CLOCK: Final = "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime" +_FIRST_MINUTE: Final = datetime.datetime(2026, 1, 1, 12, 0, 59, tzinfo=datetime.timezone.utc) +_NEXT_MINUTE: Final = datetime.datetime(2026, 1, 1, 12, 1, 0, tzinfo=datetime.timezone.utc) + + +def _rpm_limited_router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "rpm": 1, + }, + } + ], + routing_strategy="usage-based-routing-v2", + enable_pre_call_checks=True, + num_retries=0, + ) + + +_RPM_LIMITED_REQUEST: Final = { + "model": "limited", + "messages": [{"role": "user", "content": "rate limit"}], + "mock_response": "response", +} + @pytest.mark.parametrize( ("max_parallel_requests", "tpm", "rpm", "default_max_parallel_requests", "expected"), @@ -37,32 +69,26 @@ def test_scenario( @pytest.mark.asyncio async def test_max_parallel_requests_rpm_rate_limiting() -> None: - router: Final = litellm.Router( - model_list=[ - { - "model_name": "limited", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "test-key", - "rpm": 1, - }, - } - ], - routing_strategy="usage-based-routing-v2", - enable_pre_call_checks=True, - num_retries=0, - ) - request: Final = { - "model": "limited", - "messages": [{"role": "user", "content": "rate limit"}], - "mock_response": "response", - } + router: Final = _rpm_limited_router() - first_response: Final = await router.acompletion(**request) + with patch(_ROUTER_CLOCK, return_value=_FIRST_MINUTE): + first_response: Final = await router.acompletion(**_RPM_LIMITED_REQUEST) - assert first_response.choices[0].message.content == "response" - with pytest.raises(litellm.RateLimitError): - await router.acompletion(**request) + assert first_response.choices[0].message.content == "response" + with pytest.raises(litellm.RateLimitError): + await router.acompletion(**_RPM_LIMITED_REQUEST) + + +@pytest.mark.asyncio +async def test_rpm_limit_resets_in_the_next_minute() -> None: + router: Final = _rpm_limited_router() + + with patch(_ROUTER_CLOCK, return_value=_FIRST_MINUTE): + await router.acompletion(**_RPM_LIMITED_REQUEST) + with patch(_ROUTER_CLOCK, return_value=_NEXT_MINUTE): + next_minute_response: Final = await router.acompletion(**_RPM_LIMITED_REQUEST) + + assert next_minute_response.choices[0].message.content == "response" @pytest.mark.asyncio