mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
fix(bedrock): bill the regional row for inference-profile ARNs that embed a cross-region prefix (#46128)
* fix(bedrock): bill the regional row for inference-profile ARNs that embed a cross-region prefix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(bedrock): keep the ARN regional-row regression tests to the touched lines Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
cca06b27e4
commit
7265b0e889
3 changed files with 150 additions and 3 deletions
|
|
@ -5912,9 +5912,18 @@ def get_potential_model_names(model: str, custom_llm_provider: str | None) -> Po
|
|||
provider_prefixed_model_name = combined_model_name
|
||||
|
||||
if custom_llm_provider in ("bedrock", "bedrock_converse"):
|
||||
from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix
|
||||
from litellm.llms.bedrock.common_utils import (
|
||||
extract_model_name_from_bedrock_arn,
|
||||
split_bedrock_region_path,
|
||||
strip_bedrock_routing_prefix,
|
||||
)
|
||||
|
||||
split_model = strip_bedrock_routing_prefix(split_model)
|
||||
region_free_bedrock_model: Final = split_bedrock_region_path(split_model)[1]
|
||||
split_model = (
|
||||
extract_model_name_from_bedrock_arn(region_free_bedrock_model)
|
||||
if region_free_bedrock_model.startswith("arn:")
|
||||
else strip_bedrock_routing_prefix(split_model)
|
||||
)
|
||||
|
||||
region_free_split_model: Final = (
|
||||
_strip_mantle_region_prefix(split_model) if custom_llm_provider == "bedrock_mantle" else split_model
|
||||
|
|
|
|||
|
|
@ -16,7 +16,6 @@ from integration.cost_calculation.cost_tracking_case import (
|
|||
)
|
||||
from integration.cost_calculation.stream_parity.case import StreamParityTestCase
|
||||
|
||||
|
||||
# anthropic.claude-sonnet-5-v1:0
|
||||
ANTHROPIC_CLAUDE_SONNET_5_V1_0_INPUT_TEXT: Final = CostTrackingTestCase(
|
||||
name="anthropic.claude-sonnet-5-v1:0-input_text",
|
||||
|
|
@ -1096,6 +1095,119 @@ BEDROCK_CONVERSE_APAC_BARE_FALLBACK: Final = CostTrackingTestCase(
|
|||
)
|
||||
|
||||
|
||||
BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY: Final = CostTrackingTestCase(
|
||||
name="bedrock-converse-profile-arn-regional-key",
|
||||
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
model="eu.anthropic.claude-sonnet-5-v1:0",
|
||||
deployment=Deployment(
|
||||
model="bedrock/converse/arn:aws:bedrock:eu-west-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-5-v1:0"
|
||||
),
|
||||
request={
|
||||
"model": "$MODEL",
|
||||
"messages": [{"role": "user", "content": "bedrock-converse-profile-arn-regional-key"}],
|
||||
"stream": False,
|
||||
},
|
||||
response=JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "scripted response"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252},
|
||||
"metrics": {"latencyMs": 1},
|
||||
},
|
||||
),
|
||||
expected=ExactExpected(
|
||||
spend=0.01326, input_cost=0.006256, output_cost=0.007004, prompt_tokens=1840, completion_tokens=412
|
||||
),
|
||||
)
|
||||
|
||||
BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_STREAM: Final = CostTrackingTestCase(
|
||||
name="bedrock-converse-profile-arn-regional-key-stream",
|
||||
covers="quota_management.spend_tracking.scripted_wire.logs_cost",
|
||||
model="eu.anthropic.claude-sonnet-5-v1:0",
|
||||
deployment=Deployment(
|
||||
model="bedrock/converse/arn:aws:bedrock:eu-west-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-5-v1:0"
|
||||
),
|
||||
request={
|
||||
"model": "$MODEL",
|
||||
"messages": [{"role": "user", "content": "bedrock-converse-profile-arn-regional-key-stream"}],
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
},
|
||||
response=EventStreamResponse(
|
||||
content_type="application/vnd.amazon.eventstream",
|
||||
events=(
|
||||
EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}),
|
||||
EventStreamEvent(
|
||||
event_type="contentBlockDelta", payload={"delta": {"text": "scripted response"}, "contentBlockIndex": 0}
|
||||
),
|
||||
EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}),
|
||||
EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}),
|
||||
EventStreamEvent(
|
||||
event_type="metadata",
|
||||
payload={
|
||||
"usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252},
|
||||
"metrics": {"latencyMs": 1},
|
||||
},
|
||||
),
|
||||
),
|
||||
),
|
||||
expected=ExactExpected(
|
||||
spend=0.01326, input_cost=0.006256, output_cost=0.007004, prompt_tokens=1840, completion_tokens=412
|
||||
),
|
||||
)
|
||||
|
||||
BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_MESSAGES: Final = CostTrackingTestCase(
|
||||
name="bedrock-converse-profile-arn-regional-key-messages",
|
||||
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
model="eu.anthropic.claude-sonnet-5-v1:0",
|
||||
endpoint="/v1/messages",
|
||||
deployment=Deployment(
|
||||
model="bedrock/converse/arn:aws:bedrock:eu-west-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-5-v1:0"
|
||||
),
|
||||
request={
|
||||
"model": "$MODEL",
|
||||
"max_tokens": 412,
|
||||
"messages": [{"role": "user", "content": "bedrock-converse-profile-arn-regional-key-messages"}],
|
||||
},
|
||||
response=JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "scripted response"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252},
|
||||
"metrics": {"latencyMs": 1},
|
||||
},
|
||||
),
|
||||
expected=ExactExpected(
|
||||
spend=0.01326, input_cost=0.006256, output_cost=0.007004, prompt_tokens=1840, completion_tokens=412
|
||||
),
|
||||
)
|
||||
|
||||
BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_RESPONSES: Final = CostTrackingTestCase(
|
||||
name="bedrock-converse-profile-arn-regional-key-responses",
|
||||
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
model="eu.anthropic.claude-sonnet-5-v1:0",
|
||||
endpoint="/v1/responses",
|
||||
deployment=Deployment(
|
||||
model="bedrock/converse/arn:aws:bedrock:eu-west-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-5-v1:0"
|
||||
),
|
||||
request={"model": "$MODEL", "input": "bedrock-converse-profile-arn-regional-key-responses"},
|
||||
response=JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "scripted response"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 1840, "outputTokens": 412, "totalTokens": 2252},
|
||||
"metrics": {"latencyMs": 1},
|
||||
},
|
||||
),
|
||||
expected=ExactExpected(
|
||||
spend=0.01326, input_cost=0.006256, output_cost=0.007004, prompt_tokens=1840, completion_tokens=412
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# meta.llama4-maverick-17b-instruct-v1:0
|
||||
META_LLAMA4_MAVERICK_17B_INSTRUCT_V1_0_INPUT_TEXT: Final = CostTrackingTestCase(
|
||||
name="meta.llama4-maverick-17b-instruct-v1:0-input_text",
|
||||
|
|
@ -4798,6 +4910,10 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = (
|
|||
BEDROCK_CONVERSE_PROFILE_BASE_MODEL,
|
||||
BEDROCK_CONVERSE_EU_REGIONAL_KEY,
|
||||
BEDROCK_CONVERSE_APAC_BARE_FALLBACK,
|
||||
BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY,
|
||||
BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_STREAM,
|
||||
BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_MESSAGES,
|
||||
BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_RESPONSES,
|
||||
BEDROCK_CONVERSE_NOVA_2_PRO,
|
||||
BEDROCK_CONVERSE_MISTRAL_LARGE_3_STREAM,
|
||||
US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT,
|
||||
|
|
@ -4838,6 +4954,9 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = (
|
|||
)
|
||||
|
||||
PARITY: Final[tuple[StreamParityTestCase, ...]] = (
|
||||
StreamParityTestCase(
|
||||
plain=BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY, streamed=BEDROCK_CONVERSE_PROFILE_ARN_REGIONAL_KEY_STREAM
|
||||
),
|
||||
StreamParityTestCase(
|
||||
plain=ANTHROPIC_CLAUDE_SONNET_5_V1_0_TOOL_CALL, streamed=ANTHROPIC_CLAUDE_SONNET_5_V1_0_STREAM_TOOL_CALL
|
||||
),
|
||||
|
|
|
|||
|
|
@ -1425,6 +1425,25 @@ def test_get_model_info_bedrock_regional_inference_profile_pricing(local_model_c
|
|||
assert control["key"] == "eu.amazon.nova-pro-v1:0"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ["bedrock/", "bedrock/converse/", "bedrock/invoke/", "bedrock/eu-west-1/"])
|
||||
def test_get_model_info_bedrock_inference_profile_arn_keeps_regional_row(local_model_cost_map, route):
|
||||
regional = litellm.model_cost["eu.amazon.nova-pro-v1:0"]
|
||||
base = litellm.model_cost["amazon.nova-pro-v1:0"]
|
||||
assert regional["input_cost_per_token"] > base["input_cost_per_token"]
|
||||
arn_prefix = "arn:aws:bedrock:eu-west-1:123456789012:inference-profile/"
|
||||
|
||||
info = litellm.get_model_info(model=f"{route}{arn_prefix}eu.amazon.nova-pro-v1:0")
|
||||
assert info["key"] == "eu.amazon.nova-pro-v1:0"
|
||||
assert info["input_cost_per_token"] == regional["input_cost_per_token"]
|
||||
assert info["output_cost_per_token"] == regional["output_cost_per_token"]
|
||||
|
||||
bare = litellm.get_model_info(model=f"{route}{arn_prefix}amazon.nova-pro-v1:0")
|
||||
assert bare["key"] == "amazon.nova-pro-v1:0"
|
||||
assert bare["input_cost_per_token"] == base["input_cost_per_token"]
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bare_key",
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue