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: