scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

fix(pricing): retain registered rates across metadata updates

Browse source
This commit is contained in:
Joshua Valluru 2026-10-19 15:49:10 -07:00
parent cb50b02a2c
commit 512157f63c
3 changed files with 29 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/utils.py
View file

@ -3740,7 +3740,10 @@ def register_model(
if persist_across_reloads:
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost
for _registered_key, _registered_value in _registrations.items():
_runtime_registered_model_cost[_registered_key] = dict(_registered_value)
_runtime_registered_model_cost[_registered_key] = update_dictionary(
dict(_runtime_registered_model_cost.get(_registered_key, {})),
_RATE_TABLE.validate_python(_registered_value),
)
for key, value in loaded_model_cost.items():
write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name)

10
tests/unit/test_cost_calculator.py
View file

@ -9856,6 +9856,7 @@ def test_bedrock_converse_profile_preserves_registered_price_override(
}
}
)
litellm.register_model({model: {"supports_function_calling": True, "input_cost_per_token": None}})
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
assert completion_cost(
@ -9865,6 +9866,15 @@ def test_bedrock_converse_profile_preserves_registered_price_override(
optional_params={"model_id": f"us.{model}"},
) == pytest.approx(15 * rate)
updated_rate: Final = rate + 0.002
litellm.register_model({model: {"input_cost_per_token": str(updated_rate)}})
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
) == pytest.approx(10 * updated_rate + 5 * rate)
def test_bedrock_converse_profile_ignores_registered_capabilities_without_prices(
_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch

18
tests/unit/test_utils.py
View file

@ -4806,7 +4806,11 @@ async def test_builtin_string_callback_registers_when_subclass_already_active(
@pytest.mark.usefixtures("isolated_openai_model_sets")
@pytest.mark.usefixtures("isolated_openai_model_sets")
def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypatch):
@pytest.mark.parametrize("rate", (0.0, 0.000123))
@pytest.mark.parametrize("updated_rate", (None, 0.0, 0.000456))
def test_reapply_runtime_registrations_replays_register_model_overrides(
monkeypatch: pytest.MonkeyPatch, rate: float, updated_rate: float | None
) -> None:
"""
register_model is the documented way to override pricing for a model. A
price-data reload swaps litellm.model_cost for a freshly fetched catalog,
@ -4837,11 +4841,17 @@ def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypa
"openai/gpt-4o": {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 0.000123,
"input_cost_per_token": rate,
}
}
)
litellm.register_model({"openai/gpt-4o": {"supports_function_calling": True, "input_cost_per_token": None}})
litellm.register_model(
{"openai/gpt-4o": {"input_cost_per_token": str(updated_rate) if updated_rate is not None else None}}
)
litellm.register_model({"openai/gpt-4o": {"supports_function_calling": False}})
litellm.model_cost = {
"openai/gpt-4o": {
"litellm_provider": "openai",
@ -4853,7 +4863,9 @@ def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypa
_invalidate_model_cost_lowercase_map()
reapply_runtime_model_cost_registrations()
assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == 0.000123
expected_rate: Final = rate if updated_rate is None else updated_rate
assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == expected_rate
assert litellm.model_cost["openai/gpt-4o"]["supports_function_calling"] is False
assert litellm.model_cost["openai/gpt-4o"]["max_input_tokens"] == 4242
finally:
litellm.model_cost = saved_model_cost