diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4b82887e028..c24dffb7beb 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -150,7 +150,7 @@ from litellm.utils import ( TextCompletionResponse, TranscriptionResponse, cached_get_model_info_helper, - runtime_model_cost_registration, + runtime_model_cost_registrations, token_counter, unmapped_model_generalization, ) @@ -857,12 +857,14 @@ _NON_TOKEN_RATE_FIELDS: Final = frozenset( def _cost_map_entry_prices_anything(entry: Mapping[str, object]) -> bool: return any( - value is not None - and ( - field in _NON_TOKEN_RATE_FIELDS - or (field.startswith("cache_") and is_server_derived_pricing_key(field)) - or ("cost_per" in field and "token" in field) - ) + value is not None and (field in _NON_TOKEN_RATE_FIELDS or ("cost_per" in field and "token" in field)) + for field, value in entry.items() + ) + + +def _bedrock_registration_has_prices(entry: Mapping[str, object]) -> bool: + return _cost_map_entry_prices_anything(entry) or any( + value is not None and field.startswith("cache_") and is_server_derived_pricing_key(field) for field, value in entry.items() ) @@ -941,16 +943,15 @@ def select_model_name_for_cost_calc( completion_response_model if isinstance(completion_response_model, str) else None ) registered_pricing: Final = ( - runtime_model_cost_registration(configured_model, "bedrock") + runtime_model_cost_registrations(configured_model, "bedrock") if invoked_model_info is not None and isinstance(configured_model, str) - else None + else () ) - invoked_pricing_model: Final = ( - invoked_model_info["key"] - if invoked_model_info is not None - and _declares_rates(invoked_model_info) - and not (registered_pricing is not None and _cost_map_entry_prices_anything(registered_pricing)) - else None + registered_pricing_model: Final = next( + (key for key, entry in registered_pricing if _bedrock_registration_has_prices(entry)), None + ) + bedrock_pricing_model: Final = registered_pricing_model or ( + invoked_model_info["key"] if invoked_model_info is not None and _declares_rates(invoked_model_info) else None ) priced_from_response: Final = provider_response_model is not None or completion_response_model is not None priced_region: Final = ( @@ -969,8 +970,8 @@ def select_model_name_for_cost_calc( else: return_model = model - elif base_model is not None or invoked_pricing_model is not None or provider_response_model is not None: - return_model = base_model if base_model is not None else invoked_pricing_model or provider_response_model + elif base_model is not None or bedrock_pricing_model is not None or provider_response_model is not None: + return_model = base_model if base_model is not None else bedrock_pricing_model or provider_response_model elif completion_response_model is None and hidden_params is not None: if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0: diff --git a/litellm/utils.py b/litellm/utils.py index 84d91c43c59..b61b1825888 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3619,7 +3619,9 @@ def _register_on_current_map( _runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload -def runtime_model_cost_registration(model: str, custom_llm_provider: str) -> Mapping[str, object] | None: +def runtime_model_cost_registrations( + model: str, custom_llm_provider: str +) -> tuple[tuple[str, Mapping[str, object]], ...]: names: Final = get_potential_model_names(model, custom_llm_provider) candidates: Final = ( model, @@ -3628,14 +3630,13 @@ 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( - (entry for entry in entries if isinstance(entry, AliasedCostEntry) and entry.registered_fields), None + entries: Final = ((key, _builtin_cost_map().get(key)) for key in dict.fromkeys(candidates)) + registered: Final = ( + (key, entry) for key, entry in entries if isinstance(entry, AliasedCostEntry) and entry.registered_fields ) - return ( - MappingProxyType({key: value for key, value in registered.items() if key in registered.registered_fields}) - if registered is not None - else None + return tuple( + (key, MappingProxyType({name: value for name, value in entry.items() if name in entry.registered_fields})) + for key, entry in registered ) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index e4cd1944688..38d09b8642d 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -8845,6 +8845,37 @@ async def test_custom_pricing_matches_response_usage(async_mode: bool) -> None: assert response._hidden_params["response_cost"] == pytest.approx(_CUSTOM_PRICING_EXPECTED_COST) +@pytest.mark.parametrize("rate", (0.0, 0.001)) +@pytest.mark.parametrize("cache_field", ("cache_read_input_token_cost", "cache_creation_input_token_cost")) +@pytest.mark.usefixtures("isolated_pricing_registry") +@respx.mock +def test_router_cache_only_pricing_preserves_uncached_token_charges(rate: float, cache_field: str) -> None: + litellm.register_model( + {_CUSTOM_PRICING_MODEL: {**_CUSTOM_PRICING_RATES, "litellm_provider": "openai", "mode": "chat"}}, + persist_across_reloads=False, + ) + respx.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_custom_pricing_upstream) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "cache-priced-deployment", + "litellm_params": { + "model": f"openai/{_CUSTOM_PRICING_MODEL}", + "api_key": "sk-test", + cache_field: rate, + }, + "model_info": {"id": "cache-only-deployment"}, + } + ] + ) + + response: Final = router.completion( + model="cache-priced-deployment", messages=[{"role": "user", "content": "How are you?"}] + ) + + assert response._hidden_params["response_cost"] == pytest.approx(_CUSTOM_PRICING_EXPECTED_COST) + + @pytest.mark.usefixtures("isolated_pricing_registry") @respx.mock def test_streamed_custom_pricing_cost_matches_non_streamed_cost() -> None: @@ -9840,14 +9871,23 @@ 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)) +@pytest.mark.parametrize("routed_response", (True, False)) def test_bedrock_converse_profile_preserves_registered_price_override( - _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, rate: float, persist_across_reloads: bool + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + rate: float, + persist_across_reloads: bool, + routed_response: 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( + {f"bedrock/converse/{model}": {"supports_function_calling": True}}, + persist_across_reloads=persist_across_reloads, + ) litellm.register_model( { model: { @@ -9863,7 +9903,10 @@ def test_bedrock_converse_profile_preserves_registered_price_override( {model: {"supports_function_calling": True, "input_cost_per_token": None}}, persist_across_reloads=persist_across_reloads, ) - response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)) + response: Final = ModelResponse( + model=f"bedrock/converse/{model}" if routed_response else model, + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) assert completion_cost( completion_response=response, diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 3bdee29b23c..1cba6b02499 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10075,7 +10075,7 @@ def test_vertex_llama3_model_info_resolves_without_the_meta_prefix(): @pytest.mark.parametrize( "key", ("openai.test-converse", "bedrock/openai.test-converse", "bedrock/converse/openai.test-converse") ) -def test_runtime_model_cost_registration_resolves_bedrock_routing(monkeypatch: pytest.MonkeyPatch, key: str) -> None: +def test_runtime_model_cost_registrations_resolve_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 @@ -10085,9 +10085,12 @@ def test_runtime_model_cost_registration_resolves_bedrock_routing(monkeypatch: p 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 + registrations: Final = dict( + model_utils.runtime_model_cost_registrations("bedrock/converse/openai.test-converse", "bedrock") + ) + assert registrations[key] == rates + assert dict(model_utils.runtime_model_cost_registrations(alias, "bedrock"))[alias] == rates + assert model_utils.runtime_model_cost_registrations("bedrock/converse/openai.other-converse", "bedrock") == () def test_get_optional_params_keeps_a_tool_config_cache_point_out_of_the_openai_extra_body():