Treat None litellm_provider as wildcard in _check_provider_match (#28523)

Squash-merged by litellm-agent from adityasingh2400's PR.

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[

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

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