From dbcacc7ab7880d5115376cd6c22d22898c50d08f Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:34:57 -0700 Subject: [PATCH] fix(health): probe only models a partial wildcard serves (#49060) * fix(health): probe only models a partial wildcard serves A partial wildcard health check now probes the cheapest cost-map chat models the wildcard matches, keeps provider route prefixes such as Bedrock converse/ on the probe, and falls back only within the wildcard. A wildcard with no matching cost-map chat model reports an error naming model_info.health_check_model * fix(health): keep the provider-wide probe for an unmatched wildcard, skip retired models, and ignore Bedrock ARNs A partial wildcard with no served cost-map chat model falls back to the provider's cheapest chat models instead of failing, and fails with a message naming model_info.health_check_model only when the provider has none. The cheapest-model picker skips models past their cost-map deprecation_date, and a Bedrock ARN segment is never read as a route prefix. * fix(health): fall back to the provider's cheapest probe when a wildcard's own models fail A partial wildcard now probes up to three of its own cheapest chat models first and then the provider-wide cheapest ones, all from one cost-map scan, so stale or unpriced rows the wildcard matches cannot turn a healthy deployment red. Retired models are read by their resolved cost-map key, both deprecation checks use the UTC date, Bedrock only treats converse, invoke, and region names as route segments, and equal-price ties break by model id * fix(health): cap wildcard health probes at three and rank retired cost-map rows last A partial wildcard now probes two of its own cheapest cost-map chat models and then one provider-wide model, so a failing deployment makes no more calls than before. Rows past their cost-map deprecation date go last instead of being dropped, since some still serve, and the shared cheapest-model picker returns what it returned before. * fix(health): probe a wildcard's own models first, then the same three models main probes A partial wildcard's health check tries its two cheapest own cost-map models, then the provider's three cheapest models main already probes, so a row main shows healthy stays healthy. The deprecation check uses the server's local date, like the Baseten listing * fix(health): keep main's provider-wide probe order so equal-price ties match main --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../health_check_helpers.py | 80 +++- .../litellm_core_utils/llm_request_utils.py | 58 ++- .../test_health_check_helpers.py | 343 +++++++++++++++++- .../test_llm_request_utils.py | 84 ++++- .../proxy/test_health_check_max_tokens.py | 7 +- 5 files changed, 545 insertions(+), 27 deletions(-) diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 940730eeb9f..e40299a218f 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -3,6 +3,8 @@ Helper functions for health check calls. """ import base64 +import itertools +import re from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal @@ -26,6 +28,9 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = ( ) ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16 +WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT: Final = 2 +WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT: Final = 3 +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"})} ) @@ -90,6 +95,62 @@ def _ocr_health_check_document(model: str, custom_llm_provider: str) -> Document return native(model, custom_llm_provider) +def _provider_base_model(custom_llm_provider: str, model: str) -> str | None: + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + try: + model_info: Final = ProviderConfigManager.get_provider_model_info( + model=None, provider=LlmProviders(custom_llm_provider) + ) + return None if model_info is None else model_info.get_base_model(model) + except (ValueError, NotImplementedError): + return None + + +def _is_route_segment(custom_llm_provider: str, segment: str, model_after_segment: str) -> bool: + if custom_llm_provider == "bedrock": + from litellm.llms.bedrock.common_utils import get_all_bedrock_regions + + return segment in BEDROCK_ANY_MODEL_ROUTES or segment in get_all_bedrock_regions() + base_model: Final = _provider_base_model(custom_llm_provider, f"{segment}/{model_after_segment}") + 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 + from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + + provider_prefix: Final = f"{custom_llm_provider}/" + segments: Final = wildcard.removeprefix(provider_prefix).split("/") + route_segments: Final = tuple( + itertools.takewhile( + lambda segment_index: _is_route_segment( + custom_llm_provider, segments[segment_index], "/".join(segments[segment_index + 1 :]) + ), + range(len(segments) - 1), + ) + ) + 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) + ) + 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, + ) + provider_probes: Final = ( + f"{provider_prefix}{model_id}" for model_id, _ in ranked_model_ids[:WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT] + ) + return tuple(dict.fromkeys((*served_probes, *provider_probes))) + + class HealthCheckHelpers: @staticmethod async def ahealth_check_wildcard_models( @@ -99,23 +160,16 @@ class HealthCheckHelpers: litellm_logging_obj: "Logging", ) -> dict: from litellm import acompletion - from litellm.litellm_core_utils.llm_request_utils import ( - pick_cheapest_chat_models_from_llm_provider, - ) - # this is a wildcard model, we need to pick a random model from the provider - cheapest_models = pick_cheapest_chat_models_from_llm_provider(custom_llm_provider=custom_llm_provider, n=3) - if len(cheapest_models) == 0: + 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 for provider {custom_llm_provider}. Add a model on your config.yaml or contribute here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" + 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" ) - if len(cheapest_models) > 1: - fallback_models = cheapest_models[1:] # Pick the last 2 models from the shuffled list - else: - fallback_models = None - model_params["model"] = cheapest_models[0] + model_params["model"] = probe_models[0] model_params["litellm_logging_obj"] = litellm_logging_obj - model_params["fallbacks"] = fallback_models + model_params["fallbacks"] = list(probe_models[1:]) or None model_params["max_tokens"] = model_params.get("max_tokens", 16) # GPT-5 models require max_output_tokens >= 16 await acompletion(**model_params) return {} diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index db881aa16c6..9c2ec0b8984 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -1,6 +1,11 @@ +import datetime from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass from typing import Final, cast +from pydantic import TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict + import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider @@ -151,19 +156,52 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1): Returns: list[str]: A list of the n cheapest chat models. """ - if custom_llm_provider not in litellm.models_by_provider: - return [] + return [ranked_model.model for ranked_model in rank_chat_models_from_llm_provider(custom_llm_provider)[:n]] + +@dataclass(frozen=True, slots=True) +class RankedChatModel: + 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, ...]: known_models: Final[Iterable[str]] = litellm.models_by_provider.get(custom_llm_provider, ()) - ranked_rows: Final = ( - (route_rank, cost, model) - for model, cost in _chat_row_costs(known_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) if (route_rank := _route_rank(model, custom_llm_provider)) is not None ) - return [model for _, _, model in sorted(ranked_rows, key=lambda row: row[:2])[:n]] + return tuple(sorted(ranked_models, key=lambda ranked_model: (ranked_model.route_rank, ranked_model.cost))) -def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float]]: +class _CostMapDeprecation(TypedDict): + deprecation_date: NotRequired[ReadOnly[object]] + + +_COST_MAP_DEPRECATION: Final = TypeAdapter(_CostMapDeprecation) + + +def _is_past_deprecation_date(cost_map_key: str) -> bool: + deprecation_date: Final = _COST_MAP_DEPRECATION.validate_python( + litellm.model_cost.get(cost_map_key) or {} # pyright: ignore[reportUnknownMemberType] # untyped cost map JSON + ).get("deprecation_date") + if not isinstance(deprecation_date, str): + return False + try: + return ( + datetime.date.fromisoformat(deprecation_date) + <= datetime.datetime.now(datetime.timezone.utc).astimezone().date() + ) + except ValueError: + return False + + +def _chat_row_costs(models: Iterable[str], custom_llm_provider: 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) @@ -171,7 +209,11 @@ def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator continue if model_info.get("mode") != "chat": continue - yield model, (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0) + yield ( + model, + (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0), + _is_past_deprecation_date(model_info["key"]), + ) def _route_rank(model: str, custom_llm_provider: str) -> int | None: 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 6979d72ff14..950a91beb81 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -8,6 +8,7 @@ from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import responses import respx @@ -19,6 +20,7 @@ from litellm.litellm_core_utils.health_check_helpers import ( HealthCheckHelpers, default_health_check_mode, native_health_check_mode, + wildcard_health_check_models, ) from litellm.main import ahealth_check from litellm.proxy._types import UserAPIKeyAuth @@ -1449,6 +1451,339 @@ async def test_groq_wildcard_health_check_uses_a_concrete_model( assert response == {} +_WILDCARD_CATALOG: Final = MappingProxyType( + { + "databricks": ( + "databricks/databricks-claude-fixture-d", + "databricks/databricks-claude-fixture-c", + "databricks/databricks-claude-fixture-b", + "databricks/databricks-claude-fixture-a", + "databricks/databricks-cheap-chat", + "databricks/databricks-claude-fixture-retired", + ), + "bedrock": ("meta.llama-fixture", "anthropic.claude-fixture"), + "cloudflare": ("cloudflare/@cf/meta/fixture-chat",), + "vertex_ai": ("claude-fixture@default",), + } +) +_WILDCARD_CATALOG_COST: Final = MappingProxyType( + { + "databricks/databricks-claude-fixture-retired": 1e-10, + "databricks/databricks-cheap-chat": 1e-9, + "databricks/databricks-claude-fixture-a": 2e-9, + "databricks/databricks-claude-fixture-b": 3e-9, + "databricks/databricks-claude-fixture-c": 4e-9, + "databricks/databricks-claude-fixture-d": 5e-9, + "meta.llama-fixture": 1e-9, + "anthropic.claude-fixture": 2e-9, + "cloudflare/@cf/meta/fixture-chat": 1e-9, + "claude-fixture@default": 1e-9, + } +) +_DATABRICKS_FIXTURE_API_BASE: Final = "https://fixture.cloud.databricks.com/serving-endpoints" + + +@pytest.fixture +def wildcard_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + providers: Final = {model: provider for provider, models in _WILDCARD_CATALOG.items() for model in models} + litellm.register_model( + { + model: { + "litellm_provider": providers[model], + "mode": "chat", + "max_tokens": 1024, + "input_cost_per_token": cost, + "output_cost_per_token": cost, + **({"deprecation_date": "2000-01-01"} if model.endswith("-retired") else {}), + } + for model, cost in _WILDCARD_CATALOG_COST.items() + } + ) + for provider, models in _WILDCARD_CATALOG.items(): + monkeypatch.setitem(litellm.models_by_provider, provider, list(models)) + + +_DATABRICKS_PROVIDER_PROBES: Final = ( + "databricks/databricks-claude-fixture-retired", + "databricks/databricks-cheap-chat", + "databricks/databricks-claude-fixture-a", +) +_BEDROCK_PROVIDER_PROBES: Final = ("bedrock/meta.llama-fixture", "bedrock/anthropic.claude-fixture") +_DATABRICKS_SERVED_CLAUDE_PROBES: Final = ( + "databricks/databricks-claude-fixture-a", + "databricks/databricks-claude-fixture-b", +) + + +@pytest.mark.parametrize( + ("custom_llm_provider", "wildcard", "expected"), + [ + ("databricks", "databricks-claude-*", (*_DATABRICKS_SERVED_CLAUDE_PROBES, *_DATABRICKS_PROVIDER_PROBES[:2])), + ( + "databricks", + "databricks/databricks-claude-*", + (*_DATABRICKS_SERVED_CLAUDE_PROBES, *_DATABRICKS_PROVIDER_PROBES[:2]), + ), + ("databricks", "*", (*_DATABRICKS_PROVIDER_PROBES[1:], _DATABRICKS_PROVIDER_PROBES[0])), + ("databricks", "databricks-claude-fixture-r*", _DATABRICKS_PROVIDER_PROBES), + ("databricks", "databricks-llama-*", _DATABRICKS_PROVIDER_PROBES), + ("databricks", "*-fixture", _DATABRICKS_PROVIDER_PROBES), + ("databricks", "team-endpoints/*", _DATABRICKS_PROVIDER_PROBES), + ("bedrock", "converse/anthropic.*", ("bedrock/converse/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)), + ( + "bedrock", + "converse/*", + ( + "bedrock/converse/meta.llama-fixture", + "bedrock/converse/anthropic.claude-fixture", + *_BEDROCK_PROVIDER_PROBES, + ), + ), + ("bedrock", "invoke/anthropic.*", ("bedrock/invoke/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)), + ("bedrock", "us-east-1/anthropic.*", ("bedrock/us-east-1/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)), + ("bedrock", "agent/anthropic.*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "nova/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "nova-2/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "mantle/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "openai/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "chat_completions/*", _BEDROCK_PROVIDER_PROBES), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/*", _BEDROCK_PROVIDER_PROBES), + ( + "bedrock", + "converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/*", + _BEDROCK_PROVIDER_PROBES, + ), + ("cloudflare", "@cf/*", ("cloudflare/@cf/meta/fixture-chat",)), + ("vertex_ai", "claude-*", ("vertex_ai/claude-fixture@default",)), + ("vertex_ai", "meta/*", ("vertex_ai/claude-fixture@default",)), + ("fixture-gateway", "team/*", ()), + ], +) +@pytest.mark.usefixtures("wildcard_catalog") +def test_wildcard_health_check_probes_the_cheapest_served_chat_models_before_the_provider_cheapest( + custom_llm_provider: str, wildcard: str, expected: tuple[str, ...] +) -> None: + assert wildcard_health_check_models(wildcard=wildcard, custom_llm_provider=custom_llm_provider) == expected + + +def _fail_first_databricks_claude_probe(request: httpx.Request) -> httpx.Response: + model: Final = json.loads(request.content)["model"] + if model == "databricks-claude-fixture-a": + return httpx.Response(500, json={"error_code": "INTERNAL_ERROR", "message": "fixture outage"}) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-health", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "2"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_partial_wildcard_health_check_falls_back_within_the_wildcard_before_any_other_model( + 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").mock( + side_effect=_fail_first_databricks_claude_probe + ) + + 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?"}], + } + ) + + probed_models: Final = [json.loads(call.request.content)["model"] for call in upstream.calls] + assert probed_models[0] == "databricks-claude-fixture-a" + assert probed_models[-1] == "databricks-claude-fixture-b" + assert all(model.startswith("databricks-claude-") for model in probed_models) + assert "databricks-claude-fixture-a" not in probed_models[1:] + assert response == {} + + +def _answer_only_the_retired_databricks_claude_model(request: httpx.Request) -> httpx.Response: + model: Final = json.loads(request.content)["model"] + if model != "databricks-claude-fixture-retired": + return httpx.Response(404, json={"error_code": "ENDPOINT_NOT_FOUND", "message": f"{model} does not exist"}) + return _fail_first_databricks_claude_probe(request) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_partial_wildcard_health_check_stays_healthy_while_any_model_the_provider_probes_answers( + 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").mock( + side_effect=_answer_only_the_retired_databricks_claude_model + ) + + 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?"}], + } + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [ + "databricks-claude-fixture-a", + "databricks-claude-fixture-b", + "databricks-claude-fixture-retired", + ] + assert response == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_partial_wildcard_health_check_that_fails_on_every_probe_tries_its_own_models_then_the_provider_probes( + 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( + 401, json={"error_code": "UNAUTHENTICATED", "message": "fixture bad token"} + ) + + 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?"}], + } + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [ + "databricks-claude-fixture-a", + "databricks-claude-fixture-b", + "databricks-claude-fixture-retired", + "databricks-cheap-chat", + ] + assert "error" in response + + +_DATABRICKS_EQUAL_PRICE_MODELS: Final = ( + "databricks/databricks-free-d", + "databricks/databricks-free-c", + "databricks/databricks-free-b", + "databricks/databricks-free-a", +) + + +def _register_equal_price_databricks_models(monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, ...]) -> None: + litellm.register_model( + { + model: { + "litellm_provider": "databricks", + "mode": "chat", + "max_tokens": 1024, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + for model in listed_order + } + ) + monkeypatch.setitem(litellm.models_by_provider, "databricks", list(listed_order)) + + +@pytest.mark.parametrize("listed_order", [_DATABRICKS_EQUAL_PRICE_MODELS, _DATABRICKS_EQUAL_PRICE_MODELS[::-1]]) +def test_partial_wildcard_health_check_orders_its_own_equal_price_models_by_id_whatever_the_listing_order( + monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, ...] +) -> None: + _register_equal_price_databricks_models(monkeypatch, listed_order) + + assert wildcard_health_check_models(wildcard="databricks/databricks-free-*", custom_llm_provider="databricks")[ + :2 + ] == ("databricks/databricks-free-a", "databricks/databricks-free-b") + + +def _answer_only_the_first_listed_databricks_free_model(request: httpx.Request) -> httpx.Response: + model: Final = json.loads(request.content)["model"] + if model != "databricks-free-d": + return httpx.Response(404, json={"error_code": "ENDPOINT_NOT_FOUND", "message": f"{model} does not exist"}) + return _fail_first_databricks_claude_probe(request) + + +@pytest.mark.asyncio +async def test_wildcard_health_check_falls_back_to_equal_price_provider_models_in_the_provider_listing_order( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + _register_equal_price_databricks_models(monkeypatch, _DATABRICKS_EQUAL_PRICE_MODELS) + 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").mock( + side_effect=_answer_only_the_first_listed_databricks_free_model + ) + + 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?"}], + } + ) + + assert [json.loads(call.request.content)["model"] for call in upstream.calls] == ["databricks-free-d"] + assert response == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("wildcard_catalog") +async def test_partial_wildcard_health_check_without_a_served_cost_map_model_probes_the_provider_cheapest_chat_model( + 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").mock( + side_effect=_fail_first_databricks_claude_probe + ) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-llama-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + } + ) + + assert json.loads(upstream.calls[0].request.content)["model"] == "databricks-claude-fixture-retired" + assert response == {} + + +@pytest.mark.asyncio +async def test_partial_wildcard_health_check_on_a_provider_without_cost_map_chat_models_names_health_check_model( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setitem(litellm.models_by_provider, "databricks", []) + + response: Final = await ahealth_check( + { + "model": "databricks/databricks-llama-*", + "api_key": "fake-key", + "api_base": _DATABRICKS_FIXTURE_API_BASE, + "messages": [{"role": "user", "content": "What is 1 + 1?"}], + } + ) + + assert "model_info.health_check_model" in response["error"], response["error"] + assert respx_mock.calls.call_count == 0 + + @pytest.mark.asyncio async def test_cohere_rerank_health_check_returns_a_successful_response( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter @@ -1632,7 +1967,12 @@ async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_ monkeypatch.setitem( litellm.model_cost, "fake-together-chat", - {"input_cost_per_token": 1e-9, "output_cost_per_token": 1e-9, "mode": "chat", "litellm_provider": "together_ai"}, + { + "input_cost_per_token": 1e-9, + "output_cost_per_token": 1e-9, + "mode": "chat", + "litellm_provider": "together_ai", + }, ) monkeypatch.setitem(litellm.models_by_provider, "together_ai", ("fake-together-chat",)) litellm.in_memory_llm_clients_cache.flush_cache() @@ -1647,4 +1987,3 @@ 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 - diff --git a/tests/unit/litellm_core_utils/test_llm_request_utils.py b/tests/unit/litellm_core_utils/test_llm_request_utils.py index 0a2bd5cec1b..e7350fae395 100644 --- a/tests/unit/litellm_core_utils/test_llm_request_utils.py +++ b/tests/unit/litellm_core_utils/test_llm_request_utils.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_request_utils import ( ensure_extra_body_is_safe, flatten_form_field_values, pick_cheapest_chat_models_from_llm_provider, + rank_chat_models_from_llm_provider, serialize_multipart_form_fields, ) @@ -149,6 +150,24 @@ def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyP "deepseek/fake-cheapest-embedding": {"input_cost_per_token": 1e-13, "mode": "embedding"}, "openai/fake-free-chat": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"}, "deepseek/fake-cheap-chat": {"input_cost_per_token": 1e-12, "output_cost_per_token": 1e-12, "mode": "chat"}, + "deepseek/fake-retired-chat": { + "input_cost_per_token": 1e-14, + "output_cost_per_token": 1e-14, + "mode": "chat", + "deprecation_date": "2000-01-01", + }, + "deepseek/fake-retiring-chat": { + "input_cost_per_token": 2e-12, + "output_cost_per_token": 2e-12, + "mode": "chat", + "deprecation_date": "2999-01-01", + }, + "deepseek/fake-misdated-chat": { + "input_cost_per_token": 2.5e-12, + "output_cost_per_token": 2.5e-12, + "mode": "chat", + "deprecation_date": "next spring", + }, "fake-free-chat-without-provider": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"}, } for name, row in rows.items(): @@ -156,12 +175,73 @@ def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyP monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, tuple(rows)) -def test_pick_cheapest_chat_models_prefers_provider_routed_ids_and_skips_ids_routed_elsewhere( +def test_pick_cheapest_chat_models_prefers_provider_routed_ids_skips_ids_routed_elsewhere_and_keeps_retired_ones( deepseek_rows_whose_cheapest_ids_route_elsewhere: None, ) -> None: assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=10) == [ + "deepseek/fake-retired-chat", "deepseek/fake-cheap-chat", + "deepseek/fake-retiring-chat", + "deepseek/fake-misdated-chat", "deepseek/fake-pricey-chat", "fake-free-chat-without-provider", ] - assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-cheap-chat"] + assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-retired-chat"] + + +def test_rank_chat_models_flags_only_the_rows_whose_cost_map_deprecation_date_has_passed( + deepseek_rows_whose_cheapest_ids_route_elsewhere: None, +) -> None: + assert { + ranked_model.model: ranked_model.is_past_deprecation_date + for ranked_model in rank_chat_models_from_llm_provider(PICKER_PROVIDER) + } == { + "deepseek/fake-retired-chat": True, + "deepseek/fake-cheap-chat": False, + "deepseek/fake-retiring-chat": False, + "deepseek/fake-misdated-chat": False, + "deepseek/fake-pricey-chat": False, + "fake-free-chat-without-provider": False, + } + + +@pytest.mark.parametrize( + "listed_order", + [("deepseek/fake-tie-b", "deepseek/fake-tie-a"), ("deepseek/fake-tie-a", "deepseek/fake-tie-b")], +) +def test_pick_cheapest_chat_models_keeps_the_provider_listing_order_for_equal_price_models( + monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, str] +) -> None: + for name in listed_order: + monkeypatch.setitem( + litellm.model_cost, + name, + { + "input_cost_per_token": 1e-12, + "output_cost_per_token": 1e-12, + "mode": "chat", + "litellm_provider": PICKER_PROVIDER, + }, + ) + monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, listed_order) + + assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=2) == list(listed_order) + + +def test_rank_chat_models_flags_a_retired_model_the_provider_lists_without_its_cost_map_prefix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rows: Final = { + "vertex_ai/fake-retired-vertex-chat": {"input_cost_per_token": 1e-14, "deprecation_date": "2000-01-01"}, + "vertex_ai/fake-live-vertex-chat": {"input_cost_per_token": 1e-12, "deprecation_date": "2999-01-01"}, + } + for name, row in rows.items(): + monkeypatch.setitem(litellm.model_cost, name, {**row, "mode": "chat", "litellm_provider": "vertex_ai"}) + monkeypatch.setitem( + litellm.models_by_provider, "vertex_ai", tuple(name.removeprefix("vertex_ai/") for name in rows) + ) + + assert [ + (ranked_model.model, ranked_model.is_past_deprecation_date) + for ranked_model in rank_chat_models_from_llm_provider("vertex_ai") + ] == [("fake-retired-vertex-chat", True), ("fake-live-vertex-chat", False)] diff --git a/tests/unit/proxy/test_health_check_max_tokens.py b/tests/unit/proxy/test_health_check_max_tokens.py index e48bf4edc09..df0429316ed 100644 --- a/tests/unit/proxy/test_health_check_max_tokens.py +++ b/tests/unit/proxy/test_health_check_max_tokens.py @@ -11,6 +11,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.proxy import health_check as hc_module from litellm.proxy.health_check import ( _is_strategy_router_deployment, @@ -85,8 +86,10 @@ async def test_ahealth_check_wildcard_models_respects_max_tokens(): """ with ( patch( - "litellm.litellm_core_utils.llm_request_utils.pick_cheapest_chat_models_from_llm_provider", - return_value=["gpt-4o-mini"], + "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), + ), ), patch("litellm.acompletion", new_callable=AsyncMock), ):