fix(pricing): honor temporary registered prices for Bedrock profiles

litellm/litellm_core_utils/get_model_cost_map.py

@@ -558,11 +558,14 @@ def get_model_cost_map_loaded_at() -> "datetime | None":
 
 
 class AliasedCostEntry(dict[str, object]):
-    __slots__ = ("names",)
+    __slots__ = ("names", "registered_fields")
 
-    def __init__(self, entry: Mapping[str, object], names: tuple[str, ...]) -> None:
+    def __init__(
+        self, entry: Mapping[str, object], names: tuple[str, ...], registered_fields: frozenset[str] = frozenset()
+    ) -> None:
         super().__init__(entry)
         self.names: tuple[str, ...] = names
+        self.registered_fields: frozenset[str] = registered_fields
 
 
 def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, ...]]:

litellm/utils.py

@@ -3305,14 +3305,20 @@ def _builtin_cost_map() -> dict[str, Mapping[str, object]]:
     return cast(dict[str, Mapping[str, object]], litellm.model_cost)
 
 
-def _replace_cost_map_entry(cost_map: dict[str, Mapping[str, object]], key: str, entry: Mapping[str, object]) -> None:
+def _replace_cost_map_entry(
+    cost_map: dict[str, Mapping[str, object]], key: str, entry: Mapping[str, object], registration: Mapping[str, object]
+) -> None:
     current: Final = cost_map.get(key)
     merged: Final = {**(current or {}), **entry}
-    if not isinstance(current, AliasedCostEntry):
-        cost_map.update({key: merged})
-        return
-    sharing_names: Final = tuple(name for name in current.names if cost_map.get(name) is current)
-    cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, current.names)))
+    names: Final = current.names if isinstance(current, AliasedCostEntry) else ()
+    previous_fields: Final[frozenset[str]] = (
+        current.registered_fields if isinstance(current, AliasedCostEntry) else frozenset()
+    )
+    registered_fields: Final = previous_fields | frozenset(
+        name for name, value in registration.items() if value is not None
+    )
+    sharing_names: Final = tuple(name for name in names if cost_map.get(name) is current)
+    cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, names, registered_fields)))
 
 
 def _without_synthesized_costs(
@@ -3590,7 +3596,7 @@ def _registration_attempt(
 ) -> _RegistrationWrite | None:
     source: Final = _builtin_cost_map()
     write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name)
-    _replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry)
+    _replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry, value)
     return write if final_attempt or litellm.model_cost is source else None
 
 
@@ -3622,10 +3628,15 @@ def runtime_model_cost_registration(model: str, custom_llm_provider: str) -> Map
         names["split_model"],
         names["stripped_model_name"],
     )
+    entries: Final = (_builtin_cost_map().get(key) for key in candidates)
     registered: Final = next(
-        (_runtime_registered_model_cost[key] for key in candidates if key in _runtime_registered_model_cost), None
+        (entry for entry in entries if isinstance(entry, AliasedCostEntry) and entry.registered_fields), None
+    )
+    return (
+        MappingProxyType({key: value for key, value in registered.items() if key in registered.registered_fields})
+        if registered is not None
+        else None
     )
-    return MappingProxyType(registered) if registered is not None else None
 
 
 class _LiveDeploymentReplay:

tests/unit/test_cost_calculator.py

@@ -9839,12 +9839,14 @@ def test_non_bedrock_cost_ignores_bedrock_model_id(_local_model_cost_map: None)
 
 
 @pytest.mark.parametrize("rate", (0.0, 0.001))
+@pytest.mark.parametrize("persist_across_reloads", (True, False))
 def test_bedrock_converse_profile_preserves_registered_price_override(
-    _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, rate: float
+    _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, rate: float, persist_across_reloads: bool
 ) -> None:
     import litellm.utils as model_utils
 
     monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {})
+    catalog: Final = dict(litellm.model_cost)
     model: Final = "openai.gpt-5.6-sol"
     litellm.register_model(
         {
@@ -9854,9 +9856,13 @@ def test_bedrock_converse_profile_preserves_registered_price_override(
                 "litellm_provider": "bedrock_converse",
                 "mode": "chat",
             }
-        }
+        },
+        persist_across_reloads=persist_across_reloads,
+    )
+    litellm.register_model(
+        {model: {"supports_function_calling": True, "input_cost_per_token": None}},
+        persist_across_reloads=persist_across_reloads,
     )
-    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(
@@ -9867,7 +9873,9 @@ def test_bedrock_converse_profile_preserves_registered_price_override(
     ) == pytest.approx(15 * rate)
 
     updated_rate: Final = rate + 0.002
-    litellm.register_model({model: {"input_cost_per_token": str(updated_rate)}})
+    litellm.register_model(
+        {model: {"input_cost_per_token": str(updated_rate)}}, persist_across_reloads=persist_across_reloads
+    )
     assert completion_cost(
         completion_response=response,
         model=f"bedrock/converse/{model}",
@@ -9875,6 +9883,22 @@ def test_bedrock_converse_profile_preserves_registered_price_override(
         optional_params={"model_id": f"us.{model}"},
     ) == pytest.approx(10 * updated_rate + 5 * rate)
 
+    monkeypatch.setattr(litellm, "model_cost", catalog)
+    model_utils._invalidate_model_cost_lowercase_map()
+    model_utils.reapply_runtime_model_cost_registrations()
+    profile_prices: Final = litellm.get_model_info(f"us.{model}", custom_llm_provider="bedrock")
+    expected_after_reload: Final = (
+        10 * updated_rate + 5 * rate
+        if persist_across_reloads
+        else 10 * profile_prices["input_cost_per_token"] + 5 * profile_prices["output_cost_per_token"]
+    )
+    assert completion_cost(
+        completion_response=response,
+        model=f"bedrock/converse/{model}",
+        custom_llm_provider="bedrock",
+        optional_params={"model_id": f"us.{model}"},
+    ) == pytest.approx(expected_after_reload)
+
 
 def test_bedrock_converse_profile_ignores_registered_capabilities_without_prices(
     _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch

tests/unit/test_utils.py

@@ -10113,9 +10113,14 @@ def test_vertex_llama3_model_info_resolves_without_the_meta_prefix():
 )
 def test_runtime_model_cost_registration_resolves_bedrock_routing(monkeypatch: pytest.MonkeyPatch, key: str) -> None:
     from litellm import utils as model_utils
+    from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
 
     rates: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
-    monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {key: rates})
+    alias: Final = "openai.test-converse-alias"
+    entry: Final = AliasedCostEntry({"litellm_provider": "bedrock_converse", "mode": "chat"}, (key, alias))
+    monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, key: entry, alias: entry})
+    litellm.register_model({key: rates}, persist_across_reloads=False)
 
     assert model_utils.runtime_model_cost_registration("bedrock/converse/openai.test-converse", "bedrock") == rates
+    assert model_utils.runtime_model_cost_registration(alias, "bedrock") == rates
     assert model_utils.runtime_model_cost_registration("bedrock/converse/openai.other-converse", "bedrock") is None