From 9398790ba23ff03a816d40f7bb959720ecc811ca Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 28 Sep 2026 23:54:32 +0000 Subject: [PATCH] fix(health): splice wildcard prefix onto provider-stripped health check candidates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../health_check_helpers.py | 19 +++++++---- .../proxy/auth/test_model_checks.py | 33 ++++++++++++------- .../test_health_check_helpers.py | 8 +++++ 3 files changed, 42 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 16342728774..7ea9f9496d9 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Final, Literal 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 LIST_BATCHES_SUPPORTED_PROVIDERS, LlmProviders if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -39,16 +39,23 @@ def _ocr_health_check_document(model: str, custom_llm_provider: str) -> Document return native(model, custom_llm_provider) +def _strip_known_provider_prefix(model: str, known_providers: frozenset[str]) -> str: + leading, sep, model_suffix = model.partition("/") + return model_suffix if sep and leading in known_providers else model + + def _wildcard_health_check_models(wildcard_model: str, cheapest_models: Sequence[str]) -> tuple[str, ...]: - """Substitute each candidate id into the wildcard's dynamic segment so the health - check probes a name the wildcard route actually expands, instead of the bare - cost-map id (e.g. `databricks/system.ai.*` + `databricks/databricks-x` -> - `databricks/system.ai.databricks-x`).""" provider_prefix, sep, wildcard_suffix = wildcard_model.partition("/") if not sep or wildcard_suffix == "*": return tuple(cheapest_models) + known_providers: Final = frozenset(provider.value for provider in LlmProviders) + literal_prefix: Final = wildcard_suffix.replace("*", "") + stripped_ids: Final = tuple(_strip_known_provider_prefix(model, known_providers) for model in cheapest_models) return tuple( - f"{provider_prefix}/{wildcard_suffix.replace('*', model.partition('/')[-1], 1)}" for model in cheapest_models + f"{provider_prefix}/{stripped}" + if stripped.startswith(literal_prefix) + else f"{provider_prefix}/{wildcard_suffix.replace('*', stripped, 1)}" + for stripped in stripped_ids ) diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 9204a69a89e..866a8f42bdf 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -12,7 +12,9 @@ def test_get_team_models_for_all_models_and_team_only_models(): model_access_groups = {} include_model_access_groups = False - result = get_team_models(team_models, proxy_model_list, model_access_groups, include_model_access_groups) + result = get_team_models( + team_models, proxy_model_list, model_access_groups, include_model_access_groups + ) combined_models = team_models + proxy_model_list assert set(result) == set(combined_models) @@ -245,7 +247,9 @@ def test_get_key_models_does_not_mutate_input(): ), ], ) -def test_get_complete_model_list_order(key_models, team_models, proxy_model_list, model_list, expected): +def test_get_complete_model_list_order( + key_models, team_models, proxy_model_list, model_list, expected +): """ Test that get_complete_model_list preserves order """ @@ -398,7 +402,9 @@ def test_wildcard_credential_hydration_preserves_deployment_params( captured_params["api_key"] = litellm_params.api_key captured_params["api_version"] = litellm_params.api_version captured_params["credential_name"] = litellm_params.litellm_credential_name - captured_params["has_unexpected_field"] = hasattr(litellm_params, "unexpected_field") + captured_params["has_unexpected_field"] = hasattr( + litellm_params, "unexpected_field" + ) return ["gpt-4o"] monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) @@ -443,7 +449,9 @@ def test_wildcard_custom_prefix_does_not_stack_provider_prefix(monkeypatch): result = get_known_models_from_wildcard( wildcard_model="ollama_server1/*", - litellm_params=LiteLLM_Params(model="ollama_chat/*", custom_llm_provider="ollama_chat"), + litellm_params=LiteLLM_Params( + model="ollama_chat/*", custom_llm_provider="ollama_chat" + ), ) assert result == ["ollama_server1/gemma3:1b", "ollama_server1/llama3:8b"] @@ -470,18 +478,15 @@ def test_wildcard_custom_prefix_keeps_org_segment_for_non_provider_first_segment result = get_known_models_from_wildcard( wildcard_model="my_hf/*", - litellm_params=LiteLLM_Params(model="huggingface/*", custom_llm_provider="huggingface"), + litellm_params=LiteLLM_Params( + model="huggingface/*", custom_llm_provider="huggingface" + ), ) assert result == ["my_hf/meta-llama/Llama-3-8B"] def test_partial_prefix_wildcard_expands_against_stripped_provider_ids(monkeypatch): - """A partial wildcard like ``databricks/system.ai.*`` must splice the prefix before - ``*`` onto the provider-stripped model id, not onto the whole cost-map key. Prepending - to the provider-prefixed key yielded ``databricks/system.ai.databricks/databricks-x``, - an id no upstream accepts. - """ import litellm from litellm.proxy.auth import model_checks from litellm.proxy.auth.model_checks import get_known_models_from_wildcard @@ -953,7 +958,9 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): assert fake_model not in litellm.models_by_provider["vertex_ai"] try: litellm.add_known_models( - model_cost_map={fake_model: {"litellm_provider": "vertex_ai-language-models", "mode": "chat"}} + model_cost_map={ + fake_model: {"litellm_provider": "vertex_ai-language-models", "mode": "chat"} + } ) assert fake_model in litellm.models_by_provider["vertex_ai"] assert litellm.models_by_provider is captured_reference @@ -1001,4 +1008,6 @@ def test_transcribe_is_a_known_provider_for_wildcard_expansion(): assert "transcribe" in litellm.models_by_provider assert "transcribe/StartTranscriptionJob" in litellm.models_by_provider["transcribe"] assert get_provider_models("transcribe") == ["transcribe/StartTranscriptionJob"] - assert get_known_models_from_wildcard("transcribe/*") == ["transcribe/StartTranscriptionJob"] + assert get_known_models_from_wildcard("transcribe/*") == [ + "transcribe/StartTranscriptionJob" + ] 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 4925bcb0b5d..2e010bdf293 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -566,3 +566,11 @@ def test_wildcard_health_check_models_bare_provider_and_star_wildcard_unchanged( candidates: Final = ("databricks/databricks-gpt-5",) assert _wildcard_health_check_models("databricks/*", candidates) == candidates assert _wildcard_health_check_models("*", candidates) == candidates + + +def test_wildcard_health_check_models_partial_prefix_matching_literal_keeps_stripped_id(): + assert _wildcard_health_check_models("openai/gpt-4*", ["gpt-4o-mini"]) == ("openai/gpt-4o-mini",) + + +def test_wildcard_health_check_models_partial_prefix_splices_suffix_around_star(): + assert _wildcard_health_check_models("openai/ft:*", ["gpt-4o-mini"]) == ("openai/ft:gpt-4o-mini",)