diff --git a/litellm/__init__.py b/litellm/__init__.py index 959c9fb406d..9cc8965b628 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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": diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index b68db22183a..073e3a3455e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -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 diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index aa3d7942174..2acf1c69594 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -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),) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 8b6946ee562..14ae59d4f85 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -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): diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index b5328b1ba4d..284a6ffea62 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -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 diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 94c5c5d016d..5bf19844287 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -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( diff --git a/litellm/utils.py b/litellm/utils.py index fc3acd3065d..b37827a3b13 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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: diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py new file mode 100644 index 00000000000..bb24130aa6c --- /dev/null +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -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) diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 02c90e0d97b..35fbd28eae7 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -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 diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d8625ff9da3..b353dc86d5f 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -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 = { diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 87f9d58e247..a8e90fded83 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -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", [