diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6f65100a162..a464cbe493f 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1032,8 +1032,16 @@ def _cost_map_model_name(model: str, custom_llm_provider: str | None, region_nam return _strip_unregistered_leading_segments(f"{provider_prefix}/{model}", region_name) +def _cost_map_row_key(model: str, custom_llm_provider: str | None) -> str | None: + try: + return cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)["key"] + except Exception: + return None + + def _row_declares_rates(model: str, custom_llm_provider: str | None) -> bool: - return _declares_rates(_cost_map_model_info(model, custom_llm_provider)) + key: Final = _cost_map_row_key(model, custom_llm_provider) + return key is not None and cost_map_key_prices_anything(key) def _rule_declares_rates(model: str, custom_llm_provider: str | None) -> bool: @@ -3491,7 +3499,7 @@ def _cost_map_entry_declares_pricing(model_name: str, custom_llm_provider: str) lookups cannot strip that prefix, so a zero-rated override read as declaring nothing, and a session that should bill nothing fell through to the public rates instead. """ - resolved: Final = _get_model_info_or_none(model_name, custom_llm_provider) + resolved: Final = _cost_map_model_info(model_name, custom_llm_provider) entries: Final = ( litellm.model_cost.get(resolved.get("key")) if resolved is not None else None, litellm.model_cost.get(model_name), @@ -3687,9 +3695,9 @@ def handle_realtime_transcription_cost_calculation( return 0.0 model_name: Final = _get_transcription_model_name_from_results(results) or litellm_model_name - model_info: Final = _get_model_info_or_none(model_name, custom_llm_provider) + model_info: Final = _cost_map_model_info(model_name, custom_llm_provider) override_info: Final = ( - _get_model_info_or_none(custom_pricing_model, custom_llm_provider) if custom_pricing_model is not None else None + _cost_map_model_info(custom_pricing_model, custom_llm_provider) if custom_pricing_model is not None else None ) total_cost = 0.0 @@ -3720,13 +3728,6 @@ def _get_transcription_model_name_from_results( return None -def _get_model_info_or_none(model: str, custom_llm_provider: str) -> ModelInfo | None: - try: - return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - return None - - _TranscriptionRateKey: TypeAlias = Literal[ "input_cost_per_second", "input_cost_per_audio_token", "input_cost_per_token", "output_cost_per_token" ] diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 6fbc2608120..bb6cefb4345 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -581,6 +581,30 @@ def test_completion_cost_bills_requested_model_from_hidden_params_when_no_model_ assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[requested_model], usage)) +def test_completion_cost_bills_requested_model_regional_row_when_response_names_an_unmapped_display_name( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "bedrock/sol-regional-test", + {"litellm_provider": "bedrock", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + monkeypatch.setitem( + litellm.model_cost, + "bedrock/eu-central-1/sol-regional-test", + {"litellm_provider": "bedrock", "mode": "chat", "input_cost_per_token": 5e-6, "output_cost_per_token": 11e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + response: Final = ModelResponse(model="Sol Regional Display", choices=[], usage=usage) + response._hidden_params.update( + {"litellm_model_name": "bedrock/sol-regional-test", "region_name": "eu-central-1", "custom_llm_provider": "bedrock"} + ) + + cost: Final = completion_cost(completion_response=response) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["bedrock/eu-central-1/sol-regional-test"], usage)) + + def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests( _local_model_cost_map: None, ) -> None: