From 0372f6f0c7560a116d21be0c7edebc7a2767c70a Mon Sep 17 00:00:00 2001 From: Aditya Singh <60082699+adityasingh2400@users.noreply.github.com> Date: Thu, 21 May 2026 20:06:01 -0700 Subject: [PATCH] Treat None litellm_provider as wildcard in _check_provider_match (#28523) Squash-merged by litellm-agent from adityasingh2400's PR. --- litellm/utils.py | 15 +- .../test_register_model_custom_pricing.py | 161 ++++++++++++++++++ tests/test_litellm/test_utils.py | 28 +++ 3 files changed, 203 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 03ca4bac5b6..18ee811f0f1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2933,6 +2933,13 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 except Exception: existing_model = {} model_cost_key = key + # ``get_model_info`` returns ``litellm_provider: None`` when the + # provider is unknown (e.g. custom deployments registered via + # ``Router.add_deployment``). Persisting that None into + # ``litellm.model_cost`` causes ``_check_provider_match`` to drop + # custom pricing on subsequent cost lookups. + if existing_model.get("litellm_provider") is None: + existing_model.pop("litellm_provider", None) ## override / add new keys to the existing model cost dictionary updated_dictionary = _update_dictionary(existing_model, value) litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary) @@ -5522,9 +5529,15 @@ def _get_model_info_from_model_cost(key: str) -> dict: def _check_provider_match(model_info: dict, custom_llm_provider: Optional[str]) -> bool: """ Check if the model info provider matches the custom provider. + + A missing ``litellm_provider`` key and a ``litellm_provider`` set to + ``None`` both mean "no specific provider constraint" and are treated + as a wildcard match. ``register_model`` may persist ``None`` here via + ``get_model_info`` when a deployment is registered without a provider, + so normalising the two cases keeps custom pricing applied consistently. """ if custom_llm_provider and ( - "litellm_provider" in model_info + model_info.get("litellm_provider") is not None and model_info["litellm_provider"] != custom_llm_provider ): if custom_llm_provider == "vertex_ai" and model_info[ diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index 1efd698fb64..719cb8eecd2 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -190,3 +190,164 @@ def test_build_custom_pricing_entry_time_based(): assert entry["litellm_provider"] == "openai" assert entry["input_cost_per_second"] == 0.01 assert entry["output_cost_per_second"] == 0.02 + + +def test_register_model_strips_none_litellm_provider(): + """``get_model_info`` returns ``litellm_provider: None`` for deployments + registered without a provider (e.g. ``Router.add_deployment`` flows). + ``register_model`` must not persist that None into ``model_cost``, + otherwise ``_check_provider_match`` will drop custom pricing on + subsequent cost lookups. + + Regression test for https://github.com/BerriAI/litellm/issues/28336. + """ + from litellm.utils import _check_provider_match + + model_key = "test-custom-pricing-no-provider-28336" + litellm.model_cost.pop(model_key, None) + + try: + litellm.register_model( + { + model_key: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + } + ) + + registered = litellm.model_cost.get(model_key) + assert registered is not None, f"{model_key} should be in model_cost" + # The key may be absent entirely, but if present it must not be None. + assert ( + "litellm_provider" not in registered + or registered["litellm_provider"] is not None + ) + # Downstream consumers must accept this entry for any provider, + # mirroring what the cost calculator does. + assert _check_provider_match(registered, "openai") is True + assert _check_provider_match(registered, "anthropic") is True + finally: + litellm.model_cost.pop(model_key, None) + + +def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeypatch): + """Directly exercise the strip in ``register_model``. + + The companion test above hits the ``except Exception`` branch where + ``existing_model`` is an empty dict, so the ``pop`` is a no-op. This + test patches ``get_model_info`` to return the failure mode the strip + was added to handle, namely a populated dict whose ``litellm_provider`` + is ``None``. Without the strip, the merged entry in + ``litellm.model_cost`` would carry ``litellm_provider: None`` and + ``_check_provider_match`` would drop custom pricing. + + Regression test for https://github.com/BerriAI/litellm/issues/28336. + """ + from litellm import utils as litellm_utils + from litellm.utils import _check_provider_match + + model_key = "test-strip-none-provider-from-get-model-info-28336" + litellm.model_cost.pop(model_key, None) + + def _fake_get_model_info(model, *args, **kwargs): + assert model == model_key + return { + "key": model_key, + "litellm_provider": None, + "mode": "chat", + "max_tokens": 4096, + } + + # ``register_model`` calls ``get_model_info.cache_clear`` via + # ``_invalidate_model_cost_lowercase_map``, so the replacement must + # expose a no-op ``cache_clear`` attribute. + _fake_get_model_info.cache_clear = lambda: None + monkeypatch.setattr(litellm_utils, "get_model_info", _fake_get_model_info) + + try: + litellm.register_model( + { + model_key: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + } + ) + + registered = litellm.model_cost.get(model_key) + assert registered is not None, f"{model_key} should be in model_cost" + # The strip must have removed the None-valued provider that + # ``get_model_info`` returned. The key may be absent entirely, but + # it must never be present with value ``None``. + assert "litellm_provider" not in registered or ( + registered["litellm_provider"] is not None + ), ( + "register_model failed to strip litellm_provider=None returned " + f"by get_model_info, got {registered.get('litellm_provider')!r}" + ) + # Metadata from the patched ``get_model_info`` must still flow + # through, so we know the strip did not nuke the rest of the entry. + assert registered.get("mode") == "chat" + assert registered.get("max_tokens") == 4096 + # And custom pricing from the registration call must be preserved. + assert registered.get("input_cost_per_token") == 0.001 + assert registered.get("output_cost_per_token") == 0.002 + # Downstream _check_provider_match must accept any provider for + # this entry, mirroring the cost calculator path. + assert _check_provider_match(registered, "openai") is True + assert _check_provider_match(registered, "anthropic") is True + finally: + litellm.model_cost.pop(model_key, None) + + +def test_register_model_router_add_deployment_custom_pricing_applies(): + """End-to-end regression for https://github.com/BerriAI/litellm/issues/28336. + + ``Router.add_deployment`` registers custom pricing without passing + ``litellm_provider``. Cost calculation must still pick up the custom + pricing instead of falling back to the default provider price. + """ + from litellm import Router + + model_key = "router-add-deployment-custom-pricing-28336" + deployment_model = f"openai/{model_key}" + litellm.model_cost.pop(model_key, None) + litellm.model_cost.pop(deployment_model, None) + + router = Router( + model_list=[ + { + "model_name": model_key, + "litellm_params": { + "model": deployment_model, + "api_key": "fake-key-for-registration", + "input_cost_per_token": 0.00042, + "output_cost_per_token": 0.00084, + }, + "model_info": {"id": "deployment-28336"}, + } + ] + ) + + try: + # ``add_deployment`` runs as part of ``Router.__init__``; the + # registered entry must not block ``_check_provider_match`` for + # the deployment's provider. + from litellm.utils import _check_provider_match + + registered_keys = [ + k for k in (deployment_model, model_key) if k in litellm.model_cost + ] + assert registered_keys, ( + "Router.add_deployment did not register custom pricing for " + f"{model_key} / {deployment_model}" + ) + for k in registered_keys: + assert _check_provider_match(litellm.model_cost[k], "openai") is True, ( + f"custom pricing for {k} was dropped by _check_provider_match" + ) + finally: + litellm.model_cost.pop(model_key, None) + litellm.model_cost.pop(deployment_model, None) + del router diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index de286aede93..0efb3083139 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1140,6 +1140,34 @@ def test_check_provider_match(): assert litellm.utils._check_provider_match(model_info, "openai") is False +def test_check_provider_match_none_value_matches_any_provider(): + """ + A ``litellm_provider`` of None must be treated the same as a missing + key: both mean "no provider constraint" and should match any + ``custom_llm_provider``. + + Regression test for https://github.com/BerriAI/litellm/issues/28336. + Before the fix, ``register_model`` persisted ``litellm_provider: None`` + via ``get_model_info`` for deployments registered without a provider + (e.g. ``Router.add_deployment``), which caused ``_check_provider_match`` + to drop custom pricing intermittently. + """ + # Missing key already returned True; None must behave identically. + assert litellm.utils._check_provider_match({}, "openai") is True + assert ( + litellm.utils._check_provider_match({"litellm_provider": None}, "openai") + is True + ) + assert ( + litellm.utils._check_provider_match({"litellm_provider": None}, "anthropic") + is True + ) + # When custom_llm_provider is also None nothing constrains the match. + assert ( + litellm.utils._check_provider_match({"litellm_provider": None}, None) is True + ) + + def test_get_provider_rerank_config(): """ Test the get_provider_rerank_config function for various providers