From efa9f4ab869c74567133020041022d7bb55b2ead Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 11:08:29 -0800 Subject: [PATCH] fix(health): probe partial wildcards through their model_info.mode A partial wildcard health check went to the chat-only wildcard prober before the mode was read, so embedding and rerank wildcards (voyage/voyage-*, cohere/rerank-*) showed unhealthy with "the cost map has no chat model", and openai/text-embedding-* was probed with a chat call on a chat model. With an explicit non-chat mode, the wildcard now probes the cheapest served cost-map models of that mode through that mode's handler, falling back in rank order and raising the last error when all fail. Unset or chat mode, and modes with no served model of that mode, keep the existing chat probe. --- .../health_check_helpers.py | 132 +++++++++---- .../litellm_core_utils/llm_request_utils.py | 24 ++- litellm/main.py | 32 ++-- litellm/types/utils.py | 5 +- .../test_health_check_helpers.py | 174 ++++++++++++++++++ .../proxy/test_health_check_max_tokens.py | 40 +++- 6 files changed, 340 insertions(+), 67 deletions(-) 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.