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 febf9503a7e..35294e79806 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -9440,6 +9440,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() @@ -9484,6 +9485,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_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 333687b07cd..5e66191f48e 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -2720,6 +2720,7 @@ async def test_model_connection( ) health_check_params: Final = { **_OBJECT_MAPPING.validate_python(litellm_params), + "model_info": _OBJECT_MAPPING.validate_python(resolved_model_info or {}), **_admin_api_key_from_request( _OBJECT_MAPPING.validate_python(request_litellm_params), _OBJECT_MAPPING.validate_python(litellm_params), diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index dbd9ba8a32e..9293ce943f9 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(