From cb50b02a2c1afb1c1af6545b985cfe6c33b60b2e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:32:53 -0700 Subject: [PATCH] fix(bedrock): price Converse calls using the invoked model profile --- litellm/cost_calculator.py | 41 ++++- litellm/utils.py | 15 ++ .../chat/test_bedrock_converse_handler.py | 41 +++++ tests/unit/test_cost_calculator.py | 163 ++++++++++++++++++ tests/unit/test_utils.py | 13 ++ 5 files changed, 271 insertions(+), 2 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4fcd3bde843..d9e6aedf45f 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -149,6 +149,7 @@ from litellm.utils import ( TextCompletionResponse, TranscriptionResponse, cached_get_model_info_helper, + runtime_model_cost_registration, token_counter, unmapped_model_generalization, ) @@ -886,6 +887,16 @@ def _pricing_provider_for_cost_candidate( return current_provider if isinstance(current_provider, str) else None +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, @@ -894,6 +905,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 @@ -914,6 +926,26 @@ 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 + invoked_model_info: Final = ( + _cost_map_model_info(invoked_model_id, "bedrock") + if invoked_model_id and custom_llm_provider == "bedrock" and not explicit_pricing + else None + ) + configured_model: Final = model or ( + completion_response_model if isinstance(completion_response_model, str) else None + ) + registered_pricing: Final = ( + runtime_model_cost_registration(configured_model, "bedrock") + if invoked_model_info is not None and isinstance(configured_model, str) + else None + ) + invoked_pricing_model: Final = ( + invoked_model_info["key"] + if invoked_model_info is not None + and _declares_rates(invoked_model_info) + and not (registered_pricing is not None and _cost_map_entry_prices_anything(registered_pricing)) + 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 @@ -931,8 +963,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 invoked_pricing_model is not None or provider_response_model is not None: + return_model = base_model if base_model is not None else invoked_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: @@ -1552,6 +1584,10 @@ 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, + ), ) requested_model: Final = model @@ -2335,6 +2371,7 @@ 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), ) candidates: Final = (selected_model, _get_response_model(completion_response), model) resolved: Final = next( diff --git a/litellm/utils.py b/litellm/utils.py index 143c2470d33..6281be334df 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3613,6 +3613,21 @@ def _register_on_current_map( _runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload +def runtime_model_cost_registration(model: str, custom_llm_provider: str) -> Mapping[str, object] | None: + names: Final = get_potential_model_names(model, custom_llm_provider) + candidates: Final = ( + model, + names["combined_model_name"], + f"{custom_llm_provider}/{names['split_model']}", + names["split_model"], + names["stripped_model_name"], + ) + registered: Final = next( + (_runtime_registered_model_cost[key] for key in candidates if key in _runtime_registered_model_cost), None + ) + return MappingProxyType(registered) if registered is not None else None + + class _LiveDeploymentReplay: """Single-slot holder for the callback that rebuilds live router deployments. 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 1fa83254822..7d1b502df2b 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -918,3 +918,44 @@ def test_sync_converse_streaming_upstream_error_carries_the_upstream_body_as_tex assert raised.value.status_code == status_code assert raised.value.message == f"litellm.{error_class.__name__}: BedrockException - {_UPSTREAM_ERROR_BODY}" + + +@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 2d13829bf33..c0cdc246c0d 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -9720,3 +9720,166 @@ def test_responses_ws_logging_object_carries_the_completed_responses_id_tier_and assert (logged.id, logged.service_tier) == ("resp_a", "priority") assert (logged.usage.prompt_tokens, logged.usage.completion_tokens, logged.usage.total_tokens) == (10, 5, 15) assert ResponsesWebSocketTokenUsageProcessor.create_logging_object([{"type": "error", "error": {}}]).id is None + + +@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", "gpt-6.1-sol"), +) +@pytest.mark.parametrize("arn", (False, True)) +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)) +def test_bedrock_converse_profile_preserves_registered_price_override( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, rate: float +) -> None: + import litellm.utils as model_utils + + monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {}) + model: Final = "openai.gpt-5.6-sol" + litellm.register_model( + { + model: { + "input_cost_per_token": rate, + "output_cost_per_token": rate, + "litellm_provider": "bedrock_converse", + "mode": "chat", + } + } + ) + response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)) + + assert completion_cost( + completion_response=response, + model=f"bedrock/converse/{model}", + custom_llm_provider="bedrock", + optional_params={"model_id": f"us.{model}"}, + ) == pytest.approx(15 * rate) + + +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"]) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 69ef2caf8f4..c491af9de18 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10094,3 +10094,16 @@ def test_vertex_llama3_model_info_resolves_without_the_meta_prefix(): prefixed_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") assert litellm.get_model_info(model=short_name, custom_llm_provider="vertex_ai") == prefixed_info, model assert prefixed_info["key"] == f"vertex_ai/{model}", model + + +@pytest.mark.parametrize( + "key", ("openai.test-converse", "bedrock/openai.test-converse", "bedrock/converse/openai.test-converse") +) +def test_runtime_model_cost_registration_resolves_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(model_utils, "_runtime_registered_model_cost", {key: rates}) + + assert model_utils.runtime_model_cost_registration("bedrock/converse/openai.test-converse", "bedrock") == rates + assert model_utils.runtime_model_cost_registration("bedrock/converse/openai.other-converse", "bedrock") is None