diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 9f8878d36f0..428d5ee0062 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -82,7 +82,16 @@ class HealthCheckHelpers: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup _metadata_variable_name: Final = "litellm_metadata" - litellm_metadata: Final = HealthCheckHelpers._get_metadata_for_health_check_call() + existing_metadata: Final = model_params.get(_metadata_variable_name) + health_check_metadata: Final = HealthCheckHelpers._get_metadata_for_health_check_call() + existing_tags: Final = existing_metadata.get("tags", ()) if isinstance(existing_metadata, dict) else () + health_check_tags: Final = health_check_metadata["tags"] + litellm_metadata: Final = health_check_metadata.copy() + if isinstance(existing_metadata, dict): + litellm_metadata.update(existing_metadata) + litellm_metadata.update(health_check_metadata) + health_check_tags[:] = dict.fromkeys((*existing_tags, *health_check_tags)) + litellm_metadata["tags"] = health_check_tags 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/proxy/health_check.py b/litellm/proxy/health_check.py index 219f6f270ed..5f527975413 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -482,6 +482,14 @@ async def _run_model_health_check(model: dict): litellm_params, # any-ok: untyped router config dict ) litellm_params = _update_litellm_params_for_health_check(model_info, litellm_params) + existing_metadata: Final = litellm_params.get("litellm_metadata") + health_check_metadata: Final = ( + existing_metadata.copy() + if isinstance(existing_metadata, dict) + else {} # mutable-ok: health-check metadata is enriched in place + ) + health_check_metadata["model_info"] = model_info + litellm_params["litellm_metadata"] = health_check_metadata timeout: Final = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS return await run_with_timeout( diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py index 97e308d7c3c..1bbfcf2c2d8 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -481,6 +481,42 @@ async def test_run_model_health_check_threads_resolved_mode_to_ahealth_check(): assert probed_params["model"] == "amazon.titan-embed-text-v2:0" +@pytest.mark.asyncio +async def test_run_model_health_check_preserves_deployment_identity_for_spend_tracking(): + fake_ahealth_check = AsyncMock(return_value={}) + model_info = { + "id": "bedrock-health-check-deployment", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "health_check_params": {"litellm_metadata": {"operator_probe_flag": True, "tags": ["operator-health-tag"]}}, + } + model = { + "litellm_params": {"model": "bedrock/claude_platform/claude-haiku-4-5"}, + "model_info": model_info, + } + + with patch.object( # test-quality-ok: observes probe payload before provider execution + hc_module.litellm, "ahealth_check", fake_ahealth_check + ): + await hc_module._run_model_health_check(model) + + probed_params = fake_ahealth_check.call_args.args[0] + assert probed_params["model"] == "claude_platform/claude-haiku-4-5" + assert probed_params["litellm_metadata"]["operator_probe_flag"] is True + assert probed_params["litellm_metadata"]["model_info"] == model_info + + +def test_health_check_tracking_preserves_deployment_identity(): + model_info = {"id": "bedrock-health-check-deployment"} + + updated = HealthCheckHelpers._update_model_params_with_health_check_tracking_information( + {"litellm_metadata": {"model_info": model_info, "tags": ["operator-health-tag"]}} + ) + + assert updated["litellm_metadata"]["model_info"] == model_info + assert "operator-health-tag" in updated["litellm_metadata"]["tags"] + + def test_autodetected_embedding_skips_reasoning_effort(): """reasoning_effort must not leak into an embedding probe whose mode is auto-detected.