diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 695a9efb990..3adc1c25dfd 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -739,11 +739,11 @@ def _get_provider_for_cost_calc( return custom_llm_provider -def _get_provider_response_model_for_cost_calc(hidden_params: object) -> str | None: +def _get_hidden_str_for_cost_calc(hidden_params: object, key: str) -> str | None: if not isinstance(hidden_params, Mapping): return None - model: Final[object] = hidden_params.get("provider_response_model") - return model if isinstance(model, str) and model else None + value: Final[object] = hidden_params.get(key) + return value if isinstance(value, str) and value else None def _select_model_name_for_cost_calc( @@ -771,9 +771,14 @@ def _select_model_name_for_cost_calc( elif isinstance(completion_response, dict): completion_response_model = completion_response.get("model", None) hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None) - provider_response_model: Final = _get_provider_response_model_for_cost_calc(hidden_params) - region_name_value: Final[object] = hidden_params.get("region_name") if hidden_params is not None else None - region_name: str | None = region_name_value if isinstance(region_name_value, str) else None + provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model") + explicit_pricing: Final = custom_pricing is True or base_model is not None + priced_from_response: Final = provider_response_model is not None or completion_response_model is not None + region_name: Final = ( + _get_hidden_str_for_cost_calc(hidden_params, "region_name") + if not explicit_pricing and priced_from_response + else None + ) if custom_pricing is True: if router_model_id is not None and router_model_id in litellm.model_cost: diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index bee1b924397..7c2174018e8 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4095,10 +4095,7 @@ def test_select_model_name_strips_duplicated_region_segment(_local_model_cost_ma ], model="us-east-1/anthropic.claude-v2:1", ) - response._hidden_params = { - "provider_response_model": "anthropic.claude-v2:1", - "region_name": "us-east-1", - } + response._hidden_params = {"region_name": "us-east-1"} selected = _select_model_name_for_cost_calc( model=None, @@ -4109,6 +4106,53 @@ def test_select_model_name_strips_duplicated_region_segment(_local_model_cost_ma assert selected == "bedrock/us-east-1/anthropic.claude-v2:1" +def _bedrock_response_with_private_model(model: str, region_name: str) -> litellm.ModelResponse: + response = litellm.ModelResponse( + id="x", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + model=model, + ) + response._hidden_params = {"provider_response_model": model, "region_name": region_name} + return response + + +def test_select_model_name_applies_region_to_private_provider_response_model(_local_model_cost_map): + """A Bedrock stream carries its requested model as the private provider model and must keep the + request's region in the cost key, exactly as the same request does without streaming.""" + + from litellm.cost_calculator import _select_model_name_for_cost_calc + + selected = _select_model_name_for_cost_calc( + model=None, + completion_response=_bedrock_response_with_private_model("anthropic.claude-v2:1", "us-east-1"), + custom_llm_provider="bedrock", + ) + + assert selected == "bedrock/us-east-1/anthropic.claude-v2:1" + + +def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map): + """An explicit base_model keeps pricing on that model's own key even when the request carries a + region with different regional rates, so the private provider model never widens region pricing.""" + + from litellm.cost_calculator import _select_model_name_for_cost_calc + + selected = _select_model_name_for_cost_calc( + model="my-bedrock-deployment", + completion_response=_bedrock_response_with_private_model("moonshotai.kimi-k2.5", "ap-northeast-1"), + base_model="moonshotai.kimi-k2.5", + custom_llm_provider="bedrock", + ) + + assert selected == "bedrock/moonshotai.kimi-k2.5" + + def test_completion_cost_nonzero_for_slash_alias_model_name(_local_model_cost_map): """End-to-end cost through a "/"-containing alias must price above zero (#38069)."""