fix(cost): apply the priced fallback to the provider response model streaming records

This commit is contained in:
mateo-berri 2026-10-28 07:29:16 -07:00
parent 439d4b8c5c
commit aad96a42d3
2 changed files with 34 additions and 9 deletions

View file

@ -996,19 +996,18 @@ def select_model_name_for_cost_calc(
else:
return_model = 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 base_model is not None or bedrock_pricing_model is not None:
return_model = base_model if base_model is not None else bedrock_pricing_model
elif provider_response_model is not None:
return _served_or_priced_requested_model(provider_response_model, model, custom_llm_provider, priced_region)
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:
return_model = hidden_params.get("model", model)
if return_model is None and completion_response_model is not None:
return _served_or_priced_requested_model(
served=_cost_map_model_name(completion_response_model, custom_llm_provider, priced_region),
requested=_cost_map_model_name(model, custom_llm_provider, None) if model is not None else None,
custom_llm_provider=custom_llm_provider,
)
return _served_or_priced_requested_model(completion_response_model, model, custom_llm_provider, priced_region)
if return_model is None and model is not None:
return_model = model
@ -1035,8 +1034,14 @@ def _resolves_to_rates(model: str, custom_llm_provider: str | None) -> bool:
return rule_info is not None and _cost_map_entry_prices_anything(rule_info)
def _served_or_priced_requested_model(served: str, requested: str | None, custom_llm_provider: str | None) -> str:
if requested is None or requested == served or _resolves_to_rates(served, custom_llm_provider):
def _served_or_priced_requested_model(
served_model: str, requested_model: str | None, custom_llm_provider: str | None, priced_region: str | None
) -> str:
served: Final = _cost_map_model_name(served_model, custom_llm_provider, priced_region)
if requested_model is None:
return served
requested: Final = _cost_map_model_name(requested_model, custom_llm_provider, None)
if requested == served or _resolves_to_rates(served, custom_llm_provider):
return served
return requested if _resolves_to_rates(requested, custom_llm_provider) else served

View file

@ -456,6 +456,26 @@ def test_completion_cost_bills_requested_model_when_response_names_an_unmapped_d
)
def test_completion_cost_bills_requested_model_when_stream_reports_an_unmapped_display_name(
_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
) -> None:
requested_model: Final = "fireworks_ai/accounts/fireworks/models/sol-display-test"
monkeypatch.setitem(
litellm.model_cost,
requested_model,
{"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6},
)
usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
response: Final = ModelResponse(model="accounts/fireworks/models/sol-display-test", choices=[], usage=usage)
response._hidden_params["provider_response_model"] = "Sol Display Max"
cost: Final = completion_cost(
completion_response=response, model=requested_model, custom_llm_provider="fireworks_ai"
)
assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[requested_model], usage))
def test_completion_cost_bills_requested_model_when_response_row_has_no_rates(
_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
) -> None: