diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 94e74b20297..42b8a8e8d6b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13468,7 +13468,10 @@ def _enrich_model_info_with_litellm_data( except Exception: litellm_model_info = {} for k, v in litellm_model_info.items(): - if k not in model_info: + # A key that is present but None must be treated as unset. ModelInfo's + # set_model_info validator writes None over absent cost fields, so + # "k not in model_info" leaves a priced model reading as unpriced. + if model_info.get(k) is None: model_info[k] = v model["model_info"] = model_info # don't return the api key / vertex credentials @@ -14934,7 +14937,7 @@ def _get_proxy_model_info(model: dict) -> dict: except Exception: litellm_model_info = {} for k, v in litellm_model_info.items(): - if k not in model_info: + if model_info.get(k) is None: model_info[k] = v model["model_info"] = model_info # don't return the llm credentials diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e058a4f6396..d9d5f301160 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -31,7 +31,12 @@ from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.proxy_server import app, initialize +from litellm.proxy.proxy_server import ( + _enrich_model_info_with_litellm_data, + _get_proxy_model_info, + app, + initialize, +) from litellm.utils import _invalidate_model_cost_lowercase_map example_embedding_result = { @@ -13011,3 +13016,58 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp assert response.tokenizer_type == "huggingface_tokenizer" assert response.total_tokens > 0 assert_loop_stayed_free(took, lags) + + +# ModelInfo.set_model_info writes None over input_cost_per_token, +# output_cost_per_token, max_tokens, mode and base_model when the caller does not +# supply them, so an absent key becomes a key that is present and None. Cost-map +# enrichment must treat that as unset, otherwise a model priced in the built-in +# cost map reads back as unpriced and pre-call budget reservation stops bounding +# concurrent spend. An explicitly configured 0.0 is a deliberate free model. +_PRICED_MODEL = "gpt-4o" +_PRICED_BACKEND = f"openai/{_PRICED_MODEL}" + + +def _priced_model_deployment(model_info: dict) -> dict: + return { + "model_name": _PRICED_MODEL, + "litellm_params": {"model": _PRICED_BACKEND}, + "model_info": dict(model_info), + } + + +def _builtin_input_cost() -> float: + cost = litellm.get_model_info(model=_PRICED_BACKEND).get("input_cost_per_token") + if not cost: + pytest.skip(f"{_PRICED_BACKEND} carries no input_cost_per_token in the cost map") + return cost + + +@pytest.mark.parametrize( + "enrich", + [_get_proxy_model_info, _enrich_model_info_with_litellm_data], + ids=["_get_proxy_model_info", "_enrich_model_info_with_litellm_data"], +) +@pytest.mark.parametrize( + "model_info", + [{"id": "d1", "input_cost_per_token": None}, {"id": "d1"}], + ids=["present_as_none", "absent"], +) +def test_model_info_unset_cost_is_filled_from_cost_map(enrich, model_info): + """An unset cost key, whether absent or present as None, takes the cost-map value.""" + enriched = enrich(model=_priced_model_deployment(model_info)) + + assert enriched["model_info"]["input_cost_per_token"] == _builtin_input_cost() + + +@pytest.mark.parametrize( + "enrich", + [_get_proxy_model_info, _enrich_model_info_with_litellm_data], + ids=["_get_proxy_model_info", "_enrich_model_info_with_litellm_data"], +) +@pytest.mark.parametrize("configured", [0.0, 9.9e-07], ids=["free", "custom_rate"]) +def test_model_info_configured_cost_is_not_overwritten(enrich, configured): + """A cost the operator set is kept, including a deliberate 0.0.""" + enriched = enrich(model=_priced_model_deployment({"id": "d1", "input_cost_per_token": configured})) + + assert enriched["model_info"]["input_cost_per_token"] == configured