From a7052ea8a80296ca3db064e438904f562f176933 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Thu, 12 Nov 2026 15:54:56 -0800 Subject: [PATCH] fix(cost): preserve Bedrock cross-region pricing for Messages (#51479) * fix(cost): bill Bedrock cross-region profiles at the profile rate on /v1/messages Bedrock reports the region-stripped model id for a cross-region inference profile, and the cost selector preferred that served id, so streaming and non-streaming Messages spend (input, output, cache write and cache read) used the global row. A served id carrying a different cross-region prefix keeps its own price. Fixes #46286 * test(langfuse): derive the Bedrock profile export cost from the profile's model info The regional us. deployment is now billed at its profile rate, so the expected Langfuse cost comes from that profile's rates and the mock usage instead of a hardcoded global price. --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/cost_calculator.py | 18 ++ .../completion_with_bedrock_call.json | 3 - .../integrations/langfuse/test_langfuse.py | 23 +- tests/unit/test_cost_calculator.py | 200 +++++++++++++++++- 4 files changed, 237 insertions(+), 7 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 91529ae9ff5..03feb8dc1e6 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1144,11 +1144,29 @@ def _served_or_priced_requested_model( if served in (requested, requested.partition("/")[2]): return served pricing_provider: Final = _get_provider_for_cost_calc(requested, custom_llm_provider) + profile_of_served: Final = _is_bedrock_cross_region_profile_of(requested_model, served_model, pricing_provider) + if profile_of_served and _resolves_to_rates(requested, pricing_provider): + return requested if _resolves_to_rates(served, pricing_provider): return served return requested if _resolves_to_rates(requested, pricing_provider) else served +def _is_bedrock_cross_region_profile_of(profile: str, served_model: str, custom_llm_provider: str | None) -> bool: + if custom_llm_provider != "bedrock": + return False + from litellm.llms.bedrock.common_utils import ( + get_bedrock_cross_region_inference_regions, + strip_bedrock_routing_prefix, + ) + + region, _, profiled_model = strip_bedrock_routing_prefix(profile).partition(".") + return ( + region in get_bedrock_cross_region_inference_regions() + and strip_bedrock_routing_prefix(served_model) == profiled_model + ) + + def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str: """Resolve a provider-prefixed slash alias like "vertex_ai/vertex/claude-opus-5" to the registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already diff --git a/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json index 6f359380245..96548edc60f 100644 --- a/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json +++ b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json @@ -2,9 +2,6 @@ "name": "litellm-acompletion", "parent_span_id": null, "attributes": { - "langfuse.observation.cost_details": { - "total": 6e-05 - }, "langfuse.observation.input": { "messages": [ { diff --git a/tests/unit/integrations/langfuse/test_langfuse.py b/tests/unit/integrations/langfuse/test_langfuse.py index 7a374c1affe..e7fcad5c748 100644 --- a/tests/unit/integrations/langfuse/test_langfuse.py +++ b/tests/unit/integrations/langfuse/test_langfuse.py @@ -97,7 +97,9 @@ def _comparable(span: Mapping[str, object]) -> dict[str, object]: } -def _expected_generation(route: respx.Route, fixture_name: str, trace_id: str) -> dict[str, object]: +def _expected_generation( + route: respx.Route, fixture_name: str, trace_id: str, cost_total: float | None = None +) -> dict[str, object]: otel_trace_id: Final = resolve_trace_id(trace_id) generations: Final = tuple( span @@ -108,7 +110,14 @@ def _expected_generation(route: respx.Route, fixture_name: str, trace_id: str) - ) assert len(generations) == 1 expected_path: Final = Path(__file__).with_name("langfuse_expected_request_body") / fixture_name - expected: Final = TypeAdapter(dict[str, object]).validate_python(json.loads(expected_path.read_text())) + fixture: Final = TypeAdapter(dict[str, object]).validate_python(json.loads(expected_path.read_text())) + fixture_attributes: Final = TypeAdapter(dict[str, object]).validate_python(fixture["attributes"]) + priced_attributes: Final = ( + fixture_attributes + if cost_total is None + else {**fixture_attributes, "langfuse.observation.cost_details": {"total": cost_total}} + ) + expected: Final = {**fixture, "attributes": priced_attributes} actual: Final = _comparable(generations[0]) expected_comparable: Final = _comparable(expected) assert actual == expected_comparable, json.dumps( @@ -355,6 +364,12 @@ async def test_langfuse_logging_completion_with_bedrock_llm_response( langfuse_export: tuple[respx.Route, str], ) -> None: route, trace_id = langfuse_export + profile: Final = litellm.get_model_info( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", custom_llm_provider="bedrock" + ) + input_rate: Final = profile["input_cost_per_token"] + output_rate: Final = profile["output_cost_per_token"] + assert isinstance(input_rate, float) and isinstance(output_rate, float) response: Final = litellm.ModelResponse( choices=[], usage=litellm.Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20), @@ -373,7 +388,9 @@ async def test_langfuse_logging_completion_with_bedrock_llm_response( ) await _flush_langfuse() assert route.call_count == 1 - _expected_generation(route, "completion_with_bedrock_call.json", trace_id) + _expected_generation( + route, "completion_with_bedrock_call.json", trace_id, cost_total=10 * input_rate + 10 * output_rate + ) @pytest.mark.asyncio diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 71efa1352e4..23748790869 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -2,7 +2,7 @@ import json import asyncio import datetime import time -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping from pathlib import Path from types import MappingProxyType, SimpleNamespace from typing import Final, Literal, cast @@ -526,6 +526,204 @@ def test_completion_cost_bills_priced_response_model_over_priced_requested_model assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["openai/sol-requested-test-2099-01-01"], usage)) +_BEDROCK_PROFILE_GLOBAL_RATES: Final = MappingProxyType( + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 5e-6, + "cache_read_input_token_cost": 1e-7, + "cache_creation_input_token_cost": 1.25e-6, + } +) +_BEDROCK_PROFILE_MODEL: Final = "anthropic.bedrock-profile-test-v1:0" +_BEDROCK_PROFILE_RATE_MULTIPLIERS: Final = MappingProxyType( + { + _BEDROCK_PROFILE_MODEL: 1.0, + f"eu.{_BEDROCK_PROFILE_MODEL}": 1.1, + f"us.{_BEDROCK_PROFILE_MODEL}": 1.2, + f"jp.{_BEDROCK_PROFILE_MODEL}": 1.3, + f"global.{_BEDROCK_PROFILE_MODEL}": 0.9, + "anthropic.bedrock-profile-other-test-v1:0": 3.0, + "anthropic.bedrock-profile-unpriced-test-v1:0": 1.0, + "bedrock-profile-router-test": 5.0, + } +) +_BEDROCK_PROFILE_CACHE_USAGE: Final = MappingProxyType( + {"cache_creation_input_tokens": 2000, "cache_read_input_tokens": 24000} +) + + +@pytest.fixture +def _bedrock_profile_rows(_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch) -> None: + for key, multiplier in _BEDROCK_PROFILE_RATE_MULTIPLIERS.items(): + monkeypatch.setitem( + litellm.model_cost, + key, + { + "litellm_provider": "bedrock_converse", + "mode": "chat", + **{field: rate * multiplier for field, rate in _BEDROCK_PROFILE_GLOBAL_RATES.items()}, + }, + ) + monkeypatch.setitem( + litellm.model_cost, + "eu.anthropic.bedrock-profile-unpriced-test-v1:0", + {"litellm_provider": "bedrock_converse", "mode": "chat"}, + ) + + +def _bedrock_messages_body(model: str) -> dict[str, object]: + return { + "id": "msg_bdrk_profile", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 18, "output_tokens": 4, **_BEDROCK_PROFILE_CACHE_USAGE}, + } + + +def _bedrock_stream_response(model: str) -> ModelResponse: + cached: Final = _BEDROCK_PROFILE_CACHE_USAGE["cache_read_input_tokens"] + written: Final = _BEDROCK_PROFILE_CACHE_USAGE["cache_creation_input_tokens"] + return ModelResponse( + model=model, + choices=[], + usage=Usage( + prompt_tokens=18 + cached + written, + completion_tokens=4, + total_tokens=22 + cached + written, + cache_creation_input_tokens=written, + cache_read_input_tokens=cached, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached), + ), + ) + + +def _bedrock_private_route_response(model: str) -> ModelResponse: + response: Final = _bedrock_stream_response("bedrock-profile-alias") + response._hidden_params = {"custom_llm_provider": "bedrock", "provider_response_model": model} + return response + + +def _bedrock_global_cost(served_shape: Callable[[str], object]) -> float: + return completion_cost( + completion_response=served_shape(_BEDROCK_PROFILE_MODEL), + model=f"bedrock/{_BEDROCK_PROFILE_MODEL}", + custom_llm_provider="bedrock", + ) + + +_BEDROCK_SERVED_SHAPES: Final = ( + pytest.param(_bedrock_messages_body, id="non-streaming"), + pytest.param(_bedrock_stream_response, id="streaming"), + pytest.param(_bedrock_private_route_response, id="private-provider-response-model"), +) + + +@pytest.mark.parametrize("served_shape", _BEDROCK_SERVED_SHAPES) +@pytest.mark.parametrize("route", ["bedrock/", "bedrock/invoke/"]) +@pytest.mark.parametrize("region", ["eu", "us", "jp"]) +def test_completion_cost_bills_bedrock_cross_region_profile_when_bedrock_serves_the_region_stripped_id( + _bedrock_profile_rows: None, region: str, route: str, served_shape: Callable[[str], object] +) -> None: + profile: Final = f"{region}.{_BEDROCK_PROFILE_MODEL}" + + cost: Final = completion_cost( + completion_response=served_shape(_BEDROCK_PROFILE_MODEL), + model=f"{route}{profile}", + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx(_bedrock_global_cost(served_shape) * _BEDROCK_PROFILE_RATE_MULTIPLIERS[profile]) + + +@pytest.mark.parametrize("served_shape", _BEDROCK_SERVED_SHAPES) +@pytest.mark.parametrize( + ("deployment", "served_model"), + [ + pytest.param(f"bedrock/eu.{_BEDROCK_PROFILE_MODEL}", f"us.{_BEDROCK_PROFILE_MODEL}", id="other-region"), + pytest.param(f"bedrock/eu.{_BEDROCK_PROFILE_MODEL}", f"global.{_BEDROCK_PROFILE_MODEL}", id="global-profile"), + pytest.param( + f"bedrock/jp.{_BEDROCK_PROFILE_MODEL}", "anthropic.bedrock-profile-other-test-v1:0", id="other-model" + ), + pytest.param( + "bedrock/eu.anthropic.bedrock-profile-unpriced-test-v1:0", + "anthropic.bedrock-profile-unpriced-test-v1:0", + id="unpriced-profile", + ), + ], +) +def test_completion_cost_keeps_the_served_bedrock_row_unless_it_is_the_region_stripped_profile( + _bedrock_profile_rows: None, deployment: str, served_model: str, served_shape: Callable[[str], object] +) -> None: + cost: Final = completion_cost( + completion_response=served_shape(served_model), + model=deployment, + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx(_bedrock_global_cost(served_shape) * _BEDROCK_PROFILE_RATE_MULTIPLIERS[served_model]) + + +@pytest.mark.parametrize( + ("base_model", "custom_pricing", "router_model_id", "priced_row"), + [ + pytest.param( + "bedrock/anthropic.bedrock-profile-other-test-v1:0", + None, + None, + "anthropic.bedrock-profile-other-test-v1:0", + id="base-model", + ), + pytest.param(None, True, "bedrock-profile-router-test", "bedrock-profile-router-test", id="custom-pricing"), + ], +) +def test_completion_cost_explicit_pricing_precedes_the_bedrock_cross_region_profile( + _bedrock_profile_rows: None, + base_model: str | None, + custom_pricing: bool | None, + router_model_id: str | None, + priced_row: str, +) -> None: + cost: Final = completion_cost( + completion_response=_bedrock_messages_body(_BEDROCK_PROFILE_MODEL), + model=f"bedrock/eu.{_BEDROCK_PROFILE_MODEL}", + custom_llm_provider="bedrock", + base_model=base_model, + custom_pricing=custom_pricing, + router_model_id=router_model_id, + ) + + assert cost == pytest.approx( + _bedrock_global_cost(_bedrock_messages_body) * _BEDROCK_PROFILE_RATE_MULTIPLIERS[priced_row] + ) + + +def test_completion_cost_keeps_the_served_row_for_a_region_dotted_name_outside_bedrock( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "openai/eu.profile-dotted-test", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + monkeypatch.setitem( + litellm.model_cost, + "openai/profile-dotted-test", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 5e-6, "output_cost_per_token": 11e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + cost: Final = completion_cost( + completion_response=ModelResponse(model="profile-dotted-test", choices=[], usage=usage), + model="openai/eu.profile-dotted-test", + custom_llm_provider="openai", + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["openai/profile-dotted-test"], usage)) + + def test_completion_cost_keeps_response_pricing_when_requested_alias_has_no_rates( _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch ) -> None: