scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-28 06:23:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 94

fix(bedrock): resolve profile pricing before applying custom rates (#47419)

Browse source 
Backport of #47419 to stable/1.104.x.
Cherry-picked from 3124797399 (main).
Ported onto the line's register_model and router startup path without #41523's aliased-entry rewrite: registered fields come from the side table added in the previous commit. Tests adapted to the line: private names (_replay_model_cost_registrations, _response_cost_calculator), the custom-pricing test helpers from main's test_cost_calculator.py, CONVERSE_ANSWER inlined, and the alias-propagation and bedrock gpt-6.1-sol/ARN cases removed since those depend on #41523 and #46128, which are not on this line.
This commit is contained in:
joshua-berri 2026-10-20 16:39:37 -07:00 • committed by Mrinal Chanshetty
parent 7849ed2330
commit f5049dba86
11 changed files with 1220 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

99
litellm/cost_calculator.py
View file

@ -8,7 +8,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from httpx import Response
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
@ -125,6 +125,7 @@ from litellm.types.utils import (
TranscriptionUsageTokensObject,
Usage,
VectorStoreSearchResponse,
is_server_derived_pricing_key,
)
from litellm.utils import (
CallTypes,
@ -137,6 +138,8 @@ from litellm.utils import (
TextCompletionResponse,
TranscriptionResponse,
_cached_get_model_info_helper,
model_info_with_pricing_overrides,
runtime_model_cost_registrations,
token_counter,
)
@ -402,6 +405,7 @@ def cost_per_token(
### REQUEST MODEL ###
request_model: str | None = None, # original request model for router detection
custom_model_info: OCRPricing | None = None,
bedrock_model_info: ModelInfo | None = None,
) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@ -688,6 +692,15 @@ def cost_per_token(
elif custom_llm_provider == "anthropic":
return anthropic_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
elif custom_llm_provider == "bedrock":
resolved_bedrock_info: Final = bedrock_model_info or _bedrock_pricing_info(model, None, "bedrock", False)
if resolved_bedrock_info is not None:
return generic_cost_per_token(
model=model,
usage=usage_block,
custom_llm_provider="bedrock",
service_tier=service_tier,
model_info=resolved_bedrock_info,
)
return bedrock_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
elif custom_llm_provider == "openai":
return openai_cost_per_token(
@ -801,6 +814,55 @@ def _cost_map_entry_prices_anything(entry: Mapping[str, object]) -> bool:
)
_LOGGED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
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()
)
def _bedrock_pricing_info(
model: str | None,
invoked_model_id: str | None,
custom_llm_provider: str | None,
explicit_pricing: bool,
) -> ModelInfo | None:
if explicit_pricing or custom_llm_provider != "bedrock" or not model:
return None
registration: Final = next(
(
(key, entry)
for key, entry in runtime_model_cost_registrations(model, "bedrock")
if _bedrock_registration_has_prices(entry)
),
None,
)
profile: Final = _cost_map_model_info(invoked_model_id, "bedrock") if invoked_model_id else None
if profile is None or not _declares_rates(profile):
return _cost_map_model_info(registration[0], None) if registration is not None else None
if registration is None or registration[0] == profile["key"]:
return profile
return model_info_with_pricing_overrides(profile, registration[1])
def _declares_rates(model_info: ModelInfo | None) -> bool:
entry: Final = _raw_cost_map_entry(model_info["key"]) if model_info is not None else None
return entry is not None and _cost_map_entry_prices_anything(entry)
def _invoked_model_id_for_cost(
optional_params: object,
logging_obj: LitellmLoggingObject | None,
) -> str | None:
recorded: Final[object] = getattr(logging_obj, "model_call_details", None)
return _get_hidden_str_for_cost_calc(optional_params, "model_id") or _get_hidden_str_for_cost_calc(
recorded, "model_id"
)
def _select_model_name_for_cost_calc(
model: str | None,
completion_response: object | None,
@ -809,6 +871,7 @@ def _select_model_name_for_cost_calc(
custom_llm_provider: str | None = None,
router_model_id: str | None = None,
region_name: str | None = None,
invoked_model_id: str | None = None,
) -> str | None:
"""
1. If custom pricing is true, return received model name
@ -829,6 +892,13 @@ def _select_model_name_for_cost_calc(
hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", 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
configured_model: Final = model or (
completion_response_model if isinstance(completion_response_model, str) else None
)
bedrock_model_info: Final = _bedrock_pricing_info(
configured_model, invoked_model_id, custom_llm_provider, explicit_pricing
)
bedrock_pricing_model: Final = bedrock_model_info["key"] if bedrock_model_info is not None else None
priced_from_response: Final = provider_response_model is not None or completion_response_model is not None
priced_region: Final = (
_get_hidden_str_for_cost_calc(hidden_params, "region_name") or region_name
@ -846,8 +916,8 @@ def _select_model_name_for_cost_calc(
else:
return_model = model
elif base_model is not None or provider_response_model is not None:
return_model = base_model if base_model is not None else 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:
@ -1426,6 +1496,10 @@ def completion_cost(
service_tier = _normalize_service_tier(service_tier)
explicit_pricing: Final = custom_pricing is True or base_model is not None
invoked_model_id: Final = _invoked_model_id_for_cost(
_LOGGED_LITELLM_PARAMS.validate_python(optional_params) if optional_params is not None else None,
litellm_logging_obj,
)
selected_model: Final = _select_model_name_for_cost_calc(
model=model,
completion_response=completion_response,
@ -1434,6 +1508,14 @@ def completion_cost(
base_model=base_model,
router_model_id=router_model_id,
region_name=region_name,
invoked_model_id=invoked_model_id,
)
bedrock_model_info: Final = _bedrock_pricing_info(
model or _get_response_model(completion_response),
invoked_model_id,
_get_provider_for_cost_calc(model, custom_llm_provider),
explicit_pricing,
)
potential_model_names: Final = [
@ -1808,6 +1890,7 @@ def completion_cost(
vertex_location=vertex_location,
response=completion_response,
custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id),
bedrock_model_info=bedrock_model_info,
)
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
@ -1896,6 +1979,7 @@ def completion_cost(
data_residency=data_residency,
vertex_location=vertex_location,
custom_cost_per_token=custom_cost_per_token,
model_info=bedrock_model_info,
)
_reasoning_cost = _token_type_breakdown.reasoning_cost
_cache_read_cost = _token_type_breakdown.cache_read_cost
@ -2137,7 +2221,16 @@ def pricing_entry_for_cost_calc(
custom_llm_provider=custom_llm_provider,
router_model_id=router_model_id,
region_name=region_name,
invoked_model_id=_invoked_model_id_for_cost(None, litellm_logging_obj),
)
bedrock_model_info: Final = _bedrock_pricing_info(
model or _get_response_model(completion_response),
_invoked_model_id_for_cost(None, litellm_logging_obj),
_get_provider_for_cost_calc(model, custom_llm_provider),
custom_pricing is True or base_model is not None,
)
if bedrock_model_info is not None:
return bedrock_model_info["key"], bedrock_model_info
candidates: Final = (selected_model, _get_response_model(completion_response), model)
resolved: Final = next(
(info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info),

25
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -275,7 +275,7 @@ def _parse_token_threshold(threshold: str) -> float:
def _parse_above_token_threshold(key: str) -> float:
return _parse_token_threshold(key.split("_above_")[1].split("_tokens")[0])
return _parse_token_threshold(key.rsplit("_above_", 1)[1].split("_tokens")[0])
def _prompt_exceeds_threshold(prompt_tokens: int, threshold: float, inclusive: bool) -> bool:
@ -669,10 +669,12 @@ def _get_token_base_cost(
# so that the threshold detection loop only processes standard keys. The
# service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key.
tier_key_suffix: Final = _get_service_tier_cost_key("", service_tier)
threshold_prefixes: Final = tuple(f"{prefix}_above_" for prefix in _BATCH_RATE_PREFIXES)
threshold_keys: Final = [
k
for k in model_info
if k.startswith("input_cost_per_token_above_")
if k.startswith(threshold_prefixes)
and "_tokens" in k
and (
not k.endswith(_NON_STANDARD_THRESHOLD_SUFFIXES) or (tier_key_suffix != "" and k.endswith(tier_key_suffix))
)
@ -680,12 +682,12 @@ def _get_token_base_cost(
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
threshold: float | None = None
for key in sorted(threshold_keys, key=_parse_above_token_threshold, reverse=True):
for key in sorted(threshold_keys, key=_parse_above_token_threshold):
value = model_info.get(key)
if value is not None:
try:
# Handle both formats: _above_128k_tokens and _above_128_tokens
threshold_str = key.split("_above_")[1].split("_tokens")[0]
threshold_str = key.rsplit("_above_", 1)[1].split("_tokens")[0]
threshold = _parse_above_token_threshold(key)
if _prompt_exceeds_threshold(usage.prompt_tokens, threshold, threshold_is_inclusive):
# Prefer a service_tier-specific above-threshold key when available,
@ -698,7 +700,7 @@ def _get_token_base_cost(
service_tier,
)
if service_tier
else key
else f"input_cost_per_token_above_{threshold_str}_tokens"
)
prompt_base_cost = cast(
float,
@ -754,8 +756,6 @@ def _get_token_base_cost(
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
break
except (IndexError, ValueError):
continue
except Exception:
@ -1636,20 +1636,23 @@ def get_billed_token_rates(
vertex_location: str | None = None,
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
model_info: ModelInfo | None = None,
) -> BilledTokenRates | None:
"""Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type
breakdown resolve them. None when the model's pricing cannot be resolved."""
if custom_cost_per_token is not None:
return _custom_pricing_rates(custom_cost_per_token)
try:
model_info: Final = apply_provider_cache_read_default(
get_model_info(model=model, custom_llm_provider=custom_llm_provider),
resolved_model_info: Final = apply_provider_cache_read_default(
model_info
if model_info is not None
else get_model_info(model=model, custom_llm_provider=custom_llm_provider),
custom_llm_provider,
)
except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates
return None
return _cost_map_billed_rates(
model_info=model_info,
model_info=resolved_model_info,
usage=usage,
custom_llm_provider=custom_llm_provider,
service_tier=service_tier,
@ -1668,6 +1671,7 @@ def get_token_type_cost_breakdown(
vertex_location: str | None = None,
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
model_info: ModelInfo | None = None,
) -> TokenTypeCostBreakdown:
"""
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
@ -1692,6 +1696,7 @@ def get_token_type_cost_breakdown(
vertex_location=vertex_location,
current_time=current_time,
custom_cost_per_token=custom_cost_per_token,
model_info=model_info,
)
if rates is None:
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)

1
litellm/main.py
View file

@ -5722,6 +5722,7 @@ def completion(
optional_params=processed_non_default_params, # [IMPORTANT] - using processed_non_default_params ensures consistent params logged to langfuse for finetuning / eval datasets.
litellm_params=litellm_params,
custom_llm_provider=custom_llm_provider,
model_id=optional_params.get("model_id"),
)
if mock_response or mock_tool_calls or mock_timeout:
kwargs.pop("mock_timeout", None) # remove for any fallbacks triggered

85
litellm/router.py
View file

@ -318,12 +318,14 @@ from litellm.types.utils import (
StandardLoggingRoutingDecision,
Usage,
all_litellm_params,
is_server_derived_pricing_key,
shared_backend_model_info,
)
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.utils import (
CustomStreamWrapper,
EmbeddingResponse,
ModelInfoLookupMiss,
ModelResponse,
Rules,
function_setup,
@ -332,6 +334,8 @@ from litellm.utils import (
get_secret,
get_utc_datetime,
is_region_allowed,
lookup_model_info_for_registration,
model_info_with_pricing_overrides,
provider_rejectable_params,
set_live_deployment_replay,
)
@ -471,6 +475,8 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({})
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
@ -8824,39 +8830,7 @@ class Router:
litellm_params=litellm_params,
model_info=_model_info,
)
for field in CustomPricingLiteLLMParams.model_fields:
if deployment.litellm_params.get(field) is not None:
_model_info[field] = deployment.litellm_params[field]
Router._inherit_builtin_base_rates_for_off_peak(
model_info=_model_info,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
if _model_info.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=_model_info,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_service_tier_pricing(
model_info=_model_info,
backend_model=Router._cost_map_backend_model(deployment),
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_tiered_output_rate(
model_info=_model_info,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
## REGISTER MODEL INFO IN LITELLM MODEL COST MAP
Router._register_deployment_in_model_cost(
model_id=deployment.model_info.id,
model_info=_model_info,
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._register_deployment_pricing(deployment)
## Check if LLM Deployment is allowed for this deployment
if self.deployment_is_active_for_environment(deployment=deployment) is not True:
@ -9892,10 +9866,24 @@ class Router:
params_base_model: Final = deployment.litellm_params.get("base_model")
if isinstance(params_base_model, str) and params_base_model:
return params_base_model
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True))
invoked_model: Final = params.get("model_id")
if (
isinstance(invoked_model, str)
and invoked_model
and (
deployment.litellm_params.custom_llm_provider == "bedrock"
or deployment.litellm_params.model.startswith("bedrock/")
)
):
profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}"
profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock")
if not isinstance(profile_info, ModelInfoLookupMiss):
return profile_model
return deployment.litellm_params.model
@staticmethod
def _deployment_model_cost_payload(deployment: Deployment) -> dict: # mutable-ok: cost-map entry
def _deployment_model_cost_payload(deployment: Deployment) -> dict[str, object]: # mutable-ok: cost-map entry
"""The ``model_info`` a deployment contributes to ``litellm.model_cost``.
Custom pricing lives on ``litellm_params`` rather than ``model_info``, and
@ -9903,30 +9891,49 @@ class Router:
both are folded back in here. That keeps this reproducible from a
deployment alone, which is what lets a refresh rebuild the same entries.
"""
model_info: Final[dict] = deployment.model_info.model_dump(exclude_none=True) # mutable-ok: built in place
model_info: Final = dict(
_MODEL_INFO_ADAPTER.validate_python(deployment.model_info.model_dump(exclude_none=True))
)
for field in CustomPricingLiteLLMParams.model_fields:
field_value = deployment.litellm_params.get(field)
if field_value is not None:
model_info[field] = field_value
backend_model: Final = Router._cost_map_backend_model(deployment)
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True))
if (
params.get("model_id")
and backend_model != deployment.litellm_params.model
and (
deployment.litellm_params.custom_llm_provider == "bedrock"
or deployment.litellm_params.model.startswith("bedrock/")
)
):
profile_info: Final = lookup_model_info_for_registration(model=backend_model, custom_llm_provider="bedrock")
if not isinstance(profile_info, ModelInfoLookupMiss):
explicit_info: Final = _MODEL_INFO_ADAPTER.validate_python(model_info)
priced_info: Final = model_info_with_pricing_overrides(profile_info, explicit_info)
model_info.update(
{key: value for key, value in priced_info.items() if is_server_derived_pricing_key(key)}
)
Router._inherit_builtin_base_rates_for_off_peak(
model_info=model_info,
backend_model=deployment.litellm_params.model,
backend_model=backend_model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
if model_info.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=model_info,
backend_model=deployment.litellm_params.model,
backend_model=backend_model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_service_tier_pricing(
model_info=model_info,
backend_model=Router._cost_map_backend_model(deployment),
backend_model=backend_model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_tiered_output_rate(
model_info=model_info,
backend_model=deployment.litellm_params.model,
backend_model=backend_model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
return model_info

5
litellm/utils.py
View file

@ -3468,7 +3468,10 @@ def register_model(
if persist_across_reloads:
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost
for _registered_key, _registered_value in _registrations.items():
_runtime_registered_model_cost[_registered_key] = dict(_registered_value) # mutable-ok: caller-owned
_runtime_registered_model_cost[_registered_key] = _update_dictionary(
dict(_runtime_registered_model_cost.get(_registered_key, {})),
dict(_RATE_TABLE.validate_python(_registered_value)),
)
_skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO

82
tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -3936,3 +3936,85 @@ def test_get_batch_cost_rates_claude_haiku_5_5_prompt_length_tiers(
)
assert (rates.input, rates.output, rates.cache_read, rates.cache_creation) == expected
@pytest.mark.parametrize(
"field", ("output_cost_per_token", "cache_read_input_token_cost", "cache_creation_input_token_cost")
)
def test_flat_input_price_keeps_other_token_thresholds(field: str) -> None:
from typing import Final, cast
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost
info: Final = cast(
ModelInfo,
{
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
"cache_read_input_token_cost": 0.003,
"cache_creation_input_token_cost": 0.004,
f"{field}_above_300k_tokens": 0.005,
},
)
rates: Final = _get_token_base_cost(
model_info=info,
usage=Usage(prompt_tokens=300_001, completion_tokens=3, total_tokens=300_004),
)
assert rates[0] == 0.001
assert rates[1] == (0.005 if field == "output_cost_per_token" else 0.002)
assert rates[2] == (0.005 if field == "cache_creation_input_token_cost" else 0.004)
assert rates[4] == (0.005 if field == "cache_read_input_token_cost" else 0.003)
@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001, 512_000, 512_001))
@pytest.mark.parametrize("higher_rate", (0.0, 0.009))
def test_mixed_token_thresholds_preserve_each_crossed_rate(prompt_tokens: int, higher_rate: float) -> None:
from typing import Final, cast
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost
info: Final = cast(
ModelInfo,
{
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
"cache_read_input_token_cost": 0.003,
"cache_creation_input_token_cost": 0.004,
"input_cost_per_token_above_200k_tokens": 0.005,
"output_cost_per_token_above_200k_tokens": 0.006,
"cache_creation_input_token_cost_above_200k_tokens": 0.007,
"cache_read_input_token_cost_above_512k_tokens": higher_rate,
"output_cost_per_token_above_512k_tokens": higher_rate,
},
)
rates: Final = _get_token_base_cost(
model_info=info,
usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=3, total_tokens=prompt_tokens + 3),
)
assert rates[0] == (0.005 if prompt_tokens > 200_000 else 0.001)
assert rates[1] == (higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002)
assert rates[2] == (0.007 if prompt_tokens > 200_000 else 0.004)
assert rates[4] == (higher_rate if prompt_tokens > 512_000 else 0.003)
@pytest.mark.parametrize("rate", (0.0, 0.001))
def test_token_breakdown_uses_supplied_pricing_without_catalog_lookup(rate: float) -> None:
model_info: Final[ModelInfo] = {
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"cache_read_input_token_cost": rate,
"cache_creation_input_token_cost": rate,
}
usage: Final = Usage(
prompt_tokens=20, completion_tokens=3, total_tokens=23,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), cache_creation_input_tokens=5,
)
prompt_cost, output_cost = generic_cost_per_token(
model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info,
)
breakdown: Final = get_token_type_cost_breakdown(
model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info,
)
assert breakdown.cache_read_cost == pytest.approx(5 * rate)
assert breakdown.cache_creation_cost == pytest.approx(5 * rate)
assert breakdown.rates is not None
assert prompt_cost == pytest.approx(10 * breakdown.rates.input_cost_per_token + breakdown.cache_read_cost + breakdown.cache_creation_cost)
assert output_cost == pytest.approx(3 * breakdown.rates.output_cost_per_token)

43
tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -16,6 +16,8 @@ import pytest
from botocore.credentials import Credentials
from botocore.exceptions import ClientError
import litellm
from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.rust_bridge import configuration
@ -273,3 +275,44 @@ def test_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch):
sent = client.post.call_args.kwargs
assert "Credential=ASIACONVERSETAGGED/" in sent["headers"]["Authorization"]
assert "aws_session_tags" not in sent["data"]
@pytest.mark.parametrize("stream", (False, True))
def test_converse_profile_override_records_cost_without_changing_response_model(
monkeypatch: pytest.MonkeyPatch, stream: bool
) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setenv("LITELLM_RUST", "0")
configuration.rust(False)
model: Final = "openai.gpt-5.6-sol"
profile: Final = f"us.{model}"
prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock")
def respond(request: httpx.Request) -> httpx.Response:
assert request.url.path == f"/model/{profile}/converse"
return httpx.Response(200, json=CONVERSE_RESPONSE)
with httpx.Client(transport=httpx.MockTransport(respond)) as transport:
response: Final = litellm.completion(
model=f"bedrock/converse/{model}",
model_id=profile,
messages=[{"role": "user", "content": "hi"}],
api_key="test-bearer",
aws_region_name="us-east-1",
client=HTTPHandler(client=transport),
stream=stream,
fake_stream=stream,
max_tokens=16,
)
result: Final = litellm.stream_chunk_builder(list(response)) if stream else response
assert result is not None
assert result.model == model
assert result.choices[0].message.content == "hi"
cost: Final = (
response.logging_obj._response_cost_calculator(result) if stream else result._hidden_params.get("response_cost")
)
assert cost == pytest.approx(
CONVERSE_RESPONSE["usage"]["inputTokens"] * prices["input_cost_per_token"]
+ CONVERSE_RESPONSE["usage"]["outputTokens"] * prices["output_cost_per_token"]
)

620
tests/unit/test_cost_calculator.py
View file

@ -1,10 +1,13 @@
import datetime
import json
import time
from pathlib import Path
from types import MappingProxyType, SimpleNamespace
from typing import Final, cast
import httpx
import pytest
import respx
from pydantic import BaseModel
import litellm
@ -5319,3 +5322,620 @@ def test_completion_cost_is_zero_when_explicit_rates_are_zero(monkeypatch: pytes
)
assert cost == 0.0
@pytest.fixture
def isolated_pricing_registry(local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "open_ai_chat_completion_models", set(litellm.open_ai_chat_completion_models))
monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider))
_CUSTOM_PRICING_MODEL: Final = "migration-custom-pricing-chat"
_CUSTOM_PRICING_RATES: Final = MappingProxyType(
{"input_cost_per_token": 0.0000008, "output_cost_per_token": 0.0000032}
)
_CUSTOM_PRICING_EXPECTED_COST: Final = 7 * 0.0000008 + 4 * 0.0000032
def _custom_pricing_chunk(choices: list[dict[str, object]], usage: dict[str, int] | None = None) -> str:
chunk: Final = {
"id": "chatcmpl-custom-pricing",
"object": "chat.completion.chunk",
"created": 1,
"model": _CUSTOM_PRICING_MODEL,
"choices": choices,
**({"usage": usage} if usage is not None else {}),
}
return f"data: {json.dumps(chunk)}\n\n"
def _custom_pricing_upstream(request: httpx.Request) -> httpx.Response:
if json.loads(request.content).get("stream") is True:
body: Final = "".join(
(
_custom_pricing_chunk(
[{"index": 0, "delta": {"role": "assistant", "content": "I am ready to help."}, "finish_reason": None}]
),
_custom_pricing_chunk([{"index": 0, "delta": {}, "finish_reason": "stop"}]),
_custom_pricing_chunk([], {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11}),
"data: [DONE]\n\n",
)
)
return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
return httpx.Response(
200,
json={
"id": "chatcmpl-custom-pricing",
"object": "chat.completion",
"created": 1,
"model": _CUSTOM_PRICING_MODEL,
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "I am ready to help."}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11},
},
)
@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.parametrize("profile_prefix", ("us", "global"))
@pytest.mark.parametrize(
"model_name",
("gpt-5.6-sol", "gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-sol", "gpt-6-luna", "gpt-6-astra"),
)
@pytest.mark.parametrize("arn", (False,))
def test_bedrock_converse_cost_uses_invoked_profile(
_local_model_cost_map: None, model_name: str, profile_prefix: str, arn: bool
) -> None:
model: Final = f"openai.{model_name}"
profile: Final = f"{profile_prefix}.{model}"
invoked: Final = f"arn:aws:bedrock:us-east-1:123456789012:inference-profile/{profile}" if arn else profile
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock")
cost: Final = completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": invoked},
)
assert cost == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"])
assert response.model == model
@pytest.mark.parametrize("custom_pricing", (False, True))
def test_bedrock_converse_profile_preserves_explicit_base_and_custom_pricing(
_local_model_cost_map: None, custom_pricing: bool
) -> None:
model: Final = "openai.gpt-6-sol"
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
prices: Final = litellm.get_model_info(model, custom_llm_provider="bedrock")
cost: Final = completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
custom_pricing=custom_pricing,
base_model=None if custom_pricing else model,
)
assert cost == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"])
@pytest.mark.parametrize("stream", (False, True))
def test_bedrock_converse_cost_retains_profile_after_dispatch_consumes_optional_params(
_local_model_cost_map: None, stream: bool
) -> None:
from litellm.cost_calculator import pricing_entry_for_cost_calc
model: Final = "openai.gpt-5.6-sol"
profile: Final = f"us.{model}"
logging_obj: Final = Logging(
model=model,
messages=[{"role": "user", "content": "hi"}],
stream=stream,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="bedrock-profile-cost",
function_id="bedrock-profile-cost",
)
logging_obj.update_environment_variables(
litellm_params={}, optional_params={"model_id": profile}, custom_llm_provider="bedrock"
)
logging_obj.optional_params.pop("model_id")
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock")
assert logging_obj._response_cost_calculator(response) == pytest.approx(
10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
)
entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj)
assert entry is not None
assert entry[0] == prices["key"]
@pytest.mark.parametrize(
"invoked_model_id",
(None, "", 123, "unmapped-profile", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque"),
)
def test_bedrock_converse_cost_keeps_configured_rates_when_target_has_no_prices(
_local_model_cost_map: None, invoked_model_id: object
) -> None:
model: Final = "openai.gpt-6-sol"
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
expected: Final = completion_cost(completion_response=response, model=model, custom_llm_provider="bedrock")
assert (
completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
optional_params={"model_id": invoked_model_id},
)
== expected
)
def test_non_bedrock_cost_ignores_bedrock_model_id(_local_model_cost_map: None) -> None:
response: Final = ModelResponse(
model="gpt-6-sol", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
)
expected: Final = completion_cost(completion_response=response, custom_llm_provider="openai")
assert (
completion_cost(
completion_response=response,
custom_llm_provider="openai",
optional_params={"model_id": "us.openai.gpt-6-sol"},
)
== expected
)
@pytest.mark.parametrize("rate", (0.0, 0.001))
@pytest.mark.parametrize("persist_across_reloads", (True, False))
@pytest.mark.parametrize("routed_response", (True, False))
@pytest.mark.parametrize("metadata_prefix", ("bedrock/", "bedrock/converse/"))
def test_bedrock_converse_profile_preserves_registered_price_override(
_local_model_cost_map: None,
monkeypatch: pytest.MonkeyPatch,
rate: float,
persist_across_reloads: bool,
routed_response: bool,
metadata_prefix: str,
) -> None:
import litellm.utils as model_utils
from litellm.cost_calculator import pricing_entry_for_cost_calc
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"{metadata_prefix}{model}": {"supports_function_calling": True}},
persist_across_reloads=persist_across_reloads,
)
litellm.register_model(
{
model: {
"input_cost_per_token": rate,
"output_cost_per_token": rate,
"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,
)
response: Final = ModelResponse(
model=f"bedrock/converse/{model}" if routed_response else model,
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="registered-pricing",
function_id="registered-pricing",
)
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
litellm_logging_obj=logging_obj,
) == pytest.approx(15 * rate)
assert logging_obj.billed_token_rates is not None
assert logging_obj.billed_token_rates.input_cost_per_token == rate
assert logging_obj.billed_token_rates.output_cost_per_token == rate
entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, None)
assert entry is not None
assert entry[0] == model
assert entry[1]["input_cost_per_token"] == rate
assert entry[1]["output_cost_per_token"] == rate
updated_rate: Final = rate + 0.002
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}",
custom_llm_provider="bedrock",
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
) -> None:
import litellm.utils as model_utils
monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {})
model: Final = "openai.gpt-6-sol"
litellm.register_model({model: {"supports_function_calling": True}})
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
prices: Final = litellm.get_model_info(f"us.{model}", custom_llm_provider="bedrock")
assert completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
) == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"])
@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.parametrize("has_base_rates", (False, True))
@pytest.mark.parametrize("persist_across_reloads", (False, True))
@pytest.mark.parametrize("profile_prefix", ("us", "global"))
@pytest.mark.parametrize("provider", (None, "bedrock"))
def test_bedrock_converse_profile_preserves_cache_only_registered_prices(
_local_model_cost_map: None,
monkeypatch: pytest.MonkeyPatch,
rate: float,
cache_field: str,
has_base_rates: bool,
persist_across_reloads: bool,
profile_prefix: str,
provider: str | None,
) -> None:
import litellm.utils as model_utils
from litellm.cost_calculator import pricing_entry_for_cost_calc
monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {})
model: Final = "openai.cache-overlay-test"
profile: Final = f"{profile_prefix}.{model}"
catalog: Final = {
**litellm.model_cost,
model: {
"litellm_provider": "bedrock",
"mode": "chat",
**({"input_cost_per_token": 0.01, "output_cost_per_token": 0.02} if has_base_rates else {}),
},
profile: {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.02,
"output_cost_per_token": 0.04,
"cache_read_input_token_cost": 0.005,
"cache_creation_input_token_cost": 0.025,
},
}
profile_before: Final = dict(catalog[profile])
monkeypatch.setattr(litellm, "model_cost", dict(catalog))
model_utils._invalidate_model_cost_lowercase_map()
litellm.register_model({model: {cache_field: rate}}, persist_across_reloads=persist_across_reloads)
response: Final = ModelResponse(
model=model,
usage=Usage(
prompt_tokens=15,
completion_tokens=3,
total_tokens=18,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=5 if cache_field == "cache_read_input_token_cost" else 0
),
cache_creation_input_tokens=5 if cache_field == "cache_creation_input_token_cost" else 0,
),
)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="cache-overlay",
function_id="cache-overlay",
)
logging_obj.model_call_details["model_id"] = profile
expected: Final = 10 * 0.02 + 3 * 0.04 + 5 * rate
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider=provider,
optional_params={"model_id": profile},
litellm_logging_obj=logging_obj,
) == pytest.approx(expected)
breakdown_key: Final = "cache_read_cost" if cache_field == "cache_read_input_token_cost" else "cache_creation_cost"
assert logging_obj.cost_breakdown is not None
assert logging_obj.cost_breakdown.get(breakdown_key, 0.0) == pytest.approx(5 * rate)
assert logging_obj.billed_token_rates is not None
assert (
logging_obj.billed_token_rates.cache_read_input_token_cost
if cache_field == "cache_read_input_token_cost"
else logging_obj.billed_token_rates.cache_creation_input_token_cost
) == rate
entry: Final = pricing_entry_for_cost_calc(
f"bedrock/converse/{model}", response, provider, False, None, None, None, logging_obj
)
assert entry is not None
assert entry[1]["input_cost_per_token"] == 0.02
assert entry[1]["output_cost_per_token"] == 0.04
assert entry[1][cache_field] == rate
assert litellm.model_cost[profile] == profile_before
monkeypatch.setattr(litellm, "model_cost", dict(catalog))
model_utils._invalidate_model_cost_lowercase_map()
model_utils.reapply_runtime_model_cost_registrations()
expected_after_reload: Final = (
10 * 0.02 + 3 * 0.04 + 5 * (rate if persist_across_reloads else catalog[profile][cache_field])
)
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider=provider,
optional_params={"model_id": profile},
litellm_logging_obj=logging_obj,
) == pytest.approx(expected_after_reload)
@pytest.mark.parametrize("profile", ("us.openai.missing-profile-test", "us.openai.unpriced-profile-test"))
def test_bedrock_cache_registration_with_unpriced_profile_keeps_registered_rates(
_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, profile: str
) -> None:
import litellm.utils as model_utils
model: Final = "openai.registered-cache-test"
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.model_cost,
model: {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
},
"us.openai.unpriced-profile-test": {"litellm_provider": "bedrock", "mode": "chat"},
},
)
model_utils._invalidate_model_cost_lowercase_map()
litellm.register_model({model: {"cache_read_input_token_cost": 0.001}}, persist_across_reloads=False)
response: Final = ModelResponse(
model=model,
usage=Usage(
prompt_tokens=15,
completion_tokens=3,
total_tokens=18,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5),
),
)
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": profile},
) == pytest.approx(10 * 0.01 + 3 * 0.02 + 5 * 0.001)
@pytest.mark.parametrize("direction", ("input", "output"))
@pytest.mark.parametrize("rate", (0.0, 0.004))
@pytest.mark.parametrize("prompt_tokens", (100, 300_001))
def test_bedrock_partial_registration_overrides_only_its_profile_direction(
_local_model_cost_map: None,
monkeypatch: pytest.MonkeyPatch,
direction: str,
rate: float,
prompt_tokens: int,
) -> None:
from litellm.cost_calculator import pricing_entry_for_cost_calc
model: Final = "openai.partial-pricing-test"
profile: Final = f"us.{model}"
profile_entry: Final = {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"input_cost_per_token_above_300k_tokens": 0.03,
"output_cost_per_token_above_300k_tokens": 0.06,
"cache_read_input_token_cost": 0.001,
}
monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, profile: dict(profile_entry)})
litellm.register_model({model: {f"{direction}_cost_per_token": rate}}, persist_across_reloads=False)
response: Final = ModelResponse(
model=model,
usage=Usage(
prompt_tokens=prompt_tokens,
completion_tokens=3,
total_tokens=prompt_tokens + 3,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5),
),
)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="partial-pricing",
function_id="partial-pricing",
)
logging_obj.model_call_details["model_id"] = profile
input_rate: Final = rate if direction == "input" else (0.03 if prompt_tokens > 300_000 else 0.01)
output_rate: Final = rate if direction == "output" else (0.06 if prompt_tokens > 300_000 else 0.02)
assert completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
litellm_logging_obj=logging_obj,
) == pytest.approx((prompt_tokens - 5) * input_rate + 3 * output_rate + 5 * 0.001)
assert logging_obj.billed_token_rates is not None
assert logging_obj.billed_token_rates.input_cost_per_token == input_rate
assert logging_obj.billed_token_rates.output_cost_per_token == output_rate
assert logging_obj.cost_breakdown is not None
assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(5 * 0.001)
entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj)
assert entry is not None
assert entry[1][f"{direction}_cost_per_token"] == rate
assert litellm.model_cost[profile] == profile_entry
@pytest.mark.parametrize("prompt_tokens", (299_999, 300_000, 300_001))
@pytest.mark.parametrize("cache_field", ("cache_read_input_token_cost", "cache_creation_input_token_cost"))
@pytest.mark.parametrize("tier_override", (None, 0.0, 0.003))
def test_bedrock_registered_cache_prices_preserve_dynamic_threshold_rates(
_local_model_cost_map: None,
monkeypatch: pytest.MonkeyPatch,
prompt_tokens: int,
cache_field: str,
tier_override: float | None,
) -> None:
import litellm.utils as model_utils
from litellm.cost_calculator import pricing_entry_for_cost_calc
model: Final = "openai.threshold-overlay-test"
profile: Final = f"us.{model}"
tier_field: Final = f"{cache_field}_above_300k_tokens"
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.model_cost,
profile: {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"input_cost_per_token_above_300k_tokens": 0.03,
"output_cost_per_token_above_300k_tokens": 0.06,
tier_field: 0.007,
},
},
)
model_utils._invalidate_model_cost_lowercase_map()
litellm.register_model(
{
model: {
cache_field: 0.001,
**({tier_field: tier_override} if tier_override is not None else {}),
}
},
persist_across_reloads=False,
)
usage: Final = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=3,
total_tokens=prompt_tokens + 3,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=5 if cache_field == "cache_read_input_token_cost" else 0
),
cache_creation_input_tokens=5 if cache_field == "cache_creation_input_token_cost" else 0,
)
response: Final = ModelResponse(model=model, usage=usage)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="threshold-overlay",
function_id="threshold-overlay",
)
logging_obj.model_call_details["model_id"] = profile
expected_input_rate: Final = 0.03 if prompt_tokens > 300_000 else 0.01
expected_output_rate: Final = 0.06 if prompt_tokens > 300_000 else 0.02
expected_cache_rate: Final = (
(tier_override if tier_override is not None else 0.007) if prompt_tokens > 300_000 else 0.001
)
total: Final = completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": profile},
litellm_logging_obj=logging_obj,
)
assert total == pytest.approx(
(prompt_tokens - 5) * expected_input_rate + 3 * expected_output_rate + 5 * expected_cache_rate
)
breakdown_key: Final = "cache_read_cost" if cache_field == "cache_read_input_token_cost" else "cache_creation_cost"
assert logging_obj.cost_breakdown is not None
assert logging_obj.cost_breakdown.get(breakdown_key, 0.0) == pytest.approx(5 * expected_cache_rate)
assert logging_obj.billed_token_rates is not None
assert logging_obj.billed_token_rates.input_cost_per_token == expected_input_rate
assert logging_obj.billed_token_rates.output_cost_per_token == expected_output_rate
entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj)
assert entry is not None
assert entry[1]["input_cost_per_token_above_300k_tokens"] == 0.03
assert entry[1][tier_field] == (tier_override if tier_override is not None else 0.007)

95
tests/unit/test_main.py
View file

@ -4269,3 +4269,98 @@ def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice):
)
assert exc_info.value.status_code == 400
assert f"tool_choice={tool_choice}" in str(exc_info.value)
@pytest.mark.parametrize("stream", (False, True))
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.parametrize("profile,custom_price", ((None, False), ("us.openai.gpt-6-sol", False), (None, True)))
@pytest.mark.asyncio
async def test_converse_fallback_prices_the_current_attempt(
monkeypatch: pytest.MonkeyPatch, stream: bool, asynchronous: bool, profile: str | None, custom_price: bool
) -> None:
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.rust_bridge import configuration
CONVERSE_ANSWER: Final = json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
"metrics": {"latencyMs": 1},
}
).encode()
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setenv("LITELLM_RUST", "0")
configuration.rust(False)
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
failed_profile: Final = "us.amazon.nova-micro-v1:0"
model: Final = "openai.gpt-6-sol"
logging_obj: Final = LiteLLMLogging(
model="amazon.nova-micro-v1:0",
messages=[{"role": "user", "content": "hi"}],
stream=stream,
call_type="acompletion" if asynchronous else "completion",
start_time=datetime(2026, 1, 1),
litellm_call_id="bedrock-fallback-pricing",
function_id="bedrock-fallback-pricing",
)
def respond(request: httpx.Request) -> httpx.Response:
if failed_profile in request.url.path:
return httpx.Response(400, json={"message": "Input is too long for requested model."})
assert request.url.path == f"/model/{profile or model}/converse"
return httpx.Response(200, content=CONVERSE_ANSWER)
client: Final = (
AsyncHTTPHandler(transport=httpx.MockTransport(respond))
if asynchronous
else HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
)
kwargs: Final = {
"messages": [{"role": "user", "content": "hi"}],
"api_key": "test-bearer",
"aws_region_name": "us-east-1",
"client": client,
"litellm_logging_obj": logging_obj,
"num_retries": 0,
"stream": stream,
"fake_stream": stream,
}
if asynchronous:
with pytest.raises(litellm.BadRequestError):
await litellm.acompletion(model="bedrock/converse/amazon.nova-micro-v1:0", model_id=failed_profile, **kwargs)
else:
with pytest.raises(litellm.BadRequestError):
litellm.completion(model="bedrock/converse/amazon.nova-micro-v1:0", model_id=failed_profile, **kwargs)
overrides: Final = {
**({"model_id": profile} if profile is not None else {}),
**({"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} if custom_price else {}),
}
response: Final = (
await litellm.acompletion(model=f"bedrock/converse/{model}", **kwargs, **overrides)
if asynchronous
else litellm.completion(model=f"bedrock/converse/{model}", **kwargs, **overrides)
)
result: Final = (
litellm.stream_chunk_builder(
[chunk async for chunk in response] if asynchronous else list(response), logging_obj=response.logging_obj
)
if stream
else response
)
assert result is not None
prices: Final = litellm.get_model_info(profile or model, custom_llm_provider="bedrock")
expected: Final = (
0.001 + 0.002
if custom_price
else prices["input_cost_per_token"] + prices["output_cost_per_token"]
)
active_logging: Final = response.logging_obj if stream else logging_obj
assert active_logging._response_cost_calculator(result) == pytest.approx(expected)
if asynchronous:
await client.close()
else:
client.close()
configuration.reset_rust_configuration()

181
tests/unit/test_router_model_cost_isolation.py
View file

@ -761,6 +761,178 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override():
_restore_model_cost_entries(model_keys)
@pytest.mark.parametrize("service_tier", ("default", "priority", "flex"))
@pytest.mark.parametrize(
"field",
("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost", "input_cost_per_token_priority"),
)
@pytest.mark.parametrize("rate", (0.0, 0.004))
@pytest.mark.parametrize("explicit_base", (False, True))
@pytest.mark.parametrize("provider", (None, "bedrock"))
@pytest.mark.parametrize("registration_path", ("direct", "startup", "reload"))
def test_bedrock_deployment_inherits_invoked_profile_around_custom_prices(
monkeypatch: pytest.MonkeyPatch,
field: str,
rate: float,
explicit_base: bool,
provider: str | None,
registration_path: str,
service_tier: str,
) -> None:
from litellm.cost_calculator import completion_cost
from litellm.types.utils import ModelResponse, Usage, PromptTokensDetailsWrapper
model: Final = "openai.gpt-5.6-sol"
profile: Final = f"us.{model}"
base_model: Final = f"global.{model}"
profile_prices: Final = {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"cache_read_input_token_cost": 0.001,
"input_cost_per_token_priority": 0.05,
"output_cost_per_token_priority": 0.06,
"input_cost_per_token_flex": 0.007,
"output_cost_per_token_flex": 0.008,
}
base_prices: Final = {**profile_prices, "input_cost_per_token": 0.03, "output_cost_per_token": 0.04}
monkeypatch.setattr(
litellm, "model_cost", {**litellm.model_cost, profile: dict(profile_prices), base_model: dict(base_prices)}
)
deployment: Final = Deployment(
model_name="profile-priced",
litellm_params=LiteLLM_Params(
model=f"converse/{model}" if provider else f"bedrock/converse/{model}",
model_id=profile,
custom_llm_provider=provider,
),
model_info=ModelInfo(
**{
"id": "profile-priced-deployment",
field: rate,
**({"base_model": f"bedrock/{base_model}"} if explicit_base else {}),
}
),
)
if registration_path == "direct":
Router._register_deployment_pricing(deployment)
else:
router: Final = Router(model_list=[deployment.to_json(exclude_none=True)])
if registration_path == "reload":
router._replay_model_cost_registrations()
response: Final = ModelResponse(
model=model,
usage=Usage(
prompt_tokens=100,
completion_tokens=3,
total_tokens=103,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5),
),
)
catalog: Final = base_prices if explicit_base else profile_prices
input_field: Final = "input_cost_per_token" + (f"_{service_tier}" if service_tier != "default" else "")
output_field: Final = "output_cost_per_token" + (f"_{service_tier}" if service_tier != "default" else "")
input_rate: Final = rate if field == input_field else catalog[input_field]
output_rate: Final = rate if field == output_field else catalog[output_field]
cache_rate: Final = rate if field == "cache_read_input_token_cost" else 0.001
assert completion_cost(
model=model,
completion_response=response,
custom_llm_provider="bedrock",
custom_pricing=True,
router_model_id="profile-priced-deployment",
service_tier=service_tier,
) == pytest.approx(95 * input_rate + 3 * output_rate + 5 * cache_rate)
assert litellm.model_cost[profile] == profile_prices
@pytest.mark.parametrize("catalog_entry", (None, {"litellm_provider": 1, "mode": "chat"}))
@pytest.mark.parametrize("known_model", (False, True))
@pytest.mark.parametrize("explicit_base", (False, True))
@pytest.mark.parametrize("service_tier", ("default", "priority", "flex"))
@pytest.mark.parametrize("override", (None, "cache_read_input_token_cost", "input_cost_per_token_priority"))
@pytest.mark.parametrize("reload", (False, True))
def test_bedrock_deployment_keeps_custom_prices_when_profile_lookup_fails(
monkeypatch: pytest.MonkeyPatch,
catalog_entry: dict[str, object] | None,
known_model: bool,
explicit_base: bool,
service_tier: str,
override: str | None,
reload: bool,
) -> None:
from litellm.cost_calculator import completion_cost
from litellm.types.utils import ModelResponse, Usage, PromptTokensDetailsWrapper
model: Final = "anthropic.unmapped-profile-pricing-test"
profile: Final = (
"us.anthropic.invalid-profile-pricing-test"
if catalog_entry is not None
else "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque-test"
)
configured_prices: Final = {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.03,
"output_cost_per_token": 0.04,
"cache_read_input_token_cost": 0.001,
"input_cost_per_token_priority": 0.05,
"output_cost_per_token_priority": 0.06,
"input_cost_per_token_flex": 0.007,
"output_cost_per_token_flex": 0.008,
}
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.model_cost,
**({model: dict(configured_prices)} if known_model else {}),
**({profile: catalog_entry} if catalog_entry is not None else {}),
},
)
deployment: Final = Deployment(
model_name="custom-priced",
litellm_params=LiteLLM_Params(model=f"bedrock/converse/{model}", model_id=profile),
model_info=ModelInfo(
id="unmapped-profile-deployment",
input_cost_per_token=0.01,
output_cost_per_token=0.02,
**({override: 0.0} if override else {}),
**({"base_model": f"bedrock/{profile}"} if explicit_base else {}),
),
)
router: Final = Router(model_list=[deployment.to_json(exclude_none=True)])
if reload:
router._replay_model_cost_registrations()
inherit_configured: Final = known_model and not explicit_base
tier_rates: Final = inherit_configured and service_tier != "default"
input_rate: Final = configured_prices[f"input_cost_per_token_{service_tier}"] if tier_rates else 0.01
output_rate: Final = configured_prices[f"output_cost_per_token_{service_tier}"] if tier_rates else 0.02
effective_input: Final = 0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate
cache_rate: Final = 0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input)
assert completion_cost(
model=model,
completion_response=ModelResponse(
model=model,
usage=Usage(
prompt_tokens=100,
completion_tokens=3,
total_tokens=103,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=80),
),
),
custom_llm_provider="bedrock",
custom_pricing=True,
router_model_id="unmapped-profile-deployment",
service_tier=service_tier,
) == pytest.approx(20 * effective_input + 80 * cache_rate + 3 * output_rate)
if known_model:
assert {
key: value for key, value in litellm.model_cost[model].items() if key in configured_prices
} == configured_prices
def test_partial_custom_pricing_inherits_builtin_cache_pricing():
"""A deployment that overrides only input/output cost on a cache-supporting
model must still bill cache_read and cache_creation tokens. Before the
@ -1771,14 +1943,7 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None:
)
rebuilt = litellm.model_cost["priced-id"]
assert at_boot.items() <= rebuilt.items(), (
f"the rebuild changed or dropped a field the boot registration wrote: "
f"{ {k: (v, rebuilt.get(k)) for k, v in at_boot.items() if rebuilt.get(k) != v} }"
)
assert {field: rebuilt[field] for field in set(rebuilt) - set(at_boot)} == {
"db_model": False,
"member_auto_router": False,
}
assert at_boot == rebuilt
assert "member_auto_router" not in litellm.model_cost["gpt-4o"]
assert router.model_list
finally:

48
tests/unit/test_utils.py
View file

@ -4205,7 +4205,11 @@ async def test_builtin_string_callback_registers_when_subclass_already_active(
assert any(type(cb) is S3Logger for cb in litellm._async_success_callback)
def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypatch):
@pytest.mark.parametrize("rate", (0.0, 0.000123))
@pytest.mark.parametrize("updated_rate", (None, 0.0, 0.000456))
def test_reapply_runtime_registrations_replays_register_model_overrides(
monkeypatch: pytest.MonkeyPatch, rate: float, updated_rate: float | None
) -> None:
"""
register_model is the documented way to override pricing for a model. A
price-data reload swaps litellm.model_cost for a freshly fetched catalog,
@ -4236,11 +4240,17 @@ def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypa
"openai/gpt-4o": {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 0.000123,
"input_cost_per_token": rate,
}
}
)
litellm.register_model({"openai/gpt-4o": {"supports_function_calling": True, "input_cost_per_token": None}})
litellm.register_model(
{"openai/gpt-4o": {"input_cost_per_token": str(updated_rate) if updated_rate is not None else None}}
)
litellm.register_model({"openai/gpt-4o": {"supports_function_calling": False}})
litellm.model_cost = {
"openai/gpt-4o": {
"litellm_provider": "openai",
@ -4252,7 +4262,9 @@ def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypa
_invalidate_model_cost_lowercase_map()
reapply_runtime_model_cost_registrations()
assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == 0.000123
expected_rate: Final = rate if updated_rate is None else updated_rate
assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == expected_rate
assert litellm.model_cost["openai/gpt-4o"]["supports_function_calling"] is False
assert litellm.model_cost["openai/gpt-4o"]["max_input_tokens"] == 4242
finally:
litellm.model_cost = saved_model_cost
@ -6362,3 +6374,33 @@ def test_function_setup_never_logs_the_ocr_data_uri_payload() -> None:
assert logged == [{"role": "user", "content": f"data:application/pdf;base64 ({len(payload)} chars)"}]
assert payload not in str(logged)
@pytest.mark.parametrize(
"key", ("openai.test-converse", "bedrock/openai.test-converse", "bedrock/converse/openai.test-converse")
)
def test_runtime_model_cost_registrations_resolve_bedrock_routing(monkeypatch: pytest.MonkeyPatch, key: str) -> None:
from litellm import utils as model_utils
rates: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
monkeypatch.setattr(
litellm, "model_cost", {**litellm.model_cost, key: {"litellm_provider": "bedrock_converse", "mode": "chat"}}
)
monkeypatch.setattr(model_utils, "_runtime_registered_fields", {})
litellm.register_model({key: rates}, persist_across_reloads=False)
registrations: Final = dict(
model_utils.runtime_model_cost_registrations("bedrock/converse/openai.test-converse", "bedrock")
)
assert registrations[key] == rates
assert model_utils.runtime_model_cost_registrations("bedrock/converse/openai.other-converse", "bedrock") == ()
def test_registered_fields_reset_when_model_cost_is_replaced(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", {})
litellm.register_model({"registered-fields-probe": {"input_cost_per_token": 1e-06, "litellm_provider": "openai"}})
assert "registered-fields-probe" in litellm.utils._registered_fields_for_current_map()
monkeypatch.setattr(litellm, "model_cost", {})
assert litellm.utils._registered_fields_for_current_map() == {}