diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 2c7d00552b4..0e03d513b83 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -5,17 +5,22 @@ Helper functions for health check calls. import base64 import itertools import re -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.health_check_utils import create_health_check_response +from litellm.litellm_core_utils.hidden_params import get_hidden_params 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 -from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS +from litellm.types.utils import DECISIONS_MODEL_MODES, LIST_BATCHES_SUPPORTED_PROVIDERS, is_decisions_model_mode if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.llm_request_utils import RankedModel from litellm.types.utils import ImageResponse @@ -34,6 +39,7 @@ BEDROCK_ANY_MODEL_ROUTES: Final = frozenset({"converse", "invoke"}) EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( {"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})} ) +_RESPONSE_HEADERS: Final = TypeAdapter(Mapping[str, object]) def native_health_check_mode( @@ -66,6 +72,15 @@ def default_health_check_mode(requested_model: str, model: str, custom_llm_provi ) +def health_check_handler_mode(mode: str) -> str: + return "decisions" if is_decisions_model_mode(mode) else mode + + +def health_check_response(response: object) -> Mapping[str, object]: + hidden_params: Final = get_hidden_params(response) or {} + return create_health_check_response(_RESPONSE_HEADERS.validate_python(hidden_params.get("headers") or {})) + + def get_image_file_for_health_check() -> bytes: """Return the image used for health checks.""" return base64.b64decode(TEST_IMAGE_BASE64) @@ -117,8 +132,9 @@ def _is_route_segment(custom_llm_provider: str, segment: str, model_after_segmen return base_model is not None and not base_model.startswith(f"{segment}/") -def wildcard_health_check_models(wildcard: str, custom_llm_provider: str) -> tuple[str, ...]: - from litellm.litellm_core_utils.llm_request_utils import rank_chat_models_from_llm_provider +def _wildcard_served_models( + wildcard: str, custom_llm_provider: str, ranked_models: Sequence["RankedModel"] +) -> tuple[str, ...]: from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter provider_prefix: Final = f"{custom_llm_provider}/" @@ -133,22 +149,51 @@ def wildcard_health_check_models(wildcard: str, custom_llm_provider: str) -> tup ) route_prefix: Final = "".join(f"{segments[segment_index]}/" for segment_index in route_segments) model_regex: Final = re.compile(PatternMatchRouter().pattern_to_regex("/".join(segments[len(route_segments) :]))) - ranked_model_ids: Final = tuple( - (ranked_model.model.removeprefix(provider_prefix), ranked_model) - for ranked_model in rank_chat_models_from_llm_provider(custom_llm_provider) + ranked_model_ids: Final = ( + (ranked_model.model.removeprefix(provider_prefix), ranked_model) for ranked_model in ranked_models ) served_model_ids: Final = sorted( (ranked for ranked in ranked_model_ids if model_regex.fullmatch(ranked[0])), key=lambda ranked: (ranked[1].is_past_deprecation_date, ranked[1].route_rank, ranked[1].cost, ranked[0]), ) - served_probes: Final = itertools.islice( - dict.fromkeys(f"{provider_prefix}{route_prefix}{model_id}" for model_id, _ in served_model_ids), - WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT, + return tuple( + itertools.islice( + dict.fromkeys(f"{provider_prefix}{route_prefix}{model_id}" for model_id, _ in served_model_ids), + WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT, + ) ) + + +def wildcard_health_check_models(wildcard: str, custom_llm_provider: str) -> tuple[str, ...]: + from litellm.litellm_core_utils.llm_request_utils import rank_chat_models_from_llm_provider + + provider_prefix: Final = f"{custom_llm_provider}/" + ranked_models: Final = rank_chat_models_from_llm_provider(custom_llm_provider) provider_probes: Final = ( - f"{provider_prefix}{model_id}" for model_id, _ in ranked_model_ids[:WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT] + f"{provider_prefix}{ranked_model.model.removeprefix(provider_prefix)}" + for ranked_model in ranked_models[:WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT] ) - return tuple(dict.fromkeys((*served_probes, *provider_probes))) + return tuple( + dict.fromkeys((*_wildcard_served_models(wildcard, custom_llm_provider, ranked_models), *provider_probes)) + ) + + +def _wildcard_served_models_of_mode(wildcard: str, custom_llm_provider: str, handler_mode: str) -> tuple[str, ...]: + from litellm.litellm_core_utils.llm_request_utils import rank_models_from_llm_provider + + cost_map_modes: Final = DECISIONS_MODEL_MODES if handler_mode == "decisions" else frozenset({handler_mode}) + return _wildcard_served_models( + wildcard, custom_llm_provider, rank_models_from_llm_provider(custom_llm_provider, cost_map_modes) + ) + + +async def _first_answer(probes: Sequence[Callable[[], Awaitable[object]]]) -> object: + try: + return await probes[0]() + except Exception: + if len(probes) == 1: + raise + return await _first_answer(probes[1:]) class HealthCheckHelpers: @@ -158,14 +203,40 @@ class HealthCheckHelpers: custom_llm_provider: str, model_params: dict, litellm_logging_obj: "Logging", - ) -> dict: + mode: str | None = None, + prompt: str | None = None, + input: Sequence[str] | None = None, + ) -> Mapping[str, object]: from litellm import acompletion + handler_mode: Final = health_check_handler_mode(mode or "chat") + served_models: Final = ( + () if handler_mode == "chat" else _wildcard_served_models_of_mode(model, custom_llm_provider, handler_mode) + ) + provider_prefix: Final = f"{custom_llm_provider}/" + mode_probes: Final = tuple( + handler + for served_model in served_models + if ( + handler := HealthCheckHelpers.get_mode_handlers( + model=served_model.removeprefix(provider_prefix), + custom_llm_provider=custom_llm_provider, + model_params={**model_params, "model": served_model}, + prompt=prompt, + input=None if input is None else list(input), + ).get(handler_mode) + ) + is not None + ) + if mode_probes: + return health_check_response(await _first_answer(mode_probes)) + probe_models: Final = wildcard_health_check_models(wildcard=model, custom_llm_provider=custom_llm_provider) if len(probe_models) == 0: raise Exception( f"Unable to health check wildcard model {model} for provider {custom_llm_provider}: the cost map has " - "no chat model for this provider. Set model_info.health_check_model to a model this wildcard serves" + "no chat model for this provider. Set model_info.health_check_model to a model this wildcard serves, " + "or model_info.mode to the mode it serves, such as embedding or rerank" ) model_params["model"] = probe_models[0] model_params["litellm_logging_obj"] = litellm_logging_obj @@ -259,30 +330,10 @@ class HealthCheckHelpers: def get_mode_handlers( model: str, custom_llm_provider: str, - model_params: dict, + model_params: dict[str, object], prompt: str | None = None, - input: list | None = None, - ) -> dict[ - Literal[ - "chat", - "completion", - "embedding", - "audio_speech", - "audio_transcription", - "image_generation", - "image_edit", - "video_generation", - "rerank", - "realtime", - "batch", - "responses", - "anthropic_messages", - "ocr", - "decisions", - "extraction", - ], - Callable, - ]: + input: list[str] | None = None, + ) -> Mapping[str, Callable[[], Awaitable[object]]]: """ Returns a dictionary of mode handlers for health check calls. @@ -305,6 +356,7 @@ class HealthCheckHelpers: from litellm.litellm_core_utils.health_check_utils import ( DECISIONS_CALL_PARAMS, EXTRACTION_CALL_PARAMS, + OPTIONAL_STR, _filter_model_params, ) from litellm.realtime_api.main import realtime_health_check @@ -355,9 +407,9 @@ class HealthCheckHelpers: "realtime": lambda: realtime_health_check( model=model, custom_llm_provider=custom_llm_provider, - api_base=model_params.get("api_base", None), - api_key=model_params.get("api_key", None), - api_version=model_params.get("api_version", None), + api_base=OPTIONAL_STR.validate_python(model_params.get("api_base")), + api_key=OPTIONAL_STR.validate_python(model_params.get("api_key")), + api_version=OPTIONAL_STR.validate_python(model_params.get("api_version")), model_params=model_params, ), "batch": lambda: HealthCheckHelpers._batch_health_check( diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index 84f161e0c62..6f08dee9164 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -1,4 +1,5 @@ from collections.abc import Iterable, Iterator, Mapping +from collections.abc import Set as AbstractSet from dataclasses import dataclass from typing import Final, cast @@ -160,22 +161,27 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1): @dataclass(frozen=True, slots=True) -class RankedChatModel: +class RankedModel: model: str route_rank: int cost: float is_past_deprecation_date: bool -def rank_chat_models_from_llm_provider(custom_llm_provider: str) -> tuple[RankedChatModel, ...]: +_CHAT_MODES: Final = frozenset({"chat"}) + + +def rank_chat_models_from_llm_provider(custom_llm_provider: str) -> tuple[RankedModel, ...]: + return rank_models_from_llm_provider(custom_llm_provider, _CHAT_MODES) + + +def rank_models_from_llm_provider(custom_llm_provider: str, modes: AbstractSet[str]) -> tuple[RankedModel, ...]: from litellm.utils import bundled_provider_models known_models: Final = bundled_provider_models(custom_llm_provider) ranked_models: Final = ( - RankedChatModel( - model=model, route_rank=route_rank, cost=cost, is_past_deprecation_date=is_past_deprecation_date - ) - for model, cost, is_past_deprecation_date in _chat_row_costs(known_models, custom_llm_provider) + RankedModel(model=model, route_rank=route_rank, cost=cost, is_past_deprecation_date=is_past_deprecation_date) + for model, cost, is_past_deprecation_date in _row_costs(known_models, custom_llm_provider, modes) if (route_rank := _route_rank(model, custom_llm_provider)) is not None ) return tuple(sorted(ranked_models, key=lambda ranked_model: (ranked_model.route_rank, ranked_model.cost))) @@ -196,13 +202,15 @@ def _is_past_deprecation_date(cost_map_key: str) -> bool: ) -def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float, bool]]: +def _row_costs( + models: Iterable[str], custom_llm_provider: str, modes: AbstractSet[str] +) -> Iterator[tuple[str, float, bool]]: for model in models: try: model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: continue - if model_info.get("mode") != "chat": + if model_info.get("mode") not in modes: continue yield ( model, diff --git a/litellm/main.py b/litellm/main.py index 1d18100d1d5..04b5b8c8a8c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -104,10 +104,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import ( from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) -from litellm.litellm_core_utils.health_check_utils import ( - _filter_model_params, - create_health_check_response, -) +from litellm.litellm_core_utils.health_check_utils import _filter_model_params from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.mock_functions import ( @@ -175,7 +172,6 @@ from litellm.types.utils import ( ProviderRoutingHints, RawRequestTypedDict, StreamingChoices, - is_decisions_model_mode, ) from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS from litellm.utils import ( @@ -9392,7 +9388,7 @@ async def ahealth_check( model_params: dict, mode: str | None = None, prompt: str | None = None, - input: list | None = None, + input: list[str] | None = None, ): """ Support health checks for different providers. Return remaining rate limit, etc. @@ -9405,7 +9401,12 @@ 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_helpers import ( + HealthCheckHelpers, + default_health_check_mode, + health_check_handler_mode, + health_check_response, + ) from litellm.litellm_core_utils.health_check_utils import OPTIONAL_STR # Use cached import helper to lazy-load Logging class (only loads when function is called) @@ -9448,16 +9449,19 @@ async def ahealth_check( ) model_params["cache"] = {"no-cache": True} # don't used cached responses for making health check calls - mode = mode or default_health_check_mode( - requested_model=requested_model, model=model, custom_llm_provider=custom_llm_provider - ) if "*" in model: return await HealthCheckHelpers.ahealth_check_wildcard_models( model=model, custom_llm_provider=custom_llm_provider, model_params=model_params, litellm_logging_obj=litellm_logging_obj, + mode=mode, + prompt=prompt, + input=input, ) + mode = mode or default_health_check_mode( + requested_model=requested_model, model=model, custom_llm_provider=custom_llm_provider + ) mode_handlers: Final = HealthCheckHelpers.get_mode_handlers( model=model, @@ -9467,13 +9471,9 @@ async def ahealth_check( input=input, ) - handler_mode: Final = "decisions" if is_decisions_model_mode(mode) else mode + handler_mode: Final = health_check_handler_mode(mode) if handler_mode in mode_handlers: - _response: Final = await mode_handlers[handler_mode]() - _response_headers: Final = cast( # cast-ok: provider headers are stored as a string-keyed mapping - Mapping[str, object], (get_hidden_params(_response) or {}).get("headers", {}) or {} - ) - return create_health_check_response(_response_headers) + return health_check_response(await mode_handlers[handler_mode]()) else: raise Exception(f"Mode {mode} not supported. See modes here: https://docs.litellm.ai/docs/proxy/health") except Exception as e: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3ddeb118c33..5277a578b19 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -174,8 +174,11 @@ class LiteLLMCommonStrings(Enum): llm_provider_not_provided = "Unmapped LLM provider for this endpoint. You passed model={model}, custom_llm_provider={custom_llm_provider}. Check supported provider and route: https://docs.litellm.ai/docs/providers" +DECISIONS_MODEL_MODES: Final = frozenset({"decisions", "evaluation"}) + + def is_decisions_model_mode(mode: str | None) -> bool: - return mode in ("decisions", "evaluation") + return mode in DECISIONS_MODEL_MODES SupportedCacheControls: Final = ["ttl", "s-maxage", "no-cache", "no-store"] diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index 950a91beb81..00ce0d9ec70 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -4,6 +4,7 @@ import json import socket import struct import zlib +from collections.abc import Mapping from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -1781,6 +1782,7 @@ async def test_partial_wildcard_health_check_on_a_provider_without_cost_map_chat ) assert "model_info.health_check_model" in response["error"], response["error"] + assert "model_info.mode" in response["error"], response["error"] assert respx_mock.calls.call_count == 0 @@ -1987,3 +1989,175 @@ async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_ assert "error" not in result, result assert json.loads(together_route.calls.last.request.content)["model"] == "fake-together-chat" assert not other_hosts.called + + +def _register_catalog(monkeypatch: pytest.MonkeyPatch, provider: str, rows: Mapping[str, tuple[str, float]]) -> None: + for model, (mode, cost) in rows.items(): + monkeypatch.setitem( + litellm.model_cost, + model, + {"litellm_provider": provider, "mode": mode, "input_cost_per_token": cost, "output_cost_per_token": cost}, + ) + monkeypatch.setitem(litellm.models_by_provider, provider, tuple(rows)) + + +_VOYAGE_EMBEDDINGS_URL: Final = "https://api.voyageai.com/v1/embeddings" +_OPENAI_EMBEDDINGS_URL: Final = "https://api.openai.com/v1/embeddings" + + +def _embedding(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": json.loads(request.content)["model"], + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + + +@pytest.mark.asyncio +async def test_embedding_mode_wildcard_health_check_probes_the_cheapest_served_embedding_model( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + _register_catalog( + monkeypatch, + "voyage", + { + "voyage/voyage-fixture-pricey": ("embedding", 3e-9), + "voyage/voyage-fixture-rerank": ("rerank", 1e-10), + "voyage/voyage-fixture-lite": ("embedding", 1e-9), + }, + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(_VOYAGE_EMBEDDINGS_URL).mock(side_effect=_embedding) + + response: Final = await ahealth_check( + {"model": "voyage/voyage-fixture-*", "api_key": "fake-key"}, mode="embedding", input=["health check"] + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == ["voyage-fixture-lite"] + assert json.loads(upstream.calls.last.request.content)["input"] == ["health check"] + assert "error" not in response, response + + +@pytest.fixture +def openai_embedding_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + _register_catalog( + monkeypatch, + "openai", + { + "openai/gpt-fixture-nano": ("chat", 1e-10), + "openai/text-embedding-fixture-large": ("embedding", 2e-9), + "openai/text-embedding-fixture-small": ("embedding", 1e-9), + }, + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + +def _embedding_without_the_small_model(request: httpx.Request) -> httpx.Response: + if json.loads(request.content)["model"] == "text-embedding-fixture-small": + return httpx.Response(404, json={"error": {"message": "fixture model not found", "type": "invalid_request_error"}}) + return _embedding(request) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("openai_embedding_catalog") +async def test_embedding_mode_wildcard_health_check_falls_back_to_the_next_served_embedding_model( + respx_mock: respx.MockRouter, +) -> None: + upstream: Final = respx_mock.post(_OPENAI_EMBEDDINGS_URL).mock(side_effect=_embedding_without_the_small_model) + + response: Final = await ahealth_check( + {"model": "openai/text-embedding-fixture-*", "api_key": "fake-key"}, mode="embedding" + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [ + "text-embedding-fixture-small", + "text-embedding-fixture-large", + ] + assert "error" not in response, response + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("openai_embedding_catalog") +async def test_embedding_mode_wildcard_health_check_that_fails_on_every_served_model_never_probes_chat( + respx_mock: respx.MockRouter, +) -> None: + embeddings_route: Final = respx_mock.post(_OPENAI_EMBEDDINGS_URL).respond( + 401, json={"error": {"message": "fixture bad key", "type": "invalid_request_error", "code": "invalid_api_key"}} + ) + chat_route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + json=dict(CHAT_COMPLETION_BODY) + ) + + response: Final = await ahealth_check( + {"model": "openai/text-embedding-fixture-*", "api_key": "fake-key"}, mode="embedding" + ) + + assert [json.loads(call.request.content)["model"] for call in embeddings_route.calls] == [ + "text-embedding-fixture-small", + "text-embedding-fixture-large", + ] + assert chat_route.call_count == 0 + assert "fixture bad key" in response["error"], response["error"] + + +@pytest.mark.asyncio +async def test_rerank_mode_wildcard_health_check_probes_a_served_rerank_model_on_the_rerank_api( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + _register_catalog( + monkeypatch, + "cohere", + { + "rerank-fixture-pro": ("rerank", 2e-9), + "rerank-fixture-fast": ("rerank", 1e-9), + "embed-fixture": ("embedding", 1e-10), + }, + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.cohere.com/v2/rerank").respond( + json={ + "id": "rerank-health", + "results": [{"index": 0, "relevance_score": 0.7}], + "meta": {"billed_units": {"search_units": 1}}, + } + ) + + response: Final = await ahealth_check( + {"model": "cohere/rerank-fixture-*", "api_key": "fake-key"}, mode="rerank", prompt="health check" + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == ["rerank-fixture-fast"] + assert json.loads(upstream.calls.last.request.content)["query"] == "health check" + assert "error" not in response, response + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_embedding_mode_wildcard_health_check_without_a_served_embedding_model_keeps_the_chat_probe( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").respond( + json=dict(CHAT_COMPLETION_BODY) + ) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-claude-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + }, + mode="embedding", + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == ["databricks-claude-fixture-a"] + assert response == {} diff --git a/tests/unit/proxy/test_health_check_max_tokens.py b/tests/unit/proxy/test_health_check_max_tokens.py index 26d9fe2fb3d..1e90b87922c 100644 --- a/tests/unit/proxy/test_health_check_max_tokens.py +++ b/tests/unit/proxy/test_health_check_max_tokens.py @@ -12,7 +12,7 @@ import respx import litellm from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers -from litellm.litellm_core_utils.llm_request_utils import RankedChatModel +from litellm.litellm_core_utils.llm_request_utils import RankedModel from litellm.proxy import health_check as hc_module from litellm.proxy.health_check import ( _is_strategy_router_deployment, @@ -89,7 +89,7 @@ async def test_ahealth_check_wildcard_models_respects_max_tokens(): patch( "litellm.litellm_core_utils.llm_request_utils.rank_chat_models_from_llm_provider", return_value=( - RankedChatModel(model="gpt-4o-mini", route_rank=0, cost=0.0, is_past_deprecation_date=False), + RankedModel(model="gpt-4o-mini", route_rank=0, cost=0.0, is_past_deprecation_date=False), ), ), patch("litellm.acompletion", new_callable=AsyncMock), @@ -707,6 +707,42 @@ async def test_a_health_probe_of_a_custom_priced_chat_pin_leaves_plain_calls_on_ assert (chat_route.call_count, responses_route.call_count) == (1, 1) +@pytest.mark.asyncio +async def test_run_model_health_check_probes_an_embedding_mode_wildcard_on_the_embeddings_api_without_max_tokens( + monkeypatch, +): + for model, cost in (("voyage/voyage-fixture-large", 2e-9), ("voyage/voyage-fixture-lite", 1e-9)): + monkeypatch.setitem( + litellm.model_cost, model, {"litellm_provider": "voyage", "mode": "embedding", "input_cost_per_token": cost} + ) + monkeypatch.setitem( + litellm.models_by_provider, "voyage", ("voyage/voyage-fixture-large", "voyage/voyage-fixture-lite") + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=False) as respx_mock: + embeddings_route: Final = respx_mock.post("https://api.voyageai.com/v1/embeddings").respond( + json={ + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "voyage-fixture-lite", + "usage": {"total_tokens": 2}, + } + ) + result: Final = await hc_module._run_model_health_check( + { + "litellm_params": {"model": "voyage/voyage-fixture-*", "api_key": "test-key"}, + "model_info": {"mode": "embedding"}, + } + ) + + assert "error" not in result, result + sent: Final = json.loads(embeddings_route.calls.last.request.content) + assert sent["model"] == "voyage-fixture-lite" + assert "max_tokens" not in sent + + def test_autodetected_embedding_skips_reasoning_effort(): """reasoning_effort must not leak into an embedding probe whose mode is auto-detected.