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

fix(register_model)!: keep bedrock geo deployments priced across price data reloads (#41523)

Browse source 
* fix(register_model): keep bedrock geo deployments priced across price data reloads

A request-time pricing registration inserting into litellm.model_cost while a
price data reload walked it could abort the reload (dictionary changed size
during iteration in add_known_models), or trip the lookup map rebuild inside
get_model_info so the deployment replay fell through to the unmapped path and
wrote a raw bedrock/au.anthropic.<model> entry with no token pricing, a $0
shadow re-rolled on every later reload. Walk snapshots of the map, bind the
case-insensitive lookup map locally, and let the unmapped path inherit the
whole built-in entry of a provider get_model_info would accept.

* fix(register_model): keep synthesized zero costs out of inherited built-in pricing

An unmapped alias that inherits its built-in entry now takes only the token
prices the catalogue really carries, so a synthesized zero never lands next to
the caller's custom price. The router regression test reports the writer's
failures instead of swallowing them, and the local names carry Final.

* fix(register_model): match inherited built-ins by the key provider prefix when the value names none

* fix(utils): resolve a tripped lookup from the key's own catalog row before its region sibling

* fix(register_model): keep custom flat prices above inherited tier rates

* fix(register_model): keep the other direction's tier rates when a custom flat price covers one

* fix(register_model): keep a catalog row's tiers when its lookup trips, and snapshot request-path cost map walks

A tripped get_model_info on a key that is itself a catalog row now inherits that row whole, so a custom flat price bills a long prompt at the row's above-threshold rate on both lookup paths instead of flipping until the next reload. The fal_ai, fireworks, ollama and gpt-5 request-path walks over litellm.model_cost iterate a snapshot so a concurrent reload cannot raise mid-request.

* fix(register_model): price route-prefixed geo keys on a tripped lookup and replace cost map entries atomically

* fix(register_model): move every alias of a cost map entry to its registered price

* perf(register_model): read alias names from the shared entry instead of scanning the cost map

Each registration scanned every cost map entry to find the names sharing it,
about 1 ms on the 4,189-entry map, and a proxy request to a custom-priced
deployment registers twice. Aliased entries now carry the names they are
served under, so a registration reads them directly, and an alias the caller
pointed at its own entry keeps that entry. The tests this PR added lose their
docstrings

* fix(utils): keep capability flags and write through on a tripped lookup in register_model

A key the lookup cannot map now seeds its registered entry from the capability
rules its unregistered lookup resolves to, so registering a region alias no
longer turns off supports_system_messages and similar flags. A lookup that
trips mid-walk on a key naming a catalog row writes into that row instead of
leaving a shadow key that goes stale. Adds regression tests for both and for
the request-path snapshot walks.

* fix(utils): retry a tripped registration lookup and rebuild a registration a reload races

A lookup a concurrent write trips is retried instead of writing into the
catalog row its key names, so a per-request price never reprices a shared
row. The inherited entry leaves the matched row's provider behind, so a
provider-less azure_ai key stays priced, and a registration whose map a
reload swapped mid-build is rebuilt from the new map.

* fix(utils): write each registration into the map it was built from and rebuild it when a reload replaced that map

* test(integration): price data reload keeps every registered deployment priced

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-17 19:10:04 +00:00 • committed by GitHub
parent b5037f28d1
commit 9e9dccb9d4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1546 additions and 133 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/__init__.py
View file

@ -742,7 +742,7 @@ def is_openai_finetune_model(key: str) -> bool:
def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None:
for key, value in model_cost_map.items():
for key, value in tuple(model_cost_map.items()):
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
open_ai_chat_completion_models.add(key)
elif value.get("litellm_provider") == "text-completion-openai":

36
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -557,22 +557,39 @@ def get_model_cost_map_loaded_at() -> "datetime | None":
return _cost_map_source_info.loaded_at
def _expand_model_aliases(model_cost: dict) -> dict:
class AliasedCostEntry(dict[str, object]):
__slots__ = ("names",)
def __init__(self, entry: Mapping[str, object], names: tuple[str, ...]) -> None:
super().__init__(entry)
self.names: tuple[str, ...] = names
def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, ...]]:
return MappingProxyType(
{
owner: (owner, *(alias for alias, alias_owner in alias_owners.items() if alias_owner == owner))
for owner in dict.fromkeys(alias_owners.values())
}
)
def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, dict[str, object]]:
"""
Expand ``aliases`` lists in model cost entries into top-level entries.
Each alias gets a reference to the **same** dict object as the canonical
entry (zero memory overhead). The ``aliases`` key is removed from the
entry so downstream code never sees it.
Each alias gets a reference to the **same** ``AliasedCostEntry`` as the
canonical entry, which carries every name it is served under. The
``aliases`` key is removed from the entry so downstream code never sees it.
If an alias collides with an existing canonical entry the alias is
skipped and a warning is logged.
"""
aliases_to_add: Final[dict[str, dict]] = {}
alias_owners: Final[dict[str, str]] = {}
keys_with_aliases: Final[list[str]] = []
for model_name, model_info in model_cost.items():
aliases: list | None = model_info.get("aliases")
aliases: object = model_info.get("aliases")
if aliases is None:
continue
keys_with_aliases.append(model_name)
@ -594,7 +611,7 @@ def _expand_model_aliases(model_cost: dict) -> dict:
model_name,
)
continue
if alias in aliases_to_add:
if alias in alias_owners:
verbose_logger.warning(
"LiteLLM model alias conflict: alias '%s' (from '%s') "
"was already claimed by another entry — skipping.",
@ -602,13 +619,14 @@ def _expand_model_aliases(model_cost: dict) -> dict:
model_name,
)
continue
aliases_to_add[alias] = model_info # same dict reference
alias_owners[alias] = model_name
# Remove the ``aliases`` key from entries so it doesn't pollute model info
for key in keys_with_aliases:
model_cost[key].pop("aliases", None)
model_cost.update(aliases_to_add)
for owner, names in _alias_groups(alias_owners).items():
model_cost.update(dict.fromkeys(names, AliasedCostEntry(model_cost[owner], names)))
return model_cost

2
litellm/llms/fal_ai/cost_calculator.py
View file

@ -87,7 +87,7 @@ def _keyed_rows(model: str, quality: str) -> tuple[tuple[int, int, float], ...]:
suffix: Final = f"/{model}"
return tuple(
(width, height, float(raw_cost))
for key in litellm.model_cost
for key in (*litellm.model_cost,)
if isinstance(key, str) and key.startswith(prefix) and key.endswith(suffix)
for size in (key[len(prefix) : -len(suffix)],)
for dimensions in (_parse_keyed_dimensions(size),)

2
litellm/llms/fireworks_ai/chat/transformation.py
View file

@ -544,7 +544,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
return cached[2]
index: Final[list[tuple[str, dict]]] = []
for key, model_info in model_cost.items():
for key, model_info in (*model_cost.items(),):
if not key.startswith("fireworks_ai/"):
continue
if not isinstance(model_info, dict):

2
litellm/llms/ollama/common_utils.py
View file

@ -197,7 +197,7 @@ class OllamaModelInfo(BaseLLMModelInfo):
"ollama/" + stripped_model,
"ollama_chat/" + stripped_model,
}
model_cost_keys: Final = {key.lower() for key in model_cost}
model_cost_keys: Final = {key.lower() for key in (*model_cost,)}
return any(name.lower() in model_cost_keys for name in potential_model_names)
@staticmethod

4
litellm/llms/openai/chat/gpt_5_transformation.py
View file

@ -28,7 +28,9 @@ def _catalogue_declares_default_effort() -> bool:
the one litellm gave before it existed. Scanning costs ~80us on the largest published map and
only on the fallback path, which is noise beside the request it precedes.
"""
return any(isinstance(entry, dict) and "default_reasoning_effort" in entry for entry in litellm.model_cost.values())
return any(
isinstance(entry, dict) and "default_reasoning_effort" in entry for entry in (*litellm.model_cost.values(),)
)
def _normalize_reasoning_effort_for_chat_completion(

396
litellm/utils.py
View file

@ -44,7 +44,7 @@ from httpx import Proxy
from httpx._utils import get_environment_proxies
from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # OpenAI parser module is private
from openai.types.chat.completion_create_params import ResponseFormat
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, with_config
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config
import litellm
import litellm.litellm_core_utils
@ -89,6 +89,7 @@ from litellm.litellm_core_utils.fallback_generalizations import (
match_capability_generalizations,
match_fill_missing_generalizations,
)
from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
from litellm.litellm_core_utils.tokenizer import strip_special_tokens
@ -3303,46 +3304,176 @@ _BEDROCK_REGION_PREFIXES: Final = (
"ap-northeast-1.",
)
_CACHE_PRICING_FIELDS: Final = (
"cache_creation_input_token_cost",
"cache_creation_input_token_cost_above_1hr",
"cache_creation_input_token_cost_above_200k_tokens",
"cache_read_input_token_cost",
"cache_read_input_token_cost_above_200k_tokens",
_TOKEN_COST_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token")
def _builtin_cost_map() -> dict[str, Mapping[str, object]]:
return cast(dict[str, Mapping[str, object]], litellm.model_cost)
def _replace_cost_map_entry(cost_map: dict[str, Mapping[str, object]], key: str, entry: Mapping[str, object]) -> None:
current: Final = cost_map.get(key)
merged: Final = {**(current or {}), **entry}
if not isinstance(current, AliasedCostEntry):
cost_map.update({key: merged})
return
sharing_names: Final = tuple(name for name in current.names if cost_map.get(name) is current)
cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, current.names)))
def _without_synthesized_costs(
model_info: Mapping[str, object], key: str, value: Mapping[str, object]
) -> Mapping[str, object]:
"""``_get_model_info_helper`` synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` = 0 when the raw
entry has neither (#30198). Persisting those zeros would turn a name-priced entry into a free one and make
``_is_cost_explicitly_configured`` disable budget enforcement on the next re-registration."""
builtin_cost_map: Final = _builtin_cost_map()
raw_fields: Final = frozenset(builtin_cost_map.get(str(model_info["key"]), builtin_cost_map.get(key, {})))
return MappingProxyType(
{
name: field_value
for name, field_value in model_info.items()
if name not in _TOKEN_COST_FIELDS or name in raw_fields or name in value
}
)
_DIRECTION_RATE_FIELDS: Final = MappingProxyType(
{
"input_cost_per_token": ("input_cost_per_token",),
"output_cost_per_token": ("output_cost_per_token", "output_cost_per_reasoning_token"),
}
)
_RATE_TABLE: Final = TypeAdapter(Mapping[str, object])
_RATE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, object] | None:
def _is_overridden_rate(name: str, overridden_fields: tuple[str, ...]) -> bool:
is_variant: Final = name.startswith(tuple(f"{rate_field}_" for rate_field in overridden_fields))
return name in overridden_fields or (is_variant and "_cache_" not in name)
def _without_overridden_rates(rates: Mapping[str, object], overridden_fields: tuple[str, ...]) -> Mapping[str, object]:
return MappingProxyType(
{name: rate for name, rate in rates.items() if not _is_overridden_rate(name, overridden_fields)}
)
def _rate_table_with_flat_prices(
name: str, table: object, overridden_fields: tuple[str, ...], flat_prices: Mapping[str, object]
) -> object:
try:
if name == "tiered_pricing":
return [
{**_without_overridden_rates(tier, overridden_fields), **flat_prices}
for tier in _RATE_TIERS.validate_python(table)
]
if name == "off_peak_pricing":
return dict(_without_overridden_rates(_RATE_TABLE.validate_python(table), overridden_fields))
except ValidationError:
return table
return table
def _without_overridden_price_variants(
entry: Mapping[str, object], value: Mapping[str, object]
) -> Mapping[str, object]:
"""A caller's flat ``input_cost_per_token`` / ``output_cost_per_token`` is that direction's price above every
token threshold, for every service tier and batch, in every ``tiered_pricing`` tier and off-peak window, so each
inherited rate that would stand in for it there, the reasoning-token rate included for output, is dropped, and
each tier carries the flat price instead."""
flat_prices: Final = MappingProxyType(
{cost_field: value[cost_field] for cost_field in _TOKEN_COST_FIELDS if value.get(cost_field) is not None}
)
overridden_fields: Final = tuple(
itertools.chain.from_iterable(_DIRECTION_RATE_FIELDS[cost_field] for cost_field in flat_prices)
)
if not overridden_fields:
return entry
return MappingProxyType(
{
name: _rate_table_with_flat_prices(name, field_value, overridden_fields, flat_prices)
for name, field_value in entry.items()
if not _is_overridden_rate(name, overridden_fields)
}
)
@dataclass(frozen=True, slots=True)
class _BuiltinCostMatch:
key: str
entry: Mapping[str, object]
names_the_row: bool
@dataclass(frozen=True, slots=True)
class _LookupMiss:
tripped: bool
@dataclass(frozen=True, slots=True)
class _FallbackRegistration:
entry: Mapping[str, object]
matched_builtin: bool
def _fallback_registration(key: str, provider: str, value: Mapping[str, object]) -> _FallbackRegistration:
"""The key keeps its own entry, seeded from the capability rules an unregistered lookup resolves it to, with the
matched built-in entry on top. The matched row's provider stays behind, since the key's own lookup can reject it
(an ``azure_ai/`` key never resolves to an ``openai`` entry), and the key keeps the provider its caller gave."""
match: Final = _resolve_builtin_model_cost_entry(key=key, provider=provider)
generalized: Final[Mapping[str, object]] = match_capability_generalizations(key) or {}
if match is None:
return _FallbackRegistration(entry=generalized, matched_builtin=False)
inherited: Final = match.entry if match.names_the_row else _without_overridden_price_variants(match.entry, value)
own_fields: Final = {name: field_value for name, field_value in inherited.items() if name != "litellm_provider"}
return _FallbackRegistration(entry={**generalized, **own_fields}, matched_builtin=True)
def _resolve_builtin_model_cost_entry(key: str, provider: str) -> _BuiltinCostMatch | None:
"""Best-effort lookup of a built-in ``model_cost`` entry for a custom key
whose shape ``get_model_info`` cannot resolve (repeated provider prefixes
like ``bedrock/bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region
aliases).
aliases), or whose lookup tripped.
Returns a copy of the matching entry so the caller can inherit its defaults
(most importantly cache pricing) without mutating the shared built-in.
Returns ``None`` when no safe match exists.
The key itself is tried first, then the key behind its provider prefixes
and bedrock route, so a catalog row whose lookup tripped inherits from
itself rather than from a region-stripped sibling priced at the base rate.
``names_the_row`` is set when the match is the row the key names, which
``get_model_info`` resolves it to, so the caller keeps that row's rates as
an untripped registration does. Returns ``None`` when no candidate exists
under a provider ``get_model_info`` would accept for ``provider``, or, when
the value names none, for the provider prefix the key itself carries,
which is what ``get_model_info`` derives when it reads the key back.
"""
candidates: Final[list[str]] = []
from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix
segments: Final = key.split("/")
idx = 0
while idx < len(segments) - 1 and segments[idx] in LlmProvidersSet:
idx += 1
candidates.append("/".join(segments[idx:]))
base: Final = candidates[-1] if candidates else key
for region_prefix in _BEDROCK_REGION_PREFIXES:
if base.startswith(region_prefix):
candidates.append(base[len(region_prefix) :])
if provider:
stripped: Final = _strip_model_name(model=base, custom_llm_provider=provider)
if stripped != base:
candidates.append(stripped)
match_provider: Final = provider or (segments[0] if len(segments) > 1 and segments[0] in LlmProvidersSet else "")
provider_prefix_count: Final = next(
(idx for idx, segment in enumerate(segments[:-1]) if segment not in LlmProvidersSet), len(segments) - 1
)
prefix_stripped: Final = tuple("/".join(segments[idx:]) for idx in range(1, provider_prefix_count + 1))
prefix_free: Final = prefix_stripped[-1] if prefix_stripped else key
base: Final = (
strip_bedrock_routing_prefix(prefix_free) if match_provider in ("bedrock", "bedrock_converse") else prefix_free
)
region_stripped: Final = tuple(
base[len(region_prefix) :] for region_prefix in _BEDROCK_REGION_PREFIXES if base.startswith(region_prefix)
)
stripped: Final = _strip_model_name(model=base, custom_llm_provider=match_provider) if match_provider else base
row_names: Final = (key, *prefix_stripped, base, stripped)
candidates: Final = (key, *prefix_stripped, base, *region_stripped, stripped)
builtin_cost_map: Final = _builtin_cost_map()
for candidate in candidates:
entry = litellm.model_cost.get(candidate)
if entry is not None and entry.get("litellm_provider") is not None:
return dict(entry)
entry = builtin_cost_map.get(candidate)
if (
entry is not None
and entry.get("litellm_provider") is not None
and _check_provider_match(dict(entry), match_provider)
):
return _BuiltinCostMatch(key=candidate, entry=dict(entry), names_the_row=candidate in row_names)
return None
@ -3360,19 +3491,115 @@ def is_generalized_model_info(model_info: ModelInfo) -> bool:
return key not in litellm.model_cost and match_capability_generalizations(key) is not None
def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | None:
"""Resolve ``model`` to its built-in cost-map entry for registration merging.
_REGISTRATION_ATTEMPTS: Final = 3
Returns ``None`` when the lookup raises or when it resolved via a
fallback-generalization capability rule. A rule-derived entry carries no
pricing, so treating it as a hit would skip the built-in cache-pricing
inheritance for prefix-mangled keys.
"""
def _lookup_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | _LookupMiss:
try:
info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except ModelNotMappedError:
return _LookupMiss(tripped=False)
except Exception:
return None
return None if is_generalized_model_info(info) else info
return _LookupMiss(tripped=True)
return _LookupMiss(tripped=False) if is_generalized_model_info(info) else info
def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | _LookupMiss:
"""Resolve ``model`` to its built-in cost-map entry for registration merging.
Misses when the lookup cannot map the model or resolved it via a
fallback-generalization capability rule, since a rule-derived entry carries no
pricing. A lookup that raised anything else, such as a concurrent registration
changing the map mid-walk, is retried, so a transient failure registers exactly
where the untripped lookup would.
"""
lookups: Final = (_lookup_for_registration(model, custom_llm_provider) for _ in range(_REGISTRATION_ATTEMPTS))
return next(
(lookup for lookup in lookups if not (isinstance(lookup, _LookupMiss) and lookup.tripped)),
_LookupMiss(tripped=True),
)
@dataclass(frozen=True, slots=True)
class _RegistrationWrite:
key: str
entry: Mapping[str, object]
def _registration_base(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
) -> _RegistrationWrite:
provider: Final = str(value.get("litellm_provider") or "")
if provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO or any(
key.startswith(f"{p}/") for p in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
):
return _RegistrationWrite(key=key, entry=dict(_builtin_cost_map().get(key, {})))
builtin_model_info: Final = _get_builtin_model_info_for_registration(
model=key, custom_llm_provider=custom_llm_provider
)
if not isinstance(builtin_model_info, _LookupMiss):
return _RegistrationWrite(
key=builtin_model_info["key"],
entry=dict(_without_synthesized_costs(builtin_model_info, key=key, value=value)),
)
fallback: Final = _fallback_registration(key=key, provider=provider, value=value)
if not fallback.matched_builtin and (
value.get("cache_creation_input_token_cost") is None
and value.get("cache_read_input_token_cost") is None
and value.get("tiered_pricing") is None
and (value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None)
):
verbose_logger.warning(
"register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info",
warning_display_name or key,
)
return _RegistrationWrite(key=key, entry=dict(fallback.entry))
def _registration_write(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
) -> _RegistrationWrite:
base: Final = _registration_base(key, value, custom_llm_provider, warning_display_name)
# ``get_model_info`` returns ``litellm_provider: None`` when the
# provider is unknown (e.g. custom deployments registered via
# ``Router.add_deployment``). Persisting that None into
# ``litellm.model_cost`` causes ``_check_provider_match`` to drop
# custom pricing on subsequent cost lookups.
existing_model: Final = {
name: field_value
for name, field_value in base.entry.items()
if name != "litellm_provider" or field_value is not None
}
return _RegistrationWrite(key=base.key, entry=update_dictionary(existing_model, value))
def _registration_attempt(
key: str,
value: Mapping[str, object],
custom_llm_provider: str | None,
warning_display_name: str | None,
final_attempt: bool,
) -> _RegistrationWrite | None:
source: Final = _builtin_cost_map()
write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name)
_replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry)
return write if final_attempt or litellm.model_cost is source else None
def _register_on_current_map(
key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None
) -> _RegistrationWrite:
"""A price data reload that swaps the map while the entry is being built or written would otherwise get the old
map's prices written over its own, or drop the entry with the old map, so each attempt writes into the map it was
built from and the entry is rebuilt when a reload replaced that map before the write finished. The last attempt
writes into the current map, so the entry is never lost."""
attempts: Final = (
_registration_attempt(
key, value, custom_llm_provider, warning_display_name, final_attempt=attempt == _REGISTRATION_ATTEMPTS - 1
)
for attempt in range(_REGISTRATION_ATTEMPTS)
)
return next(write for write in attempts if write is not None)
_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload
@ -3492,79 +3719,15 @@ def register_model(
for _registered_key, _registered_value in _registrations.items():
_runtime_registered_model_cost[_registered_key] = dict(_registered_value)
_skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
for key, value in loaded_model_cost.items():
## get model info ##
provider = value.get("litellm_provider", "")
_key_str = str(key)
if provider in _skip_get_model_info_providers or any(
_key_str.startswith(f"{p}/") for p in _skip_get_model_info_providers
):
existing_model = litellm.model_cost.get(key, {})
model_cost_key = key
else:
builtin_model_info = _get_builtin_model_info_for_registration(
model=_key_str, custom_llm_provider=custom_llm_provider
)
if builtin_model_info is not None:
existing_model = cast(dict, builtin_model_info)
model_cost_key = existing_model["key"]
else:
# An exact entry ends the lookup ladder before the capability rules are
# consulted, so seed from them: otherwise registering an unmapped model
# shadows the very defaults it would have resolved to unregistered.
existing_model = dict(match_capability_generalizations(_key_str) or {})
model_cost_key = key
builtin_entry = _resolve_builtin_model_cost_entry(key=_key_str, provider=provider)
if builtin_entry is not None:
for field in _CACHE_PRICING_FIELDS:
if value.get(field) is None and builtin_entry.get(field) is not None:
existing_model[field] = builtin_entry[field]
elif (
value.get("cache_creation_input_token_cost") is None
and value.get("cache_read_input_token_cost") is None
and value.get("tiered_pricing") is None
and (
value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None
)
):
verbose_logger.warning(
"register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info",
warning_display_name or key,
)
# ``get_model_info`` returns ``litellm_provider: None`` when the
# provider is unknown (e.g. custom deployments registered via
# ``Router.add_deployment``). Persisting that None into
# ``litellm.model_cost`` causes ``_check_provider_match`` to drop
# custom pricing on subsequent cost lookups.
if existing_model.get("litellm_provider") is None:
existing_model.pop("litellm_provider", None)
# Same pattern for cost fields (#30198): ``_get_model_info_helper``
# synthesizes ``input_cost_per_token`` / ``output_cost_per_token``
# = 0 when they are absent from the raw entry. Writing those zeros
# back flips a sparse entry from "no cost keys" (priced via name)
# to "cost keys = 0" (free), which makes
# ``_is_cost_explicitly_configured`` return True and silently
# disables budget enforcement on the next re-registration.
_raw_entry = litellm.model_cost.get(model_cost_key)
if _raw_entry is None:
_raw_entry = litellm.model_cost.get(key)
if _raw_entry is None:
_raw_entry = {}
for _cost_field in ("input_cost_per_token", "output_cost_per_token"):
if _cost_field not in _raw_entry and _cost_field not in value:
existing_model.pop(_cost_field, None)
## override / add new keys to the existing model cost dictionary
updated_dictionary = update_dictionary(existing_model, value)
litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary)
write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name)
# Invalidate case-insensitive lookup map since model_cost was modified
_invalidate_model_cost_lowercase_map()
verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", model_cost_key, model_cost_key)
verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", write.key, write.key)
if add_to_provider_model_sets and value.get("litellm_provider") == "openrouter":
litellm.openrouter_models.add(_key_str.split("/", 1)[-1])
litellm.openrouter_models.add(str(key).split("/", 1)[-1])
if add_to_provider_model_sets:
litellm.add_known_models(loaded_model_cost)
return model_cost
@ -5732,7 +5895,7 @@ def _rebuild_model_cost_lowercase_map() -> dict[str, str]:
The rebuilt map (guaranteed to be not None).
"""
global _model_cost_lowercase_map
_model_cost_lowercase_map = {k.lower(): k for k in litellm.model_cost}
_model_cost_lowercase_map = {k.lower(): k for k in tuple(litellm.model_cost)}
return _model_cost_lowercase_map
@ -5747,9 +5910,7 @@ def _handle_stale_map_entry_rebuild(
Returns:
The matched key if found after rebuild, None otherwise.
"""
global _model_cost_lowercase_map
_model_cost_lowercase_map = _rebuild_model_cost_lowercase_map()
matched_key: Final = _model_cost_lowercase_map.get(potential_key_lower)
matched_key: Final = _rebuild_model_cost_lowercase_map().get(potential_key_lower)
if matched_key is not None and matched_key in litellm.model_cost:
return matched_key
return None
@ -5767,7 +5928,7 @@ def _handle_new_key_with_scan(
The matched key if found, None otherwise.
"""
global _model_cost_lowercase_map
for key in litellm.model_cost:
for key in tuple(_builtin_cost_map()):
if key.lower() == potential_key_lower:
_model_cost_lowercase_map = _rebuild_model_cost_lowercase_map()
return key
@ -5789,18 +5950,19 @@ def get_model_cost_key(potential_key: str) -> str | None:
called and confirmed not to cause performance issues, add it to the allowed_helpers
list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py
"""
global _model_cost_lowercase_map
# Exact match (O(1))
if potential_key in litellm.model_cost:
return potential_key
# Case-insensitive lookup via map (O(1))
if _model_cost_lowercase_map is None:
_model_cost_lowercase_map = _rebuild_model_cost_lowercase_map()
# Case-insensitive lookup via map (O(1)); bound locally because another thread's
# registration can invalidate the module-level map at any point
cached_lowercase_map: Final = _model_cost_lowercase_map
lowercase_map: Final = (
cached_lowercase_map if cached_lowercase_map is not None else _rebuild_model_cost_lowercase_map()
)
potential_key_lower: Final = potential_key.lower()
matched_key = _model_cost_lowercase_map.get(potential_key_lower)
matched_key = lowercase_map.get(potential_key_lower)
# Verify key exists (O(1) - handles model_cost.pop() case)
if matched_key is not None and matched_key in litellm.model_cost:

433
tests/integration/pricing/test_price_data_reload_prices.py Normal file
View file

@ -0,0 +1,433 @@
import itertools
import json
import re
import threading
import uuid
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
_ROW: Final = "anthropic.claude-audit-reload-v1:0"
_GEO: Final = f"au.{_ROW}"
_FINETUNE: Final = "ft:gpt-audit-reload-2025-01-01"
_CLAUDE: Final = "claude-audit-reload-sonnet"
_MINI: Final = "gpt-audit-reload-mini"
_USAGE: Final = re.compile(r"bill (\d+) (\d+) (\d+)")
_CUSTOM: Final = MappingProxyType({"input_cost_per_token": 1.1e-06, "output_cost_per_token": 2.3e-06})
_COST_MAP_LIMITS: Final = MappingProxyType({"MODEL_COST_MAP_MIN_MODEL_COUNT": "1", "MODEL_COST_MAP_MAX_SHRINK_RATIO": "0"})
_PUBLISHED_SCALE: Final = 2.0
_CATALOG: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType(
{
_ROW: {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token_above_200k_tokens": 6e-06,
"output_cost_per_token_above_200k_tokens": 2.25e-05,
},
_GEO: {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
},
_FINETUNE: {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 3.75e-06,
"output_cost_per_token": 1.5e-05,
"cache_read_input_token_cost": 1.875e-06,
},
_CLAUDE: {
"litellm_provider": "anthropic",
"mode": "chat",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"cache_read_input_token_cost": 3e-07,
},
_MINI: {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 5e-07,
"output_cost_per_token": 2e-06,
},
}
)
@dataclass(frozen=True, slots=True)
class _Usage:
prompt: int
cached: int
completion: int
def text(self) -> str:
return f"bill {self.prompt} {self.cached} {self.completion}"
@dataclass(frozen=True, slots=True)
class _Rates:
input: float
output: float
cache_read: float
def cost(self, usage: _Usage) -> float:
return (
(usage.prompt - usage.cached) * self.input + usage.cached * self.cache_read + usage.completion * self.output
)
@dataclass(frozen=True, slots=True)
class _Deployment:
name: str
model: str
prices: Mapping[str, float]
usage: _Usage
rates: Callable[[float], _Rates | None]
@dataclass(frozen=True, slots=True)
class _Charge:
request_id: str
expected: float
def _catalog(scale: float) -> dict[str, dict[str, JsonValue]]:
return {
model: {name: value * scale if isinstance(value, float) else value for name, value in row.items()}
for model, row in _CATALOG.items()
}
def _rate(model: str, field: str, scale: float) -> float:
rate: Final = _CATALOG[model][field]
assert isinstance(rate, float), rate
return rate * scale
def _catalog_rates(model: str) -> Callable[[float], _Rates]:
return lambda scale: _Rates(
input=_rate(model, "input_cost_per_token", scale),
output=_rate(model, "output_cost_per_token", scale),
cache_read=_rate(model, "cache_read_input_token_cost", scale)
if "cache_read_input_token_cost" in _CATALOG[model]
else _rate(model, "input_cost_per_token", scale),
)
def _custom_rates(_scale: float) -> _Rates:
return _Rates(
input=_CUSTOM["input_cost_per_token"],
output=_CUSTOM["output_cost_per_token"],
cache_read=_CUSTOM["input_cost_per_token"],
)
def _unpriced(_scale: float) -> None:
return None
def _expected_cost(rates: _Rates | None, usage: _Usage) -> float:
return 0.0 if rates is None else rates.cost(usage)
def _expected_listing(rates: _Rates | None) -> tuple[object, object]:
if rates is None:
return None, None
return pytest.approx(rates.input), pytest.approx(rates.output)
def _connection(model: str, wire_url: str) -> dict[str, JsonValue]:
if model.startswith("bedrock/"):
return {"api_base": wire_url, "api_key": "synthetic-bedrock-bearer", "aws_region_name": "ap-northeast-1"}
if model.startswith("azure/"):
return {"api_base": wire_url, "api_key": "synthetic-azure-key", "api_version": "2025-01-01-preview"}
return {"api_base": f"{wire_url}/v1", "api_key": "synthetic-openai-key"}
def _config(directory: Path, deployments: tuple[_Deployment, ...], wire_url: str) -> Path:
config: Final = directory / "price_data_reload_config.yaml"
config.write_text(
json.dumps(
{
"model_list": [
{
"model_name": deployment.name,
"litellm_params": {
"model": deployment.model,
**_connection(deployment.model, wire_url),
**deployment.prices,
},
}
for deployment in deployments
],
"general_settings": {
"master_key": "os.environ/LITELLM_MASTER_KEY",
"database_url": "os.environ/DATABASE_URL",
},
"router_settings": {"disable_cooldowns": True},
}
)
)
return config
def _usage_in(request: Request) -> _Usage:
found: Final = _USAGE.search(request.body.decode())
assert found is not None, request
return _Usage(prompt=int(found.group(1)), cached=int(found.group(2)), completion=int(found.group(3)))
def _reply(request: Request) -> Reply:
usage: Final = _usage_in(request)
if "/converse" in request.target:
return Reply(
body=json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "priced"}]}},
"stopReason": "end_turn",
"usage": {
"inputTokens": usage.prompt,
"outputTokens": usage.completion,
"totalTokens": usage.prompt + usage.completion,
},
"metrics": {"latencyMs": 1},
}
).encode()
)
if "/invoke" in request.target:
return Reply(
body=json.dumps(
{
"id": f"msg_{uuid.uuid4().hex}",
"type": "message",
"role": "assistant",
"model": "claude-scripted",
"content": [{"type": "text", "text": "priced"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": usage.prompt, "output_tokens": usage.completion},
}
).encode()
)
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "scripted",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "priced"}, "finish_reason": "stop"}
],
"usage": {
"prompt_tokens": usage.prompt,
"completion_tokens": usage.completion,
"total_tokens": usage.prompt + usage.completion,
"prompt_tokens_details": {"cached_tokens": usage.cached},
},
}
).encode()
)
def _responder(published: threading.Event) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.target == "/model_prices.json":
return Reply(body=json.dumps(_catalog(_PUBLISHED_SCALE if published.is_set() else 1.0)).encode())
if request.method == "GET":
return Reply(body=json.dumps({"object": "list", "data": []}).encode())
return _reply(request)
return respond
def _listed_rates(proxy: Gateway, name: str) -> tuple[JsonValue, JsonValue]:
entries: Final = proxy.get("/model/info")["data"]
assert isinstance(entries, list), entries
info: Final = next(
object_value(object_value(entry)["model_info"])
for entry in entries
if object_value(entry)["model_name"] == name
)
return info.get("input_cost_per_token"), info.get("output_cost_per_token")
def _billed(client: httpx.Client, key: str, deployment: _Deployment) -> tuple[str, float | None]:
response: Final = client.post(
"/v1/chat/completions",
json={"model": deployment.name, "messages": [{"role": "user", "content": deployment.usage.text()}]},
headers={"Authorization": f"Bearer {key}"},
)
assert response.status_code == 200, response.text
reported: Final = response.headers.get("x-litellm-response-cost")
return string_value(object_value(response.json())["id"]), None if reported is None else float(reported)
def _charged_as_listed(proxy: Gateway, deployment: _Deployment, scale: float) -> _Charge:
rates: Final = deployment.rates(scale)
assert _listed_rates(proxy, deployment.name) == _expected_listing(rates), deployment.name
request_id, charged = _billed(proxy.client, proxy.key, deployment)
expected: Final = _expected_cost(rates, deployment.usage)
assert charged == (None if rates is None else pytest.approx(expected, rel=1e-6)), deployment.name
return _Charge(request_id=request_id, expected=expected)
def _spend(request_id: str) -> float:
rows: Final = eventually(
lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)),
lambda values: len(values) == 1,
seconds=70,
)
spend: Final = rows[0]["spend"]
assert isinstance(spend, (int, float)), rows
return float(spend)
def _proxy_environment(wire_url: str) -> dict[str, str]:
return {**_COST_MAP_LIMITS, "LITELLM_MODEL_COST_MAP_URL": f"{wire_url}/model_prices.json"}
def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments: tuple[_Deployment, ...]) -> None:
published: Final = threading.Event()
with (
wire_server(_responder(published)) as wire,
owned_proxy(
gateway,
tmp_path,
_proxy_environment(wire.url),
config=_config(tmp_path, deployments, wire.url),
remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",),
) as proxy,
):
before: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in deployments)
published.set()
reload: Final = proxy.request("POST", "/reload/model_cost_map")
assert reload.status_code == 200, reload.text
after: Final = tuple(_charged_as_listed(proxy, deployment, _PUBLISHED_SCALE) for deployment in deployments)
for charge in (*before, *after):
assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6)
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
def test_price_data_reload_reprices_catalog_deployments_and_keeps_custom_prices(
gateway: Gateway, tmp_path: Path
) -> None:
_assert_priced_across_a_reload(
gateway,
tmp_path,
(
_Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)),
_Deployment("mini-custom", f"openai/{_MINI}", _CUSTOM, _Usage(1000, 0, 500), _custom_rates),
),
)
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
@pytest.mark.parametrize(
"deployment",
(
_Deployment(
"region-unpriced", f"bedrock/ap-northeast-1.{_ROW}", {}, _Usage(1000, 0, 500), _catalog_rates(_ROW)
),
_Deployment(
"finetune-unpriced",
f"openai/{_FINETUNE}:org:suffix:y",
{},
_Usage(2000, 1000, 1000),
_catalog_rates(_FINETUNE),
),
_Deployment("finetune-custom", f"openai/{_FINETUNE}:org:suffix:z", _CUSTOM, _Usage(2000, 0, 1000), _custom_rates),
_Deployment("cross-provider-unpriced", f"azure/{_CLAUDE}", {}, _Usage(2000, 1000, 1000), _unpriced),
_Deployment(
"region-custom-above-threshold",
f"bedrock/ap-northeast-1.{_ROW}",
_CUSTOM,
_Usage(250_000, 0, 1000),
_custom_rates,
),
),
ids=lambda deployment: deployment.name,
)
def test_deployment_bills_the_catalog_row_its_registration_resolves_to_across_a_reload(
gateway: Gateway, tmp_path: Path, deployment: _Deployment
) -> None:
_assert_priced_across_a_reload(gateway, tmp_path, (deployment,))
def _storm_deployment(index: int) -> _Deployment:
prices: Final = MappingProxyType(
{"input_cost_per_token": (index + 1) * 1e-07, "output_cost_per_token": (index + 1) * 3e-07}
)
return _Deployment(
f"storm-{index}",
f"openai/gpt-audit-reload-storm-{index}",
prices,
_Usage(100, 0, 50),
lambda _scale: _Rates(
input=prices["input_cost_per_token"],
output=prices["output_cost_per_token"],
cache_read=prices["input_cost_per_token"],
),
)
def _storm_traffic(
base_url: str, key: str, deployments: tuple[_Deployment, ...], offset: int, finished: threading.Event
) -> tuple[tuple[float, float], ...]:
rotation: Final = (
deployments[(offset + turn) % len(deployments)] for turn, _ in enumerate(iter(finished.is_set, True))
)
with httpx.Client(base_url=base_url, timeout=60, trust_env=False) as client:
return tuple(
(_billed(client, key, deployment)[1], _expected_cost(deployment.rates(1.0), deployment.usage))
for deployment in rotation
)
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
def test_price_data_reloads_under_custom_priced_traffic_keep_every_deployment_priced(
gateway: Gateway, tmp_path: Path
) -> None:
storm: Final = tuple(_storm_deployment(index) for index in range(24))
geo: Final = _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO))
finished: Final = threading.Event()
with (
wire_server(_responder(threading.Event())) as wire,
owned_proxy(
gateway,
tmp_path,
_proxy_environment(wire.url),
config=_config(tmp_path, (geo, *storm), wire.url),
remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",),
) as proxy,
ThreadPoolExecutor(max_workers=8) as pool,
):
traffic: Final = tuple(
pool.submit(_storm_traffic, str(proxy.client.base_url), proxy.key, storm, offset, finished)
for offset in range(8)
)
reloads: Final = tuple(proxy.request("POST", "/reload/model_cost_map") for _ in range(40))
finished.set()
charges: Final = tuple(itertools.chain.from_iterable(calls.result() for calls in traffic))
assert tuple(reload.status_code for reload in reloads) == (200,) * 40, tuple(
reload.text for reload in reloads if reload.status_code != 200
)
assert len(charges) >= 8
assert tuple(charged for charged, _ in charges) == pytest.approx(tuple(expected for _, expected in charges))
settled: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in (geo, *storm))
for charge in settled:
assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6)

608
tests/unit/test_register_model_custom_pricing.py
View file

@ -10,13 +10,16 @@ calculations for DB-sourced models with prompt caching pricing.
"""
import copy
import itertools
import os
from typing import Final
from collections.abc import Callable, Mapping
from typing import Final, NoReturn
import pytest
import litellm
from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases, adopt_model_cost_map
from litellm.main import _build_custom_pricing_entry
from litellm.utils import _invalidate_model_cost_lowercase_map
@ -796,6 +799,609 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key():
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize("route", ["", "converse/", "invoke/"])
def test_register_model_prices_a_geo_alias_from_its_builtin_when_the_lookup_fails(
monkeypatch: pytest.MonkeyPatch, route: str
) -> None:
from litellm.types.utils import ModelResponse, Usage
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
_invalidate_model_cost_lowercase_map()
canonical_key: Final = "au.anthropic.claude-sonnet-5"
alias: Final = f"bedrock/{route}{canonical_key}"
canonical: Final = litellm.model_cost[canonical_key]
def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration)
litellm.register_model(
{alias: {"litellm_provider": "bedrock", "mode": "chat", "supports_vision": True}},
persist_across_reloads=False,
)
usage: Final = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100)
alias_cost: Final = litellm.completion_cost(
completion_response=ModelResponse(model=alias, usage=usage), model=alias, custom_llm_provider="bedrock"
)
canonical_cost: Final = litellm.completion_cost(
completion_response=ModelResponse(model=canonical_key, usage=usage),
model=canonical_key,
custom_llm_provider="bedrock",
)
try:
assert alias_cost == canonical_cost > 0
assert litellm.get_model_info(alias)["input_cost_per_token"] == canonical["input_cost_per_token"]
assert litellm.get_model_info(alias)["supports_vision"] is True
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None:
geo_key: Final = "au.anthropic.lit-5853-geo-model"
geo_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
"cache_read_input_token_cost": 3.3e-07,
}
sibling_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"cache_read_input_token_cost": 3e-07,
}
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
geo_key: dict(geo_row),
"anthropic.lit-5853-geo-model": sibling_row,
},
)
_invalidate_model_cost_lowercase_map()
def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration)
litellm.register_model({geo_key: {"mode": "chat"}}, persist_across_reloads=False)
try:
for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"):
assert litellm.model_cost[geo_key][field] == geo_row[field]
assert litellm.get_model_info(geo_key)[field] == geo_row[field]
assert litellm.model_cost[geo_key]["litellm_provider"] == geo_row["litellm_provider"]
finally:
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize("route", ["", "converse/", "invoke/"])
def test_register_model_leaves_no_shadow_key_when_the_lookup_trips_once(
monkeypatch: pytest.MonkeyPatch, route: str
) -> None:
catalog_key: Final = "au.anthropic.lit-5853-geo-model"
alias: Final = f"bedrock/{route}{catalog_key}"
repriced_input_cost: Final = 9e-06
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
catalog_key: {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
},
},
)
_invalidate_model_cost_lowercase_map()
real_lookup: Final = litellm.utils.get_model_info
lookups: Final = itertools.count()
def lookup_tripped_once_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> object:
if next(lookups) == 0:
raise RuntimeError("dictionary changed size during iteration")
return real_lookup(model=model, custom_llm_provider=custom_llm_provider)
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_once_by_a_concurrent_registration)
litellm.register_model({alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False)
litellm.register_model({catalog_key: {"input_cost_per_token": repriced_input_cost}}, persist_across_reloads=False)
try:
assert litellm.get_model_info(alias)["key"] == catalog_key
assert litellm.get_model_info(alias)["input_cost_per_token"] == repriced_input_cost
finally:
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize("key_prefix", ["", "bedrock/", "bedrock/converse/"])
def test_register_model_bills_a_catalog_rows_tier_rates_the_same_when_the_lookup_fails(
monkeypatch: pytest.MonkeyPatch, key_prefix: str
) -> None:
from litellm.types.utils import ModelResponse, Usage
catalog_key: Final = "au.anthropic.lit-5853-tiered-geo-model"
registered_key: Final = f"{key_prefix}{catalog_key}"
catalog_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
"input_cost_per_token_above_200k_tokens": 6.6e-06,
"output_cost_per_token_above_200k_tokens": 2.475e-05,
}
builtin_cost_map: Final = litellm.get_model_cost_map(url="")
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
def long_prompt_cost_after_registering_with(lookup: Callable[..., object]) -> float:
monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)})
_invalidate_model_cost_lowercase_map()
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup)
litellm.register_model(
{
registered_key: {
"litellm_provider": "bedrock",
"input_cost_per_token": 1e-06,
"output_cost_per_token": 2e-06,
}
},
persist_across_reloads=False,
)
return litellm.completion_cost(
completion_response=ModelResponse(
model=registered_key,
usage=Usage(prompt_tokens=250_000, completion_tokens=1_000, total_tokens=251_000),
),
model=registered_key,
custom_llm_provider="bedrock",
)
try:
assert long_prompt_cost_after_registering_with(lookup_tripped_by_a_concurrent_registration) == pytest.approx(
long_prompt_cost_after_registering_with(litellm.utils.get_model_info)
)
finally:
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize("lookup_trips", [False, True])
def test_register_model_prices_a_provider_less_azure_ai_key_without_repricing_the_openai_row(
monkeypatch: pytest.MonkeyPatch, lookup_trips: bool
) -> None:
openai_key: Final = "lit-5853-openai-model"
azure_ai_key: Final = f"azure_ai/{openai_key}"
openai_row: Final = {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 5e-06,
"output_cost_per_token": 1.5e-05,
}
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setenv("AZURE_AI_API_BASE", "https://lit-5853.openai.azure.com")
monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), openai_key: dict(openai_row)})
monkeypatch.setattr(
litellm, "open_ai_chat_completion_models", {*litellm.open_ai_chat_completion_models, openai_key}
)
_invalidate_model_cost_lowercase_map()
def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn:
raise RuntimeError("dictionary changed size during iteration")
with pytest.MonkeyPatch.context() as lookup_patch:
if lookup_trips:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration)
litellm.register_model(
{azure_ai_key: {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}},
persist_across_reloads=False,
)
try:
assert litellm.model_cost[openai_key] == openai_row
assert litellm.cost_per_token(model=azure_ai_key, prompt_tokens=1000, completion_tokens=1000) == pytest.approx(
(1e-03, 2e-03)
)
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monkeypatch: pytest.MonkeyPatch) -> None:
catalog_key: Final = "au.anthropic.lit-5853-reloaded-model"
deployment_key: Final = f"bedrock/{catalog_key}"
catalog_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
}
reloaded_input_cost: Final = 6.6e-06
builtin_cost_map: Final = litellm.get_model_cost_map(url="")
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)})
_invalidate_model_cost_lowercase_map()
reloaded_cost_map: Final = {
**builtin_cost_map,
catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost},
}
real_lookup: Final = litellm.utils.get_model_info
def lookup_racing_a_price_data_reload(model: str, custom_llm_provider: str | None = None) -> object:
info: Final = real_lookup(model=model, custom_llm_provider=custom_llm_provider)
if litellm.model_cost is not reloaded_cost_map:
adopt_model_cost_map(reloaded_cost_map)
return info
with pytest.MonkeyPatch.context() as lookup_patch:
lookup_patch.setattr(litellm.utils, "get_model_info", lookup_racing_a_price_data_reload)
litellm.register_model({deployment_key: {"litellm_provider": "bedrock"}}, persist_across_reloads=False)
try:
assert litellm.model_cost[catalog_key]["input_cost_per_token"] == reloaded_input_cost
assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_keeps_its_price_when_a_reload_adopts_a_map_as_it_writes(
monkeypatch: pytest.MonkeyPatch,
) -> None:
catalog_key: Final = "au.anthropic.lit-5853-reloaded-as-written-model"
deployment_key: Final = f"bedrock/{catalog_key}"
catalog_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3.3e-06,
"output_cost_per_token": 1.65e-05,
}
reloaded_input_cost: Final = 6.6e-06
registered_output_cost: Final = 2e-05
builtin_cost_map: Final = litellm.get_model_cost_map(url="")
reloaded_cost_map: Final = {
**builtin_cost_map,
catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost},
}
class CostMapAReloadReplacesAsItIsWritten(dict[str, dict[str, object]]):
def update(self, other: Mapping[str, dict[str, object]], /) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # registration writes pass one mapping
if litellm.model_cost is self:
adopt_model_cost_map(reloaded_cost_map)
super().update(other)
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
CostMapAReloadReplacesAsItIsWritten({**builtin_cost_map, catalog_key: dict(catalog_row)}),
)
_invalidate_model_cost_lowercase_map()
litellm.register_model(
{deployment_key: {"litellm_provider": "bedrock", "output_cost_per_token": registered_output_cost}},
persist_across_reloads=False,
)
try:
assert litellm.model_cost is reloaded_cost_map
assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost
assert litellm.get_model_info(deployment_key)["output_cost_per_token"] == registered_output_cost
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_leaves_the_entry_a_lookup_is_walking_intact(monkeypatch: pytest.MonkeyPatch) -> None:
catalog_key: Final = "lit-5853-live-entry-model"
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
catalog_key: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 2e-06},
},
)
_invalidate_model_cost_lowercase_map()
walk: Final = iter(litellm.model_cost[catalog_key].items())
first_field: Final = next(walk)
try:
litellm.register_model(
{catalog_key: {"litellm_provider": "openai", "input_cost_per_token": 1e-06, "supports_vision": True}},
persist_across_reloads=False,
)
assert dict((first_field, *walk)) == {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 2e-06,
}
assert litellm.get_model_info(catalog_key)["input_cost_per_token"] == 1e-06
assert litellm.get_model_info(catalog_key)["supports_vision"] is True
finally:
_invalidate_model_cost_lowercase_map()
_CANONICAL_KEY: Final = "lit-5853-canonical-model"
_ALIAS_KEY: Final = "lit-5853-alias-model"
def _install_aliased_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
_expand_model_aliases(
{
**litellm.get_model_cost_map(url=""),
_CANONICAL_KEY: {
"aliases": [_ALIAS_KEY],
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 2e-06,
"output_cost_per_token": 8e-06,
},
}
),
)
_invalidate_model_cost_lowercase_map()
def _token_prices(name: str) -> tuple[float | None, float | None]:
info: Final = litellm.get_model_info(name)
return info["input_cost_per_token"], info["output_cost_per_token"]
@pytest.mark.parametrize("registered_key", [_CANONICAL_KEY, _ALIAS_KEY])
def test_register_model_prices_every_alias_of_the_registered_entry(
monkeypatch: pytest.MonkeyPatch, registered_key: str
) -> None:
_install_aliased_cost_map(monkeypatch)
try:
litellm.register_model(
{registered_key: {"litellm_provider": "openai", "input_cost_per_token": 1e-06}},
persist_across_reloads=False,
)
litellm.register_model(
{registered_key: {"litellm_provider": "openai", "output_cost_per_token": 4e-06}},
persist_across_reloads=False,
)
assert {name: _token_prices(name) for name in (_CANONICAL_KEY, _ALIAS_KEY)} == {
_CANONICAL_KEY: (1e-06, 4e-06),
_ALIAS_KEY: (1e-06, 4e-06),
}
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_leaves_an_alias_repointed_at_its_own_entry(monkeypatch: pytest.MonkeyPatch) -> None:
_install_aliased_cost_map(monkeypatch)
litellm.model_cost[_ALIAS_KEY] = {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 5e-06,
"output_cost_per_token": 6e-06,
}
try:
litellm.register_model(
{_CANONICAL_KEY: {"litellm_provider": "openai", "input_cost_per_token": 1e-06}},
persist_across_reloads=False,
)
assert {name: _token_prices(name) for name in (_CANONICAL_KEY, _ALIAS_KEY)} == {
_CANONICAL_KEY: (1e-06, 8e-06),
_ALIAS_KEY: (5e-06, 6e-06),
}
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_inherits_builtin_token_pricing_for_unmapped_key(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
_invalidate_model_cost_lowercase_map()
builtin_key: Final = "us.anthropic.claude-sonnet-4-6"
registered_key: Final = f"bedrock/bedrock/bedrock/{builtin_key}"
builtin: Final = litellm.model_cost[builtin_key]
try:
litellm.register_model(
{registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False
)
info: Final = litellm.get_model_info(registered_key)
assert info["input_cost_per_token"] == builtin["input_cost_per_token"] > 0
assert info["output_cost_per_token"] == builtin["output_cost_per_token"] > 0
assert info["cache_read_input_token_cost"] == builtin["cache_read_input_token_cost"]
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_keeps_the_capabilities_an_unregistered_region_alias_resolves_to(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.litellm_core_utils.fallback_generalizations import (
get_fallback_generalization_rules,
set_fallback_generalizations,
)
sibling_key: Final = "anthropic.lit-5853-flagless-model"
region_alias: Final = f"bedrock/ap-northeast-1.{sibling_key}"
sibling_row: Final = {
"litellm_provider": "bedrock_converse",
"mode": "chat",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
}
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), sibling_key: sibling_row})
_invalidate_model_cost_lowercase_map()
previous_rules: Final = list(get_fallback_generalization_rules())
set_fallback_generalizations(
[
{
"name": "lit-5853-flagless",
"pattern": "lit-5853-flagless",
"model_info": {"supports_system_messages": True, "supports_parallel_function_calling": True},
}
]
)
try:
unregistered: Final = litellm.get_model_info(region_alias)
litellm.register_model(
{region_alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False
)
registered: Final = litellm.get_model_info(region_alias)
assert registered["supports_system_messages"] is unregistered["supports_system_messages"] is True
assert (
registered["supports_parallel_function_calling"]
is unregistered["supports_parallel_function_calling"]
is True
)
assert registered["input_cost_per_token"] == sibling_row["input_cost_per_token"]
finally:
set_fallback_generalizations(previous_rules)
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize(
("builtin_prices", "custom_prices", "usage", "expected_cost"),
[
pytest.param(
{
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"input_cost_per_token_above_200k_tokens": 6e-06,
"output_cost_per_token_above_200k_tokens": 2.25e-05,
},
{"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06},
{"prompt_tokens": 250_000, "completion_tokens": 1_000},
250_000 * 1e-06 + 1_000 * 2e-06,
id="above-200k-tier",
),
pytest.param(
{
"tiered_pricing": [
{"input_cost_per_token": 5e-08, "output_cost_per_token": 4e-07, "range": [0, 256_000]},
{"input_cost_per_token": 2.5e-07, "output_cost_per_token": 2e-06, "range": [256_000, 1_000_000]},
]
},
{"input_cost_per_token": 1e-06},
{"prompt_tokens": 300_000, "completion_tokens": 1_000},
300_000 * 1e-06 + 1_000 * 2e-06,
id="tiered-pricing-input-only",
),
pytest.param(
{"input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05, "output_cost_per_reasoning_token": 3e-05},
{"output_cost_per_token": 2e-06},
{"prompt_tokens": 100, "completion_tokens": 1_000, "reasoning_tokens": 600},
100 * 3e-06 + 1_000 * 2e-06,
id="reasoning-rate-output-only",
),
pytest.param(
{
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"off_peak_pricing": {
"input_cost_per_token": 1e-06,
"output_cost_per_token": 5e-06,
"hours_utc": "00:00-00:00",
},
},
{"output_cost_per_token": 2e-06},
{"prompt_tokens": 100, "completion_tokens": 1_000},
100 * 1e-06 + 1_000 * 2e-06,
id="off-peak-output-only",
),
],
)
def test_register_model_keeps_custom_flat_prices_over_inherited_rate_variants(
monkeypatch: pytest.MonkeyPatch,
builtin_prices: dict[str, object],
custom_prices: dict[str, float],
usage: dict[str, int],
expected_cost: float,
) -> None:
from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, Usage
builtin_key: Final = "anthropic.lit-5853-tiered-model-v1:0"
alias: Final = f"bedrock/ap-northeast-1.{builtin_key}"
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
builtin_key: {"litellm_provider": "bedrock", "mode": "chat", **builtin_prices},
},
)
_invalidate_model_cost_lowercase_map()
try:
litellm.register_model(
{alias: {"litellm_provider": "bedrock", **custom_prices}},
persist_across_reloads=False,
)
response_usage: Final = Usage(
prompt_tokens=usage["prompt_tokens"],
completion_tokens=usage["completion_tokens"],
total_tokens=usage["prompt_tokens"] + usage["completion_tokens"],
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=usage.get("reasoning_tokens", 0)),
)
cost: Final = litellm.completion_cost(
completion_response=ModelResponse(model=alias, usage=response_usage),
model=alias,
custom_llm_provider="bedrock",
)
assert cost == pytest.approx(expected_cost)
finally:
_invalidate_model_cost_lowercase_map()
def test_register_model_never_inherits_pricing_across_providers(monkeypatch: pytest.MonkeyPatch) -> None:
openai_key: Final = "lit-5853-openai-only-model"
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.get_model_cost_map(url=""),
openai_key: {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 2.5e-06,
"output_cost_per_token": 1e-05,
"cache_read_input_token_cost": 1.25e-06,
},
},
)
_invalidate_model_cost_lowercase_map()
registered_key: Final = f"bedrock/{openai_key}"
try:
litellm.register_model(
{registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False
)
registered: Final = litellm.model_cost[registered_key]
assert "cache_read_input_token_cost" not in registered
assert "input_cost_per_token" not in registered
assert "output_cost_per_token" not in registered
finally:
_invalidate_model_cost_lowercase_map()
def test_update_dictionary_merges_nested_dicts_without_aliasing():
"""A nested dict must be merged copy-on-write: the pre-existing nested dict
object stays untouched, and the caller's incoming nested dict is never

75
tests/unit/test_router_model_cost_isolation.py
View file

@ -9,9 +9,14 @@ should still use the built-in pricing.
import asyncio
import copy
import itertools
import logging
import os
import queue
import re
import sys
import threading
from types import MappingProxyType
from typing import Final
from unittest.mock import Mock, patch
@ -3241,6 +3246,76 @@ def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info(
assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price
def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
catalog: Final = litellm.get_model_cost_map(url="")
monkeypatch.setattr(litellm, "model_cost", dict(catalog))
_invalidate_model_cost_lowercase_map()
backends: Final = tuple(
f"bedrock/au.anthropic.{model}"
for model in ("claude-sonnet-5", "claude-opus-5", "claude-haiku-4-5-20251001-v1:0", "claude-sonnet-4-6")
)
canonical_prices: Final = MappingProxyType(
{backend: catalog[backend.removeprefix("bedrock/")]["input_cost_per_token"] for backend in backends}
)
assert all(price > 0 for price in canonical_prices.values())
router: Final = Router(
model_list=[
{
"model_name": backend.rsplit(".", 1)[-1],
"litellm_params": {"model": backend, "aws_region_name": "ap-southeast-2"},
"model_info": {"id": f"lit-5853-dep-{index}"},
}
for index, backend in enumerate(backends)
]
)
stop: Final = threading.Event()
writer_outcome: Final[queue.SimpleQueue[int | Exception]] = queue.SimpleQueue()
def register_until_stopped() -> None:
try:
for i in itertools.count():
if stop.is_set():
writer_outcome.put(i)
return
litellm.register_model(
{
f"lit-5853-writer-{i}": {
"litellm_provider": "lit-5853-writer",
"mode": "chat",
"input_cost_per_token": 1e-07,
"output_cost_per_token": 4e-07,
"cache_read_input_token_cost": 1e-08,
}
},
persist_across_reloads=False,
)
except Exception as failure:
writer_outcome.put(failure)
writer: Final = threading.Thread(target=register_until_stopped, daemon=True)
saved_switch_interval: Final = sys.getswitchinterval()
sys.setswitchinterval(1e-4)
writer.start()
try:
for _ in range(10):
_simulate_price_data_reload_with_provider_sets(monkeypatch, dict(catalog))
finally:
stop.set()
writer.join()
sys.setswitchinterval(saved_switch_interval)
registrations_landed: Final = writer_outcome.get(timeout=5)
try:
assert isinstance(registrations_landed, int) and registrations_landed > 0, registrations_landed
assert tuple(backend for backend in backends if backend in litellm.model_cost) == ()
for backend in backends:
assert litellm.get_model_info(backend)["input_cost_per_token"] == canonical_prices[backend]
assert tuple(router.get_model_names()) == tuple(backend.rsplit(".", 1)[-1] for backend in backends)
finally:
_invalidate_model_cost_lowercase_map()
_COLLIDING_BUILTIN_KEY: Final = "baseten/zai-org/GLM-5.2"
_COLLIDING_SHARED_OPENAI_KEY: Final = "openai/zai-org/GLM-5.2"
_COLLIDING_MODEL_INFO: Final = {

119
tests/unit/test_utils.py
View file

@ -14,7 +14,7 @@ from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import Future, ThreadPoolExecutor
from datetime import datetime, timedelta, timezone
from pathlib import Path, PurePath
from typing import Final, cast
from typing import Final, SupportsIndex, cast
from unittest import mock
from unittest.mock import AsyncMock, MagicMock, Mock, patch
@ -4887,6 +4887,104 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey
_invalidate_model_cost_lowercase_map()
def test_model_cost_walks_survive_a_registration_landing_mid_walk(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.utils import _invalidate_model_cost_lowercase_map, _rebuild_model_cost_lowercase_map
class _KeyRegisteredMidWalk(str):
def lower(self) -> str:
litellm.model_cost.setdefault("lit-5853-registered-during-the-rebuild", {"mode": "chat"})
return super().lower()
class _EntryRegisteredMidWalk(dict):
def get(self, field: str, default: object = None) -> object:
litellm.model_cost.setdefault("lit-5853-registered-during-add-known-models", {"mode": "chat"})
return super().get(field, default)
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
litellm.model_cost[_KeyRegisteredMidWalk("LIT-5853-WALKED-LAST")] = _EntryRegisteredMidWalk(mode="chat")
try:
assert _rebuild_model_cost_lowercase_map()["lit-5853-walked-last"] == "LIT-5853-WALKED-LAST"
litellm.add_known_models(model_cost_map=litellm.model_cost)
assert "lit-5853-registered-during-the-rebuild" in litellm.model_cost
assert "lit-5853-registered-during-add-known-models" in litellm.model_cost
finally:
_invalidate_model_cost_lowercase_map()
class _KeyWhoseLoweringLosesTheRace(str):
"""Lowering the key is the one Python-level step between the lookup map's rebuild and
its read, so another thread's registration invalidating the map lands right there."""
def lower(self) -> str:
from litellm.utils import _invalidate_model_cost_lowercase_map
_invalidate_model_cost_lowercase_map()
return super().lower()
def test_case_insensitive_model_cost_lookup_survives_an_invalidation_mid_lookup() -> None:
from litellm.utils import _get_model_cost_key, _invalidate_model_cost_lowercase_map
_invalidate_model_cost_lowercase_map()
try:
assert _get_model_cost_key(_KeyWhoseLoweringLosesTheRace("GPT-4O")) == "gpt-4o"
finally:
_invalidate_model_cost_lowercase_map()
def _register_a_model_mid_walk() -> None:
litellm.model_cost.setdefault("lit-5853-registered-mid-walk", {"mode": "chat"})
class _KeyReadDuringARegistration(str):
def startswith(
self,
prefix: str | tuple[str, ...],
start: SupportsIndex | None = None,
end: SupportsIndex | None = None,
) -> bool:
_register_a_model_mid_walk()
return super().startswith(prefix, start, end)
def lower(self) -> str:
_register_a_model_mid_walk()
return super().lower()
class _EntryReadDuringARegistration(dict[str, object]):
def __contains__(self, key: object) -> bool:
_register_a_model_mid_walk()
return super().__contains__(key)
@pytest.mark.parametrize("walk", ["fal_ai", "fireworks_ai", "ollama", "gpt-5", "case-insensitive scan"])
def test_request_path_cost_map_walks_survive_a_registration_mid_walk(monkeypatch: pytest.MonkeyPatch, walk: str) -> None:
from litellm.llms.fal_ai.cost_calculator import _keyed_rows
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
from litellm.llms.ollama.common_utils import OllamaModelInfo
from litellm.llms.openai.chat.gpt_5_transformation import _catalogue_declares_default_effort
from litellm.utils import _handle_new_key_with_scan, _invalidate_model_cost_lowercase_map
walks: Final[Mapping[str, Callable[[], object]]] = {
"fal_ai": lambda: _keyed_rows("lit-5853-model", "standard"),
"fireworks_ai": FireworksAIConfig._get_fireworks_index,
"ollama": lambda: OllamaModelInfo._is_static_ollama_model("ollama/lit-5853-model"),
"gpt-5": _catalogue_declares_default_effort,
"case-insensitive scan": lambda: _handle_new_key_with_scan("lit-5853-model"),
}
monkeypatch.setattr(
litellm,
"model_cost",
{_KeyReadDuringARegistration("lit-5853-catalog-model"): _EntryReadDuringARegistration(mode="chat")},
)
try:
walks[walk]()
assert "lit-5853-registered-mid-walk" in litellm.model_cost
finally:
_invalidate_model_cost_lowercase_map()
class _JsonCapture(logging.Handler):
def __init__(self):
super().__init__()
@ -7091,6 +7189,25 @@ def test_get_model_info_gemini(monkeypatch):
assert info.get("rpm") is not None, f"{model} does not have rpm"
def test_register_model_without_a_provider_never_inherits_another_providers_entry(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.utils import _invalidate_model_cost_lowercase_map
key: Final = "azure/claude-sonnet-4-6"
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
litellm.model_cost.pop(key, None)
try:
litellm.register_model(
{key: {"input_cost_per_token": 1e-07, "output_cost_per_token": 4e-07}}, persist_across_reloads=False
)
assert litellm.model_cost[key].get("litellm_provider") is None
info: Final = litellm.get_model_info(key)
assert info["litellm_provider"] == "azure"
assert info["input_cost_per_token"] == 1e-07
assert info["output_cost_per_token"] == 4e-07
finally:
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize(
"model, litellm_params, expects_native_config",
[