diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index c99381c8f26..cadd39e7a7e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -439,6 +439,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. @@ -740,6 +741,14 @@ 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": + if bedrock_model_info is not None: + return generic_cost_per_token( + model=model, + usage=usage_block, + custom_llm_provider="bedrock", + service_tier=service_tier, + model_info=bedrock_model_info, + ) registered_pricing_model: Final = _bedrock_registered_pricing_model(model) if registered_pricing_model is not None: return generic_cost_per_token( @@ -889,6 +898,38 @@ def _bedrock_registered_pricing_model(model: str) -> str | None: ) +_MODEL_INFO_ADAPTER: Final = TypeAdapter(ModelInfo) + + +def _bedrock_cache_overlay( + 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 or not invoked_model_id: + return None + registration: Final = next( + ( + entry + for _, entry in runtime_model_cost_registrations(model, "bedrock") + if _bedrock_registration_has_prices(entry) + ), + None, + ) + if registration is None or _cost_map_entry_prices_anything(registration): + return None + profile: Final = _cost_map_model_info(invoked_model_id, "bedrock", prefer_registration=False) + if profile is None or not _declares_rates(profile): + return None + overrides: Final = { + field: value + for field, value in registration.items() + if value is not None and field.startswith("cache_") and is_server_derived_pricing_key(field) + } + return _MODEL_INFO_ADAPTER.validate_python({**profile, **overrides}) + + 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) @@ -955,7 +996,7 @@ def select_model_name_for_cost_calc( 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 invoked_model_info: Final = ( - _cost_map_model_info(invoked_model_id, "bedrock") + _cost_map_model_info(invoked_model_id, "bedrock", prefer_registration=False) if invoked_model_id and custom_llm_provider == "bedrock" and not explicit_pricing else None ) @@ -967,9 +1008,12 @@ def select_model_name_for_cost_calc( if invoked_model_info is not None and isinstance(configured_model, str) else None ) - bedrock_pricing_model: Final = registered_pricing_model or ( - invoked_model_info["key"] if invoked_model_info is not None and _declares_rates(invoked_model_info) else None + cache_overlay: Final = _bedrock_cache_overlay( + configured_model, invoked_model_id, custom_llm_provider, explicit_pricing ) + bedrock_pricing_model: Final = ( + cache_overlay["key"] if cache_overlay is not None else registered_pricing_model + ) or (invoked_model_info["key"] if invoked_model_info is not None and _declares_rates(invoked_model_info) 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 @@ -1602,6 +1646,10 @@ def completion_cost( if base_model is not None and custom_pricing is not True else 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, @@ -1610,12 +1658,15 @@ def completion_cost( base_model=base_model, router_model_id=router_model_id, region_name=region_name, - invoked_model_id=_invoked_model_id_for_cost( - _LOGGED_LITELLM_PARAMS.validate_python(optional_params) if optional_params is not None else None, - litellm_logging_obj, - ), + invoked_model_id=invoked_model_id, ) + bedrock_model_info: Final = _bedrock_cache_overlay( + model or _get_response_model(completion_response), + invoked_model_id, + _get_provider_for_cost_calc(model, custom_llm_provider), + explicit_pricing, + ) requested_model: Final = model request_pricing_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None raw_candidates: Final = ( @@ -2004,6 +2055,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) @@ -2344,10 +2396,14 @@ def _layered_ocr_pricing(*sources: Mapping[str, object] | None) -> OCRPricing: ) -def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None: +def _cost_map_model_info( + model: str, custom_llm_provider: str | None, *, prefer_registration: bool = True +) -> ModelInfo | None: try: registered_pricing_model: Final = ( - _bedrock_registered_pricing_model(model) if custom_llm_provider == "bedrock" else None + _bedrock_registered_pricing_model(model) + if custom_llm_provider == "bedrock" and prefer_registration + else None ) if registered_pricing_model is not None: return litellm.get_model_info(model=registered_pricing_model) @@ -2404,6 +2460,14 @@ def pricing_entry_for_cost_calc( region_name=region_name, invoked_model_id=_invoked_model_id_for_cost(None, litellm_logging_obj), ) + cache_overlay: Final = _bedrock_cache_overlay( + 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 cache_overlay is not None: + return cache_overlay["key"], cache_overlay 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/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 3759f3e6cdc..3d7f1d9f8b0 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -9973,28 +9973,133 @@ def test_bedrock_converse_profile_ignores_registered_capabilities_without_prices @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, rate: float, cache_field: str + _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: - model: Final = "openai.gpt-6-sol" - litellm.register_model({model: {cache_field: rate}}, persist_across_reloads=False) + 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=5, - completion_tokens=0, - total_tokens=5, + 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, ), ) - expected: Final = completion_cost(completion_response=response, model=model, custom_llm_provider="bedrock") - assert expected == pytest.approx(5 * rate) + 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=model, - custom_llm_provider="bedrock", - optional_params={"model_id": f"us.{model}"}, + model=f"bedrock/converse/{model}", + custom_llm_provider=provider, + optional_params={"model_id": profile}, ) == pytest.approx(expected) + 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}, + ) == 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)