diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 845c62d91d6..55aeaaa10ff 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict, assert_never import litellm @@ -125,6 +125,7 @@ from litellm.types.utils import ( TranscriptionUsageTokensObject, Usage, VectorStoreSearchResponse, + is_server_derived_pricing_key, ) from litellm.utils import ( CallTypes, @@ -137,6 +138,8 @@ from litellm.utils import ( TextCompletionResponse, TranscriptionResponse, _cached_get_model_info_helper, + model_info_with_pricing_overrides, + runtime_model_cost_registrations, token_counter, ) @@ -402,6 +405,7 @@ def cost_per_token( ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection custom_model_info: OCRPricing | None = None, + bedrock_model_info: ModelInfo | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -688,6 +692,15 @@ def cost_per_token( elif custom_llm_provider == "anthropic": return anthropic_cost_per_token(model=model, usage=usage_block, service_tier=service_tier) elif custom_llm_provider == "bedrock": + resolved_bedrock_info: Final = bedrock_model_info or _bedrock_pricing_info(model, None, "bedrock", False) + if resolved_bedrock_info is not None: + return generic_cost_per_token( + model=model, + usage=usage_block, + custom_llm_provider="bedrock", + service_tier=service_tier, + model_info=resolved_bedrock_info, + ) return bedrock_cost_per_token(model=model, usage=usage_block, service_tier=service_tier) elif custom_llm_provider == "openai": return openai_cost_per_token( @@ -801,6 +814,55 @@ def _cost_map_entry_prices_anything(entry: Mapping[str, object]) -> bool: ) +_LOGGED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + +def _bedrock_registration_has_prices(entry: Mapping[str, object]) -> bool: + return _cost_map_entry_prices_anything(entry) or any( + value is not None and field.startswith("cache_") and is_server_derived_pricing_key(field) + for field, value in entry.items() + ) + + +def _bedrock_pricing_info( + model: str | None, + invoked_model_id: str | None, + custom_llm_provider: str | None, + explicit_pricing: bool, +) -> ModelInfo | None: + if explicit_pricing or custom_llm_provider != "bedrock" or not model: + return None + registration: Final = next( + ( + (key, entry) + for key, entry in runtime_model_cost_registrations(model, "bedrock") + if _bedrock_registration_has_prices(entry) + ), + None, + ) + profile: Final = _cost_map_model_info(invoked_model_id, "bedrock") if invoked_model_id else None + if profile is None or not _declares_rates(profile): + return _cost_map_model_info(registration[0], None) if registration is not None else None + if registration is None or registration[0] == profile["key"]: + return profile + return model_info_with_pricing_overrides(profile, registration[1]) + + +def _declares_rates(model_info: ModelInfo | None) -> bool: + entry: Final = _raw_cost_map_entry(model_info["key"]) if model_info is not None else None + return entry is not None and _cost_map_entry_prices_anything(entry) + + +def _invoked_model_id_for_cost( + optional_params: object, + logging_obj: LitellmLoggingObject | None, +) -> str | None: + recorded: Final[object] = getattr(logging_obj, "model_call_details", None) + return _get_hidden_str_for_cost_calc(optional_params, "model_id") or _get_hidden_str_for_cost_calc( + recorded, "model_id" + ) + + def _select_model_name_for_cost_calc( model: str | None, completion_response: object | None, @@ -809,6 +871,7 @@ def _select_model_name_for_cost_calc( custom_llm_provider: str | None = None, router_model_id: str | None = None, region_name: str | None = None, + invoked_model_id: str | None = None, ) -> str | None: """ 1. If custom pricing is true, return received model name @@ -829,6 +892,13 @@ def _select_model_name_for_cost_calc( hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None) provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model") explicit_pricing: Final = custom_pricing is True or base_model is not None + configured_model: Final = model or ( + completion_response_model if isinstance(completion_response_model, str) else None + ) + bedrock_model_info: Final = _bedrock_pricing_info( + configured_model, invoked_model_id, custom_llm_provider, explicit_pricing + ) + bedrock_pricing_model: Final = bedrock_model_info["key"] if bedrock_model_info is not None else None priced_from_response: Final = provider_response_model is not None or completion_response_model is not None priced_region: Final = ( _get_hidden_str_for_cost_calc(hidden_params, "region_name") or region_name @@ -846,8 +916,8 @@ def _select_model_name_for_cost_calc( else: return_model = model - elif base_model is not None or provider_response_model is not None: - return_model = base_model if base_model is not None else provider_response_model + elif base_model is not None or bedrock_pricing_model is not None or provider_response_model is not None: + return_model = base_model if base_model is not None else bedrock_pricing_model or provider_response_model elif completion_response_model is None and hidden_params is not None: if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0: @@ -1426,6 +1496,10 @@ def completion_cost( service_tier = _normalize_service_tier(service_tier) explicit_pricing: Final = custom_pricing is True or base_model is not None + invoked_model_id: Final = _invoked_model_id_for_cost( + _LOGGED_LITELLM_PARAMS.validate_python(optional_params) if optional_params is not None else None, + litellm_logging_obj, + ) selected_model: Final = _select_model_name_for_cost_calc( model=model, completion_response=completion_response, @@ -1434,6 +1508,14 @@ def completion_cost( base_model=base_model, router_model_id=router_model_id, region_name=region_name, + invoked_model_id=invoked_model_id, + ) + + bedrock_model_info: Final = _bedrock_pricing_info( + model or _get_response_model(completion_response), + invoked_model_id, + _get_provider_for_cost_calc(model, custom_llm_provider), + explicit_pricing, ) potential_model_names: Final = [ @@ -1808,6 +1890,7 @@ def completion_cost( vertex_location=vertex_location, response=completion_response, custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id), + bedrock_model_info=bedrock_model_info, ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) @@ -1896,6 +1979,7 @@ def completion_cost( data_residency=data_residency, vertex_location=vertex_location, custom_cost_per_token=custom_cost_per_token, + model_info=bedrock_model_info, ) _reasoning_cost = _token_type_breakdown.reasoning_cost _cache_read_cost = _token_type_breakdown.cache_read_cost @@ -2137,7 +2221,16 @@ def pricing_entry_for_cost_calc( custom_llm_provider=custom_llm_provider, router_model_id=router_model_id, region_name=region_name, + invoked_model_id=_invoked_model_id_for_cost(None, litellm_logging_obj), ) + bedrock_model_info: Final = _bedrock_pricing_info( + model or _get_response_model(completion_response), + _invoked_model_id_for_cost(None, litellm_logging_obj), + _get_provider_for_cost_calc(model, custom_llm_provider), + custom_pricing is True or base_model is not None, + ) + if bedrock_model_info is not None: + return bedrock_model_info["key"], bedrock_model_info candidates: Final = (selected_model, _get_response_model(completion_response), model) resolved: Final = next( (info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info), diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index d25b70e2450..8ca79874564 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -275,7 +275,7 @@ def _parse_token_threshold(threshold: str) -> float: def _parse_above_token_threshold(key: str) -> float: - return _parse_token_threshold(key.split("_above_")[1].split("_tokens")[0]) + return _parse_token_threshold(key.rsplit("_above_", 1)[1].split("_tokens")[0]) def _prompt_exceeds_threshold(prompt_tokens: int, threshold: float, inclusive: bool) -> bool: @@ -669,10 +669,12 @@ def _get_token_base_cost( # so that the threshold detection loop only processes standard keys. The # service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key. tier_key_suffix: Final = _get_service_tier_cost_key("", service_tier) + threshold_prefixes: Final = tuple(f"{prefix}_above_" for prefix in _BATCH_RATE_PREFIXES) threshold_keys: Final = [ k for k in model_info - if k.startswith("input_cost_per_token_above_") + if k.startswith(threshold_prefixes) + and "_tokens" in k and ( not k.endswith(_NON_STANDARD_THRESHOLD_SUFFIXES) or (tier_key_suffix != "" and k.endswith(tier_key_suffix)) ) @@ -680,12 +682,12 @@ def _get_token_base_cost( # Only sort the threshold keys (typically 1-2 keys instead of 66+) threshold: float | None = None - for key in sorted(threshold_keys, key=_parse_above_token_threshold, reverse=True): + for key in sorted(threshold_keys, key=_parse_above_token_threshold): value = model_info.get(key) if value is not None: try: # Handle both formats: _above_128k_tokens and _above_128_tokens - threshold_str = key.split("_above_")[1].split("_tokens")[0] + threshold_str = key.rsplit("_above_", 1)[1].split("_tokens")[0] threshold = _parse_above_token_threshold(key) if _prompt_exceeds_threshold(usage.prompt_tokens, threshold, threshold_is_inclusive): # Prefer a service_tier-specific above-threshold key when available, @@ -698,7 +700,7 @@ def _get_token_base_cost( service_tier, ) if service_tier - else key + else f"input_cost_per_token_above_{threshold_str}_tokens" ) prompt_base_cost = cast( float, @@ -754,8 +756,6 @@ def _get_token_base_cost( ) cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost) - - break except (IndexError, ValueError): continue except Exception: @@ -1636,20 +1636,23 @@ def get_billed_token_rates( vertex_location: str | None = None, current_time: datetime | None = None, custom_cost_per_token: CostPerToken | None = None, + model_info: ModelInfo | None = None, ) -> BilledTokenRates | None: """Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type breakdown resolve them. None when the model's pricing cannot be resolved.""" if custom_cost_per_token is not None: return _custom_pricing_rates(custom_cost_per_token) try: - model_info: Final = apply_provider_cache_read_default( - get_model_info(model=model, custom_llm_provider=custom_llm_provider), + resolved_model_info: Final = apply_provider_cache_read_default( + model_info + if model_info is not None + else get_model_info(model=model, custom_llm_provider=custom_llm_provider), custom_llm_provider, ) except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates return None return _cost_map_billed_rates( - model_info=model_info, + model_info=resolved_model_info, usage=usage, custom_llm_provider=custom_llm_provider, service_tier=service_tier, @@ -1668,6 +1671,7 @@ def get_token_type_cost_breakdown( vertex_location: str | None = None, current_time: datetime | None = None, custom_cost_per_token: CostPerToken | None = None, + model_info: ModelInfo | None = None, ) -> TokenTypeCostBreakdown: """ Provider-agnostic cost of reasoning and cache tokens, derived from the usage @@ -1692,6 +1696,7 @@ def get_token_type_cost_breakdown( vertex_location=vertex_location, current_time=current_time, custom_cost_per_token=custom_cost_per_token, + model_info=model_info, ) if rates is None: return TokenTypeCostBreakdown(0.0, 0.0, 0.0) diff --git a/litellm/main.py b/litellm/main.py index c3cf278633b..4aa4159d771 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5722,6 +5722,7 @@ def completion( optional_params=processed_non_default_params, # [IMPORTANT] - using processed_non_default_params ensures consistent params logged to langfuse for finetuning / eval datasets. litellm_params=litellm_params, custom_llm_provider=custom_llm_provider, + model_id=optional_params.get("model_id"), ) if mock_response or mock_tool_calls or mock_timeout: kwargs.pop("mock_timeout", None) # remove for any fallbacks triggered diff --git a/litellm/router.py b/litellm/router.py index aba87177a38..5e84f54c65d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -318,12 +318,14 @@ from litellm.types.utils import ( StandardLoggingRoutingDecision, Usage, all_litellm_params, + is_server_derived_pricing_key, shared_backend_model_info, ) from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, + ModelInfoLookupMiss, ModelResponse, Rules, function_setup, @@ -332,6 +334,8 @@ from litellm.utils import ( get_secret, get_utc_datetime, is_region_allowed, + lookup_model_info_for_registration, + model_info_with_pricing_overrides, provider_rejectable_params, set_live_deployment_replay, ) @@ -471,6 +475,8 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]: @@ -8824,39 +8830,7 @@ class Router: litellm_params=litellm_params, model_info=_model_info, ) - for field in CustomPricingLiteLLMParams.model_fields: - if deployment.litellm_params.get(field) is not None: - _model_info[field] = deployment.litellm_params[field] - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=_model_info, - backend_model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) - if _model_info.get("input_cost_per_token") is not None: - Router._inherit_builtin_cache_pricing( - model_info=_model_info, - backend_model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) - Router._inherit_builtin_service_tier_pricing( - model_info=_model_info, - backend_model=Router._cost_map_backend_model(deployment), - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) - Router._inherit_builtin_tiered_output_rate( - model_info=_model_info, - backend_model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) - - ## REGISTER MODEL INFO IN LITELLM MODEL COST MAP - Router._register_deployment_in_model_cost( - model_id=deployment.model_info.id, - model_info=_model_info, - model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) + Router._register_deployment_pricing(deployment) ## Check if LLM Deployment is allowed for this deployment if self.deployment_is_active_for_environment(deployment=deployment) is not True: @@ -9892,10 +9866,24 @@ class Router: params_base_model: Final = deployment.litellm_params.get("base_model") if isinstance(params_base_model, str) and params_base_model: return params_base_model + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) + invoked_model: Final = params.get("model_id") + if ( + isinstance(invoked_model, str) + and invoked_model + and ( + deployment.litellm_params.custom_llm_provider == "bedrock" + or deployment.litellm_params.model.startswith("bedrock/") + ) + ): + profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}" + profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock") + if not isinstance(profile_info, ModelInfoLookupMiss): + return profile_model return deployment.litellm_params.model @staticmethod - def _deployment_model_cost_payload(deployment: Deployment) -> dict: # mutable-ok: cost-map entry + def _deployment_model_cost_payload(deployment: Deployment) -> dict[str, object]: # mutable-ok: cost-map entry """The ``model_info`` a deployment contributes to ``litellm.model_cost``. Custom pricing lives on ``litellm_params`` rather than ``model_info``, and @@ -9903,30 +9891,49 @@ class Router: both are folded back in here. That keeps this reproducible from a deployment alone, which is what lets a refresh rebuild the same entries. """ - model_info: Final[dict] = deployment.model_info.model_dump(exclude_none=True) # mutable-ok: built in place + model_info: Final = dict( + _MODEL_INFO_ADAPTER.validate_python(deployment.model_info.model_dump(exclude_none=True)) + ) for field in CustomPricingLiteLLMParams.model_fields: field_value = deployment.litellm_params.get(field) if field_value is not None: model_info[field] = field_value + backend_model: Final = Router._cost_map_backend_model(deployment) + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) + if ( + params.get("model_id") + and backend_model != deployment.litellm_params.model + and ( + deployment.litellm_params.custom_llm_provider == "bedrock" + or deployment.litellm_params.model.startswith("bedrock/") + ) + ): + profile_info: Final = lookup_model_info_for_registration(model=backend_model, custom_llm_provider="bedrock") + if not isinstance(profile_info, ModelInfoLookupMiss): + explicit_info: Final = _MODEL_INFO_ADAPTER.validate_python(model_info) + priced_info: Final = model_info_with_pricing_overrides(profile_info, explicit_info) + model_info.update( + {key: value for key, value in priced_info.items() if is_server_derived_pricing_key(key)} + ) Router._inherit_builtin_base_rates_for_off_peak( model_info=model_info, - backend_model=deployment.litellm_params.model, + backend_model=backend_model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) if model_info.get("input_cost_per_token") is not None: Router._inherit_builtin_cache_pricing( model_info=model_info, - backend_model=deployment.litellm_params.model, + backend_model=backend_model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) Router._inherit_builtin_service_tier_pricing( model_info=model_info, - backend_model=Router._cost_map_backend_model(deployment), + backend_model=backend_model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) Router._inherit_builtin_tiered_output_rate( model_info=model_info, - backend_model=deployment.litellm_params.model, + backend_model=backend_model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) return model_info diff --git a/litellm/utils.py b/litellm/utils.py index 50e82b570fb..5f32820df5e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3468,7 +3468,10 @@ def register_model( if persist_across_reloads: _registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost for _registered_key, _registered_value in _registrations.items(): - _runtime_registered_model_cost[_registered_key] = dict(_registered_value) # mutable-ok: caller-owned + _runtime_registered_model_cost[_registered_key] = _update_dictionary( + dict(_runtime_registered_model_cost.get(_registered_key, {})), + dict(_RATE_TABLE.validate_python(_registered_value)), + ) _skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index e599a4b6353..3fc70ffec5c 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -3936,3 +3936,85 @@ def test_get_batch_cost_rates_claude_haiku_5_5_prompt_length_tiers( ) assert (rates.input, rates.output, rates.cache_read, rates.cache_creation) == expected + + +@pytest.mark.parametrize( + "field", ("output_cost_per_token", "cache_read_input_token_cost", "cache_creation_input_token_cost") +) +def test_flat_input_price_keeps_other_token_thresholds(field: str) -> None: + from typing import Final, cast + from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost + + info: Final = cast( + ModelInfo, + { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "cache_read_input_token_cost": 0.003, + "cache_creation_input_token_cost": 0.004, + f"{field}_above_300k_tokens": 0.005, + }, + ) + rates: Final = _get_token_base_cost( + model_info=info, + usage=Usage(prompt_tokens=300_001, completion_tokens=3, total_tokens=300_004), + ) + assert rates[0] == 0.001 + assert rates[1] == (0.005 if field == "output_cost_per_token" else 0.002) + assert rates[2] == (0.005 if field == "cache_creation_input_token_cost" else 0.004) + assert rates[4] == (0.005 if field == "cache_read_input_token_cost" else 0.003) + + +@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001, 512_000, 512_001)) +@pytest.mark.parametrize("higher_rate", (0.0, 0.009)) +def test_mixed_token_thresholds_preserve_each_crossed_rate(prompt_tokens: int, higher_rate: float) -> None: + from typing import Final, cast + from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost + + info: Final = cast( + ModelInfo, + { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "cache_read_input_token_cost": 0.003, + "cache_creation_input_token_cost": 0.004, + "input_cost_per_token_above_200k_tokens": 0.005, + "output_cost_per_token_above_200k_tokens": 0.006, + "cache_creation_input_token_cost_above_200k_tokens": 0.007, + "cache_read_input_token_cost_above_512k_tokens": higher_rate, + "output_cost_per_token_above_512k_tokens": higher_rate, + }, + ) + rates: Final = _get_token_base_cost( + model_info=info, + usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=3, total_tokens=prompt_tokens + 3), + ) + assert rates[0] == (0.005 if prompt_tokens > 200_000 else 0.001) + assert rates[1] == (higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002) + assert rates[2] == (0.007 if prompt_tokens > 200_000 else 0.004) + assert rates[4] == (higher_rate if prompt_tokens > 512_000 else 0.003) + + +@pytest.mark.parametrize("rate", (0.0, 0.001)) +def test_token_breakdown_uses_supplied_pricing_without_catalog_lookup(rate: float) -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "cache_read_input_token_cost": rate, + "cache_creation_input_token_cost": rate, + } + usage: Final = Usage( + prompt_tokens=20, completion_tokens=3, total_tokens=23, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), cache_creation_input_tokens=5, + ) + prompt_cost, output_cost = generic_cost_per_token( + model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info, + ) + breakdown: Final = get_token_type_cost_breakdown( + model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info, + ) + assert breakdown.cache_read_cost == pytest.approx(5 * rate) + assert breakdown.cache_creation_cost == pytest.approx(5 * rate) + assert breakdown.rates is not None + assert prompt_cost == pytest.approx(10 * breakdown.rates.input_cost_per_token + breakdown.cache_read_cost + breakdown.cache_creation_cost) + assert output_cost == pytest.approx(3 * breakdown.rates.output_cost_per_token) diff --git a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index 08bcac33a35..93320b62083 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -16,6 +16,8 @@ import pytest from botocore.credentials import Credentials from botocore.exceptions import ClientError +import litellm + from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import configuration @@ -273,3 +275,44 @@ def test_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch): sent = client.post.call_args.kwargs assert "Credential=ASIACONVERSETAGGED/" in sent["headers"]["Authorization"] assert "aws_session_tags" not in sent["data"] + + +@pytest.mark.parametrize("stream", (False, True)) +def test_converse_profile_override_records_cost_without_changing_response_model( + monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setenv("LITELLM_RUST", "0") + configuration.rust(False) + model: Final = "openai.gpt-5.6-sol" + profile: Final = f"us.{model}" + prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock") + + def respond(request: httpx.Request) -> httpx.Response: + assert request.url.path == f"/model/{profile}/converse" + return httpx.Response(200, json=CONVERSE_RESPONSE) + + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + response: Final = litellm.completion( + model=f"bedrock/converse/{model}", + model_id=profile, + messages=[{"role": "user", "content": "hi"}], + api_key="test-bearer", + aws_region_name="us-east-1", + client=HTTPHandler(client=transport), + stream=stream, + fake_stream=stream, + max_tokens=16, + ) + result: Final = litellm.stream_chunk_builder(list(response)) if stream else response + + assert result is not None + assert result.model == model + assert result.choices[0].message.content == "hi" + cost: Final = ( + response.logging_obj._response_cost_calculator(result) if stream else result._hidden_params.get("response_cost") + ) + assert cost == pytest.approx( + CONVERSE_RESPONSE["usage"]["inputTokens"] * prices["input_cost_per_token"] + + CONVERSE_RESPONSE["usage"]["outputTokens"] * prices["output_cost_per_token"] + ) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 99dea6366f9..1e9b7f596bf 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -1,10 +1,13 @@ import datetime +import json import time from pathlib import Path from types import MappingProxyType, SimpleNamespace from typing import Final, cast +import httpx import pytest +import respx from pydantic import BaseModel import litellm @@ -5319,3 +5322,620 @@ def test_completion_cost_is_zero_when_explicit_rates_are_zero(monkeypatch: pytes ) assert cost == 0.0 + + +@pytest.fixture +def isolated_pricing_registry(local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "open_ai_chat_completion_models", set(litellm.open_ai_chat_completion_models)) + monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider)) + + +_CUSTOM_PRICING_MODEL: Final = "migration-custom-pricing-chat" + + +_CUSTOM_PRICING_RATES: Final = MappingProxyType( + {"input_cost_per_token": 0.0000008, "output_cost_per_token": 0.0000032} +) + + +_CUSTOM_PRICING_EXPECTED_COST: Final = 7 * 0.0000008 + 4 * 0.0000032 + + +def _custom_pricing_chunk(choices: list[dict[str, object]], usage: dict[str, int] | None = None) -> str: + chunk: Final = { + "id": "chatcmpl-custom-pricing", + "object": "chat.completion.chunk", + "created": 1, + "model": _CUSTOM_PRICING_MODEL, + "choices": choices, + **({"usage": usage} if usage is not None else {}), + } + return f"data: {json.dumps(chunk)}\n\n" + + +def _custom_pricing_upstream(request: httpx.Request) -> httpx.Response: + if json.loads(request.content).get("stream") is True: + body: Final = "".join( + ( + _custom_pricing_chunk( + [{"index": 0, "delta": {"role": "assistant", "content": "I am ready to help."}, "finish_reason": None}] + ), + _custom_pricing_chunk([{"index": 0, "delta": {}, "finish_reason": "stop"}]), + _custom_pricing_chunk([], {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11}), + "data: [DONE]\n\n", + ) + ) + return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"}) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-custom-pricing", + "object": "chat.completion", + "created": 1, + "model": _CUSTOM_PRICING_MODEL, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "I am ready to help."}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11}, + }, + ) + + +@pytest.mark.parametrize("rate", (0.0, 0.001)) +@pytest.mark.parametrize("cache_field", ("cache_read_input_token_cost", "cache_creation_input_token_cost")) +@pytest.mark.usefixtures("isolated_pricing_registry") +@respx.mock +def test_router_cache_only_pricing_preserves_uncached_token_charges(rate: float, cache_field: str) -> None: + litellm.register_model( + {_CUSTOM_PRICING_MODEL: {**_CUSTOM_PRICING_RATES, "litellm_provider": "openai", "mode": "chat"}}, + persist_across_reloads=False, + ) + respx.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_custom_pricing_upstream) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "cache-priced-deployment", + "litellm_params": { + "model": f"openai/{_CUSTOM_PRICING_MODEL}", + "api_key": "sk-test", + cache_field: rate, + }, + "model_info": {"id": "cache-only-deployment"}, + } + ] + ) + + response: Final = router.completion( + model="cache-priced-deployment", messages=[{"role": "user", "content": "How are you?"}] + ) + + assert response._hidden_params["response_cost"] == pytest.approx(_CUSTOM_PRICING_EXPECTED_COST) + + +@pytest.mark.parametrize("profile_prefix", ("us", "global")) +@pytest.mark.parametrize( + "model_name", + ("gpt-5.6-sol", "gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-sol", "gpt-6-luna", "gpt-6-astra"), +) +@pytest.mark.parametrize("arn", (False,)) +def test_bedrock_converse_cost_uses_invoked_profile( + _local_model_cost_map: None, model_name: str, profile_prefix: str, arn: bool +) -> None: + model: Final = f"openai.{model_name}" + profile: Final = f"{profile_prefix}.{model}" + invoked: Final = f"arn:aws:bedrock:us-east-1:123456789012:inference-profile/{profile}" if arn else profile + response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)) + prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock") + + cost: Final = completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider="bedrock", + optional_params={"model_id": invoked}, + ) + + assert cost == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]) + assert response.model == model + + +@pytest.mark.parametrize("custom_pricing", (False, True)) +def test_bedrock_converse_profile_preserves_explicit_base_and_custom_pricing( + _local_model_cost_map: None, custom_pricing: bool +) -> None: + model: Final = "openai.gpt-6-sol" + response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)) + prices: Final = litellm.get_model_info(model, custom_llm_provider="bedrock") + + cost: Final = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + optional_params={"model_id": f"us.{model}"}, + custom_pricing=custom_pricing, + base_model=None if custom_pricing else model, + ) + + assert cost == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]) + + +@pytest.mark.parametrize("stream", (False, True)) +def test_bedrock_converse_cost_retains_profile_after_dispatch_consumes_optional_params( + _local_model_cost_map: None, stream: bool +) -> None: + from litellm.cost_calculator import pricing_entry_for_cost_calc + + model: Final = "openai.gpt-5.6-sol" + profile: Final = f"us.{model}" + logging_obj: Final = Logging( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=stream, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="bedrock-profile-cost", + function_id="bedrock-profile-cost", + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={"model_id": profile}, custom_llm_provider="bedrock" + ) + logging_obj.optional_params.pop("model_id") + response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)) + prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock") + + assert logging_obj._response_cost_calculator(response) == pytest.approx( + 10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + ) + entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj) + assert entry is not None + assert entry[0] == prices["key"] + + +@pytest.mark.parametrize( + "invoked_model_id", + (None, "", 123, "unmapped-profile", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque"), +) +def test_bedrock_converse_cost_keeps_configured_rates_when_target_has_no_prices( + _local_model_cost_map: None, invoked_model_id: object +) -> None: + model: Final = "openai.gpt-6-sol" + response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)) + expected: Final = completion_cost(completion_response=response, model=model, custom_llm_provider="bedrock") + + assert ( + completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + optional_params={"model_id": invoked_model_id}, + ) + == expected + ) + + +def test_non_bedrock_cost_ignores_bedrock_model_id(_local_model_cost_map: None) -> None: + response: Final = ModelResponse( + model="gpt-6-sol", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + ) + expected: Final = completion_cost(completion_response=response, custom_llm_provider="openai") + + assert ( + completion_cost( + completion_response=response, + custom_llm_provider="openai", + optional_params={"model_id": "us.openai.gpt-6-sol"}, + ) + == expected + ) + + +@pytest.mark.parametrize("rate", (0.0, 0.001)) +@pytest.mark.parametrize("persist_across_reloads", (True, False)) +@pytest.mark.parametrize("routed_response", (True, False)) +@pytest.mark.parametrize("metadata_prefix", ("bedrock/", "bedrock/converse/")) +def test_bedrock_converse_profile_preserves_registered_price_override( + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + rate: float, + persist_across_reloads: bool, + routed_response: bool, + metadata_prefix: str, +) -> None: + import litellm.utils as model_utils + from litellm.cost_calculator import pricing_entry_for_cost_calc + + monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {}) + catalog: Final = dict(litellm.model_cost) + model: Final = "openai.gpt-5.6-sol" + litellm.register_model( + {f"{metadata_prefix}{model}": {"supports_function_calling": True}}, + persist_across_reloads=persist_across_reloads, + ) + litellm.register_model( + { + model: { + "input_cost_per_token": rate, + "output_cost_per_token": rate, + "litellm_provider": "bedrock_converse", + "mode": "chat", + } + }, + persist_across_reloads=persist_across_reloads, + ) + litellm.register_model( + {model: {"supports_function_calling": True, "input_cost_per_token": None}}, + persist_across_reloads=persist_across_reloads, + ) + response: Final = ModelResponse( + model=f"bedrock/converse/{model}" if routed_response else model, + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + logging_obj: Final = Logging( + model=model, + messages=[], + stream=False, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="registered-pricing", + function_id="registered-pricing", + ) + assert completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider="bedrock", + optional_params={"model_id": f"us.{model}"}, + litellm_logging_obj=logging_obj, + ) == pytest.approx(15 * rate) + + assert logging_obj.billed_token_rates is not None + assert logging_obj.billed_token_rates.input_cost_per_token == rate + assert logging_obj.billed_token_rates.output_cost_per_token == rate + entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, None) + assert entry is not None + assert entry[0] == model + assert entry[1]["input_cost_per_token"] == rate + assert entry[1]["output_cost_per_token"] == rate + + updated_rate: Final = rate + 0.002 + litellm.register_model( + {model: {"input_cost_per_token": str(updated_rate)}}, persist_across_reloads=persist_across_reloads + ) + assert completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider="bedrock", + optional_params={"model_id": f"us.{model}"}, + ) == pytest.approx(10 * updated_rate + 5 * rate) + + monkeypatch.setattr(litellm, "model_cost", catalog) + model_utils._invalidate_model_cost_lowercase_map() + model_utils.reapply_runtime_model_cost_registrations() + profile_prices: Final = litellm.get_model_info(f"us.{model}", custom_llm_provider="bedrock") + expected_after_reload: Final = ( + 10 * updated_rate + 5 * rate + if persist_across_reloads + else 10 * profile_prices["input_cost_per_token"] + 5 * profile_prices["output_cost_per_token"] + ) + assert completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider="bedrock", + optional_params={"model_id": f"us.{model}"}, + ) == pytest.approx(expected_after_reload) + + +def test_bedrock_converse_profile_ignores_registered_capabilities_without_prices( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + import litellm.utils as model_utils + + monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {}) + model: Final = "openai.gpt-6-sol" + litellm.register_model({model: {"supports_function_calling": True}}) + response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)) + prices: Final = litellm.get_model_info(f"us.{model}", custom_llm_provider="bedrock") + + assert completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + optional_params={"model_id": f"us.{model}"}, + ) == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]) + + +@pytest.mark.parametrize("rate", (0.0, 0.001)) +@pytest.mark.parametrize("cache_field", ("cache_read_input_token_cost", "cache_creation_input_token_cost")) +@pytest.mark.parametrize("has_base_rates", (False, True)) +@pytest.mark.parametrize("persist_across_reloads", (False, True)) +@pytest.mark.parametrize("profile_prefix", ("us", "global")) +@pytest.mark.parametrize("provider", (None, "bedrock")) +def test_bedrock_converse_profile_preserves_cache_only_registered_prices( + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + rate: float, + cache_field: str, + has_base_rates: bool, + persist_across_reloads: bool, + profile_prefix: str, + provider: str | None, +) -> None: + import litellm.utils as model_utils + from litellm.cost_calculator import pricing_entry_for_cost_calc + + monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {}) + model: Final = "openai.cache-overlay-test" + profile: Final = f"{profile_prefix}.{model}" + catalog: Final = { + **litellm.model_cost, + model: { + "litellm_provider": "bedrock", + "mode": "chat", + **({"input_cost_per_token": 0.01, "output_cost_per_token": 0.02} if has_base_rates else {}), + }, + profile: { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.02, + "output_cost_per_token": 0.04, + "cache_read_input_token_cost": 0.005, + "cache_creation_input_token_cost": 0.025, + }, + } + profile_before: Final = dict(catalog[profile]) + monkeypatch.setattr(litellm, "model_cost", dict(catalog)) + model_utils._invalidate_model_cost_lowercase_map() + litellm.register_model({model: {cache_field: rate}}, persist_across_reloads=persist_across_reloads) + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=15, + completion_tokens=3, + total_tokens=18, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=5 if cache_field == "cache_read_input_token_cost" else 0 + ), + cache_creation_input_tokens=5 if cache_field == "cache_creation_input_token_cost" else 0, + ), + ) + logging_obj: Final = Logging( + model=model, + messages=[], + stream=False, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="cache-overlay", + function_id="cache-overlay", + ) + logging_obj.model_call_details["model_id"] = profile + expected: Final = 10 * 0.02 + 3 * 0.04 + 5 * rate + assert completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider=provider, + optional_params={"model_id": profile}, + litellm_logging_obj=logging_obj, + ) == pytest.approx(expected) + breakdown_key: Final = "cache_read_cost" if cache_field == "cache_read_input_token_cost" else "cache_creation_cost" + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown.get(breakdown_key, 0.0) == pytest.approx(5 * rate) + assert logging_obj.billed_token_rates is not None + assert ( + logging_obj.billed_token_rates.cache_read_input_token_cost + if cache_field == "cache_read_input_token_cost" + else logging_obj.billed_token_rates.cache_creation_input_token_cost + ) == rate + entry: Final = pricing_entry_for_cost_calc( + f"bedrock/converse/{model}", response, provider, False, None, None, None, logging_obj + ) + assert entry is not None + assert entry[1]["input_cost_per_token"] == 0.02 + assert entry[1]["output_cost_per_token"] == 0.04 + assert entry[1][cache_field] == rate + assert litellm.model_cost[profile] == profile_before + + monkeypatch.setattr(litellm, "model_cost", dict(catalog)) + model_utils._invalidate_model_cost_lowercase_map() + model_utils.reapply_runtime_model_cost_registrations() + expected_after_reload: Final = ( + 10 * 0.02 + 3 * 0.04 + 5 * (rate if persist_across_reloads else catalog[profile][cache_field]) + ) + assert completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider=provider, + optional_params={"model_id": profile}, + litellm_logging_obj=logging_obj, + ) == pytest.approx(expected_after_reload) + + +@pytest.mark.parametrize("profile", ("us.openai.missing-profile-test", "us.openai.unpriced-profile-test")) +def test_bedrock_cache_registration_with_unpriced_profile_keeps_registered_rates( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, profile: str +) -> None: + import litellm.utils as model_utils + + model: Final = "openai.registered-cache-test" + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + model: { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + "us.openai.unpriced-profile-test": {"litellm_provider": "bedrock", "mode": "chat"}, + }, + ) + model_utils._invalidate_model_cost_lowercase_map() + litellm.register_model({model: {"cache_read_input_token_cost": 0.001}}, persist_across_reloads=False) + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=15, + completion_tokens=3, + total_tokens=18, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), + ), + ) + assert completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider="bedrock", + optional_params={"model_id": profile}, + ) == pytest.approx(10 * 0.01 + 3 * 0.02 + 5 * 0.001) + + +@pytest.mark.parametrize("direction", ("input", "output")) +@pytest.mark.parametrize("rate", (0.0, 0.004)) +@pytest.mark.parametrize("prompt_tokens", (100, 300_001)) +def test_bedrock_partial_registration_overrides_only_its_profile_direction( + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + direction: str, + rate: float, + prompt_tokens: int, +) -> None: + from litellm.cost_calculator import pricing_entry_for_cost_calc + + model: Final = "openai.partial-pricing-test" + profile: Final = f"us.{model}" + profile_entry: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "input_cost_per_token_above_300k_tokens": 0.03, + "output_cost_per_token_above_300k_tokens": 0.06, + "cache_read_input_token_cost": 0.001, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, profile: dict(profile_entry)}) + litellm.register_model({model: {f"{direction}_cost_per_token": rate}}, persist_across_reloads=False) + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=3, + total_tokens=prompt_tokens + 3, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), + ), + ) + logging_obj: Final = Logging( + model=model, + messages=[], + stream=False, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="partial-pricing", + function_id="partial-pricing", + ) + logging_obj.model_call_details["model_id"] = profile + input_rate: Final = rate if direction == "input" else (0.03 if prompt_tokens > 300_000 else 0.01) + output_rate: Final = rate if direction == "output" else (0.06 if prompt_tokens > 300_000 else 0.02) + assert completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + litellm_logging_obj=logging_obj, + ) == pytest.approx((prompt_tokens - 5) * input_rate + 3 * output_rate + 5 * 0.001) + assert logging_obj.billed_token_rates is not None + assert logging_obj.billed_token_rates.input_cost_per_token == input_rate + assert logging_obj.billed_token_rates.output_cost_per_token == output_rate + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(5 * 0.001) + entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj) + assert entry is not None + assert entry[1][f"{direction}_cost_per_token"] == rate + assert litellm.model_cost[profile] == profile_entry + + +@pytest.mark.parametrize("prompt_tokens", (299_999, 300_000, 300_001)) +@pytest.mark.parametrize("cache_field", ("cache_read_input_token_cost", "cache_creation_input_token_cost")) +@pytest.mark.parametrize("tier_override", (None, 0.0, 0.003)) +def test_bedrock_registered_cache_prices_preserve_dynamic_threshold_rates( + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + prompt_tokens: int, + cache_field: str, + tier_override: float | None, +) -> None: + import litellm.utils as model_utils + from litellm.cost_calculator import pricing_entry_for_cost_calc + + model: Final = "openai.threshold-overlay-test" + profile: Final = f"us.{model}" + tier_field: Final = f"{cache_field}_above_300k_tokens" + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + profile: { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "input_cost_per_token_above_300k_tokens": 0.03, + "output_cost_per_token_above_300k_tokens": 0.06, + tier_field: 0.007, + }, + }, + ) + model_utils._invalidate_model_cost_lowercase_map() + litellm.register_model( + { + model: { + cache_field: 0.001, + **({tier_field: tier_override} if tier_override is not None else {}), + } + }, + persist_across_reloads=False, + ) + usage: Final = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=3, + total_tokens=prompt_tokens + 3, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=5 if cache_field == "cache_read_input_token_cost" else 0 + ), + cache_creation_input_tokens=5 if cache_field == "cache_creation_input_token_cost" else 0, + ) + response: Final = ModelResponse(model=model, usage=usage) + logging_obj: Final = Logging( + model=model, + messages=[], + stream=False, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="threshold-overlay", + function_id="threshold-overlay", + ) + logging_obj.model_call_details["model_id"] = profile + expected_input_rate: Final = 0.03 if prompt_tokens > 300_000 else 0.01 + expected_output_rate: Final = 0.06 if prompt_tokens > 300_000 else 0.02 + expected_cache_rate: Final = ( + (tier_override if tier_override is not None else 0.007) if prompt_tokens > 300_000 else 0.001 + ) + total: Final = completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider="bedrock", + optional_params={"model_id": profile}, + litellm_logging_obj=logging_obj, + ) + assert total == pytest.approx( + (prompt_tokens - 5) * expected_input_rate + 3 * expected_output_rate + 5 * expected_cache_rate + ) + breakdown_key: Final = "cache_read_cost" if cache_field == "cache_read_input_token_cost" else "cache_creation_cost" + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown.get(breakdown_key, 0.0) == pytest.approx(5 * expected_cache_rate) + assert logging_obj.billed_token_rates is not None + assert logging_obj.billed_token_rates.input_cost_per_token == expected_input_rate + assert logging_obj.billed_token_rates.output_cost_per_token == expected_output_rate + entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj) + assert entry is not None + assert entry[1]["input_cost_per_token_above_300k_tokens"] == 0.03 + assert entry[1][tier_field] == (tier_override if tier_override is not None else 0.007) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 2c99228013e..0302c3d648b 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4269,3 +4269,98 @@ def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): ) assert exc_info.value.status_code == 400 assert f"tool_choice={tool_choice}" in str(exc_info.value) + + +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("profile,custom_price", ((None, False), ("us.openai.gpt-6-sol", False), (None, True))) +@pytest.mark.asyncio +async def test_converse_fallback_prices_the_current_attempt( + monkeypatch: pytest.MonkeyPatch, stream: bool, asynchronous: bool, profile: str | None, custom_price: bool +) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.rust_bridge import configuration + + CONVERSE_ANSWER: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + "metrics": {"latencyMs": 1}, + } + ).encode() + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setenv("LITELLM_RUST", "0") + configuration.rust(False) + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + failed_profile: Final = "us.amazon.nova-micro-v1:0" + model: Final = "openai.gpt-6-sol" + logging_obj: Final = LiteLLMLogging( + model="amazon.nova-micro-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=stream, + call_type="acompletion" if asynchronous else "completion", + start_time=datetime(2026, 1, 1), + litellm_call_id="bedrock-fallback-pricing", + function_id="bedrock-fallback-pricing", + ) + + def respond(request: httpx.Request) -> httpx.Response: + if failed_profile in request.url.path: + return httpx.Response(400, json={"message": "Input is too long for requested model."}) + assert request.url.path == f"/model/{profile or model}/converse" + return httpx.Response(200, content=CONVERSE_ANSWER) + + client: Final = ( + AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + if asynchronous + else HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + ) + kwargs: Final = { + "messages": [{"role": "user", "content": "hi"}], + "api_key": "test-bearer", + "aws_region_name": "us-east-1", + "client": client, + "litellm_logging_obj": logging_obj, + "num_retries": 0, + "stream": stream, + "fake_stream": stream, + } + if asynchronous: + with pytest.raises(litellm.BadRequestError): + await litellm.acompletion(model="bedrock/converse/amazon.nova-micro-v1:0", model_id=failed_profile, **kwargs) + else: + with pytest.raises(litellm.BadRequestError): + litellm.completion(model="bedrock/converse/amazon.nova-micro-v1:0", model_id=failed_profile, **kwargs) + + overrides: Final = { + **({"model_id": profile} if profile is not None else {}), + **({"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} if custom_price else {}), + } + response: Final = ( + await litellm.acompletion(model=f"bedrock/converse/{model}", **kwargs, **overrides) + if asynchronous + else litellm.completion(model=f"bedrock/converse/{model}", **kwargs, **overrides) + ) + result: Final = ( + litellm.stream_chunk_builder( + [chunk async for chunk in response] if asynchronous else list(response), logging_obj=response.logging_obj + ) + if stream + else response + ) + assert result is not None + prices: Final = litellm.get_model_info(profile or model, custom_llm_provider="bedrock") + expected: Final = ( + 0.001 + 0.002 + if custom_price + else prices["input_cost_per_token"] + prices["output_cost_per_token"] + ) + active_logging: Final = response.logging_obj if stream else logging_obj + assert active_logging._response_cost_calculator(result) == pytest.approx(expected) + if asynchronous: + await client.close() + else: + client.close() + configuration.reset_rust_configuration() diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 6a6e4c5b4d4..2a7d9813925 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -761,6 +761,178 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): _restore_model_cost_entries(model_keys) +@pytest.mark.parametrize("service_tier", ("default", "priority", "flex")) +@pytest.mark.parametrize( + "field", + ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost", "input_cost_per_token_priority"), +) +@pytest.mark.parametrize("rate", (0.0, 0.004)) +@pytest.mark.parametrize("explicit_base", (False, True)) +@pytest.mark.parametrize("provider", (None, "bedrock")) +@pytest.mark.parametrize("registration_path", ("direct", "startup", "reload")) +def test_bedrock_deployment_inherits_invoked_profile_around_custom_prices( + monkeypatch: pytest.MonkeyPatch, + field: str, + rate: float, + explicit_base: bool, + provider: str | None, + registration_path: str, + service_tier: str, +) -> None: + from litellm.cost_calculator import completion_cost + from litellm.types.utils import ModelResponse, Usage, PromptTokensDetailsWrapper + + model: Final = "openai.gpt-5.6-sol" + profile: Final = f"us.{model}" + base_model: Final = f"global.{model}" + profile_prices: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "cache_read_input_token_cost": 0.001, + "input_cost_per_token_priority": 0.05, + "output_cost_per_token_priority": 0.06, + "input_cost_per_token_flex": 0.007, + "output_cost_per_token_flex": 0.008, + } + base_prices: Final = {**profile_prices, "input_cost_per_token": 0.03, "output_cost_per_token": 0.04} + monkeypatch.setattr( + litellm, "model_cost", {**litellm.model_cost, profile: dict(profile_prices), base_model: dict(base_prices)} + ) + deployment: Final = Deployment( + model_name="profile-priced", + litellm_params=LiteLLM_Params( + model=f"converse/{model}" if provider else f"bedrock/converse/{model}", + model_id=profile, + custom_llm_provider=provider, + ), + model_info=ModelInfo( + **{ + "id": "profile-priced-deployment", + field: rate, + **({"base_model": f"bedrock/{base_model}"} if explicit_base else {}), + } + ), + ) + if registration_path == "direct": + Router._register_deployment_pricing(deployment) + else: + router: Final = Router(model_list=[deployment.to_json(exclude_none=True)]) + if registration_path == "reload": + router._replay_model_cost_registrations() + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=100, + completion_tokens=3, + total_tokens=103, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), + ), + ) + catalog: Final = base_prices if explicit_base else profile_prices + input_field: Final = "input_cost_per_token" + (f"_{service_tier}" if service_tier != "default" else "") + output_field: Final = "output_cost_per_token" + (f"_{service_tier}" if service_tier != "default" else "") + input_rate: Final = rate if field == input_field else catalog[input_field] + output_rate: Final = rate if field == output_field else catalog[output_field] + cache_rate: Final = rate if field == "cache_read_input_token_cost" else 0.001 + assert completion_cost( + model=model, + completion_response=response, + custom_llm_provider="bedrock", + custom_pricing=True, + router_model_id="profile-priced-deployment", + service_tier=service_tier, + ) == pytest.approx(95 * input_rate + 3 * output_rate + 5 * cache_rate) + assert litellm.model_cost[profile] == profile_prices + + +@pytest.mark.parametrize("catalog_entry", (None, {"litellm_provider": 1, "mode": "chat"})) +@pytest.mark.parametrize("known_model", (False, True)) +@pytest.mark.parametrize("explicit_base", (False, True)) +@pytest.mark.parametrize("service_tier", ("default", "priority", "flex")) +@pytest.mark.parametrize("override", (None, "cache_read_input_token_cost", "input_cost_per_token_priority")) +@pytest.mark.parametrize("reload", (False, True)) +def test_bedrock_deployment_keeps_custom_prices_when_profile_lookup_fails( + monkeypatch: pytest.MonkeyPatch, + catalog_entry: dict[str, object] | None, + known_model: bool, + explicit_base: bool, + service_tier: str, + override: str | None, + reload: bool, +) -> None: + from litellm.cost_calculator import completion_cost + from litellm.types.utils import ModelResponse, Usage, PromptTokensDetailsWrapper + + model: Final = "anthropic.unmapped-profile-pricing-test" + profile: Final = ( + "us.anthropic.invalid-profile-pricing-test" + if catalog_entry is not None + else "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque-test" + ) + configured_prices: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.03, + "output_cost_per_token": 0.04, + "cache_read_input_token_cost": 0.001, + "input_cost_per_token_priority": 0.05, + "output_cost_per_token_priority": 0.06, + "input_cost_per_token_flex": 0.007, + "output_cost_per_token_flex": 0.008, + } + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + **({model: dict(configured_prices)} if known_model else {}), + **({profile: catalog_entry} if catalog_entry is not None else {}), + }, + ) + deployment: Final = Deployment( + model_name="custom-priced", + litellm_params=LiteLLM_Params(model=f"bedrock/converse/{model}", model_id=profile), + model_info=ModelInfo( + id="unmapped-profile-deployment", + input_cost_per_token=0.01, + output_cost_per_token=0.02, + **({override: 0.0} if override else {}), + **({"base_model": f"bedrock/{profile}"} if explicit_base else {}), + ), + ) + router: Final = Router(model_list=[deployment.to_json(exclude_none=True)]) + if reload: + router._replay_model_cost_registrations() + inherit_configured: Final = known_model and not explicit_base + tier_rates: Final = inherit_configured and service_tier != "default" + input_rate: Final = configured_prices[f"input_cost_per_token_{service_tier}"] if tier_rates else 0.01 + output_rate: Final = configured_prices[f"output_cost_per_token_{service_tier}"] if tier_rates else 0.02 + effective_input: Final = 0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate + cache_rate: Final = 0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input) + assert completion_cost( + model=model, + completion_response=ModelResponse( + model=model, + usage=Usage( + prompt_tokens=100, + completion_tokens=3, + total_tokens=103, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=80), + ), + ), + custom_llm_provider="bedrock", + custom_pricing=True, + router_model_id="unmapped-profile-deployment", + service_tier=service_tier, + ) == pytest.approx(20 * effective_input + 80 * cache_rate + 3 * output_rate) + if known_model: + assert { + key: value for key, value in litellm.model_cost[model].items() if key in configured_prices + } == configured_prices + + def test_partial_custom_pricing_inherits_builtin_cache_pricing(): """A deployment that overrides only input/output cost on a cache-supporting model must still bill cache_read and cache_creation tokens. Before the @@ -1771,14 +1943,7 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None: ) rebuilt = litellm.model_cost["priced-id"] - assert at_boot.items() <= rebuilt.items(), ( - f"the rebuild changed or dropped a field the boot registration wrote: " - f"{ {k: (v, rebuilt.get(k)) for k, v in at_boot.items() if rebuilt.get(k) != v} }" - ) - assert {field: rebuilt[field] for field in set(rebuilt) - set(at_boot)} == { - "db_model": False, - "member_auto_router": False, - } + assert at_boot == rebuilt assert "member_auto_router" not in litellm.model_cost["gpt-4o"] assert router.model_list finally: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index c40121e284f..b9cf6abd360 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -4205,7 +4205,11 @@ async def test_builtin_string_callback_registers_when_subclass_already_active( assert any(type(cb) is S3Logger for cb in litellm._async_success_callback) -def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypatch): +@pytest.mark.parametrize("rate", (0.0, 0.000123)) +@pytest.mark.parametrize("updated_rate", (None, 0.0, 0.000456)) +def test_reapply_runtime_registrations_replays_register_model_overrides( + monkeypatch: pytest.MonkeyPatch, rate: float, updated_rate: float | None +) -> None: """ register_model is the documented way to override pricing for a model. A price-data reload swaps litellm.model_cost for a freshly fetched catalog, @@ -4236,11 +4240,17 @@ def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypa "openai/gpt-4o": { "litellm_provider": "openai", "mode": "chat", - "input_cost_per_token": 0.000123, + "input_cost_per_token": rate, } } ) + litellm.register_model({"openai/gpt-4o": {"supports_function_calling": True, "input_cost_per_token": None}}) + litellm.register_model( + {"openai/gpt-4o": {"input_cost_per_token": str(updated_rate) if updated_rate is not None else None}} + ) + litellm.register_model({"openai/gpt-4o": {"supports_function_calling": False}}) + litellm.model_cost = { "openai/gpt-4o": { "litellm_provider": "openai", @@ -4252,7 +4262,9 @@ def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypa _invalidate_model_cost_lowercase_map() reapply_runtime_model_cost_registrations() - assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == 0.000123 + expected_rate: Final = rate if updated_rate is None else updated_rate + assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == expected_rate + assert litellm.model_cost["openai/gpt-4o"]["supports_function_calling"] is False assert litellm.model_cost["openai/gpt-4o"]["max_input_tokens"] == 4242 finally: litellm.model_cost = saved_model_cost @@ -6362,3 +6374,33 @@ def test_function_setup_never_logs_the_ocr_data_uri_payload() -> None: assert logged == [{"role": "user", "content": f"data:application/pdf;base64 ({len(payload)} chars)"}] assert payload not in str(logged) + + +@pytest.mark.parametrize( + "key", ("openai.test-converse", "bedrock/openai.test-converse", "bedrock/converse/openai.test-converse") +) +def test_runtime_model_cost_registrations_resolve_bedrock_routing(monkeypatch: pytest.MonkeyPatch, key: str) -> None: + from litellm import utils as model_utils + + rates: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} + monkeypatch.setattr( + litellm, "model_cost", {**litellm.model_cost, key: {"litellm_provider": "bedrock_converse", "mode": "chat"}} + ) + monkeypatch.setattr(model_utils, "_runtime_registered_fields", {}) + litellm.register_model({key: rates}, persist_across_reloads=False) + + registrations: Final = dict( + model_utils.runtime_model_cost_registrations("bedrock/converse/openai.test-converse", "bedrock") + ) + assert registrations[key] == rates + assert model_utils.runtime_model_cost_registrations("bedrock/converse/openai.other-converse", "bedrock") == () + + +def test_registered_fields_reset_when_model_cost_is_replaced(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_cost", {}) + litellm.register_model({"registered-fields-probe": {"input_cost_per_token": 1e-06, "litellm_provider": "openai"}}) + assert "registered-fields-probe" in litellm.utils._registered_fields_for_current_map() + + monkeypatch.setattr(litellm, "model_cost", {}) + + assert litellm.utils._registered_fields_for_current_map() == {}