From dfac6289e41f03c83d699cf2d09d09dd7cd6df98 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:23:14 -0700 Subject: [PATCH] fix(bedrock): resolve the region prefix the same way on every Bedrock surface One region precedence for every call: aws_region_name, then the ARN region (model_id or model), then the region prefix in the model id, then env, with an empty aws_region_name treated as unset. The converse handler stops injecting the prefix region itself, the invoke config inherits the BaseAWSLLM provider resolver and transforms with the bare model, and count_tokens, batch jobs, realtime sessions, and the OpenAI-compatible invoke route send the bare id to AWS. get_bedrock_base_model strips the routing prefix, the region prefix, the ARN, the throughput suffix, and the cross-region prefix in one pass, and the cost lookup keeps the us./eu. profile row for a region-prefixed profile id. mx-central-1 leaves the region list because Bedrock does not run there. --- basedpyright-code-budget.json | 4 +- litellm/llms/bedrock/base_aws_llm.py | 11 +-- .../llms/bedrock/batches/transformation.py | 3 +- litellm/llms/bedrock/chat/converse_handler.py | 3 - .../amazon_openai_transformation.py | 10 +-- .../base_invoke_transformation.py | 73 +++---------------- litellm/llms/bedrock/common_utils.py | 19 ++--- .../count_tokens/bedrock_token_counter.py | 7 +- litellm/llms/bedrock/realtime/handler.py | 4 +- litellm/utils.py | 4 +- ruff-strict-budget.json | 6 +- .../bedrock/batches/test_transformation.py | 24 ++++++ .../test_amazon_openai_transformation.py | 39 ++++++++++ .../test_base_invoke_transformation.py | 12 +++ .../test_bedrock_token_counter.py | 35 +++++++++ .../realtime/test_bedrock_realtime_handler.py | 19 +++++ .../rerank/test_bedrock_rerank_handler.py | 2 +- .../llms/bedrock/test_base_aws_llm.py | 17 ++++- .../llms/bedrock/test_bedrock_common_utils.py | 17 +++++ tests/test_litellm/test_utils.py | 9 +++ type-discipline-budget.json | 10 +-- 21 files changed, 210 insertions(+), 118 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_openai_transformation.py create mode 100644 tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_token_counter.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index b876cf2d69a..8bc07b16195 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 13429 + "limit": 13427 }, "reportArgumentType": { "limit": 2198 @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1808 + "limit": 1807 }, "reportRedeclaration": { "limit": 8 diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 26ac5fd249e..aab1eb3dac2 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -582,12 +582,8 @@ class BaseAWSLLM: Returns: str: The AWS region name """ - aws_region_name = optional_params.get("aws_region_name", None) + aws_region_name = optional_params.get("aws_region_name") or None self._validate_aws_region_name(aws_region_name) - if aws_region_name is None and model: - from litellm.llms.bedrock.common_utils import split_bedrock_region_prefix, strip_bedrock_routing_prefix - - aws_region_name = split_bedrock_region_prefix(strip_bedrock_routing_prefix(model))[0] ### SET REGION NAME ### if aws_region_name is None: # check model arn # @@ -595,6 +591,11 @@ class BaseAWSLLM: aws_region_name = self._get_aws_region_from_model_arn(model_id) else: aws_region_name = self._get_aws_region_from_model_arn(model) + if aws_region_name is None and model: + from litellm.llms.bedrock.common_utils import split_bedrock_region_prefix, strip_bedrock_routing_prefix + + aws_region_name = split_bedrock_region_prefix(strip_bedrock_routing_prefix(model))[0] + if aws_region_name is None: # check env # litellm_aws_region_name: Final = get_secret("AWS_REGION_NAME", None) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 7729cdfdb0d..b76f10e2ca5 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -33,6 +33,7 @@ from ..common_utils import ( CommonBatchFilesUtils, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id, + split_bedrock_region_prefix, ) if TYPE_CHECKING: @@ -221,7 +222,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # Create Bedrock batch request with proper typing bedrock_request: Final[BedrockCreateBatchRequest] = { - "modelId": model, + "modelId": split_bedrock_region_prefix(model)[1], "jobName": job_name, "inputDataConfig": input_data_config, "outputDataConfig": output_data_config, diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 544eea0fbcc..841e8061bee 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -308,9 +308,6 @@ class BedrockConverseLLM(BaseAWSLLM): _model_for_id = _model_for_id.replace(_nova_prefix, "", 1) break modelId = self.encode_model_id(model_id=_model_for_id) - # Inject region extracted from model path so _get_aws_region_name picks it up - if _region_from_model is not None and not optional_params.get("aws_region_name"): - optional_params["aws_region_name"] = _region_from_model fake_stream = litellm.AmazonConverseConfig().should_fake_stream( fake_stream=fake_stream, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index 1671585be2d..d072bab300e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM -from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.common_utils import BedrockError, split_bedrock_region_prefix from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, merge_bedrock_invoke_headers, @@ -59,13 +59,7 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): Input format: bedrock/openai/ Returns: """ - # Remove bedrock/ prefix if present - model = model.removeprefix("bedrock/") - - # Remove openai/ prefix - model = model.removeprefix("openai/") - - return model + return split_bedrock_region_prefix(model.removeprefix("bedrock/").removeprefix("openai/"))[1] def get_complete_url( self, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 37121d2ece7..5a4912bfdce 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -1,7 +1,7 @@ import copy import json import time -from typing import TYPE_CHECKING, Any, Final, cast, get_args +from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import TypeAdapter, ValidationError @@ -18,7 +18,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call -from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.common_utils import BedrockError, split_bedrock_region_prefix from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, merge_bedrock_invoke_headers, @@ -184,11 +184,12 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): optional_params.pop("stream_chunk_size", None) custom_prompt_dict: Final[dict] = litellm_params.pop("custom_prompt_dict", None) or {} hf_model_name: Final = litellm_params.get("hf_model_name", None) + bare_model: Final = split_bedrock_region_prefix(model.removeprefix("invoke/"))[1] provider: Final = self.get_bedrock_invoke_provider(model) prompt, chat_history = self.convert_messages_to_prompt( - model=hf_model_name or model, + model=hf_model_name or bare_model, messages=messages, provider=provider, custom_prompt_dict=custom_prompt_dict, @@ -197,7 +198,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): inference_params = {k: v for k, v in inference_params.items() if k not in self.aws_authentication_params} request_data: dict = {} if provider == "cohere": - if model.startswith("cohere.command-r"): + if bare_model.startswith("cohere.command-r"): ## LOAD CONFIG config = litellm.AmazonCohereChatConfig().get_config() self._apply_config_to_params(config, inference_params) @@ -214,7 +215,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): request_data = {"prompt": prompt, **inference_params} elif provider == "anthropic": transformed_request: Final = litellm.AmazonAnthropicClaudeConfig().transform_request( - model=model, + model=bare_model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, @@ -224,7 +225,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): return transformed_request elif provider == "nova": return litellm.AmazonInvokeNovaConfig().transform_request( - model=model, + model=bare_model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, @@ -255,7 +256,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): request_data = {"prompt": prompt, **inference_params} elif provider == "twelvelabs": return litellm.AmazonTwelveLabsPegasusConfig().transform_request( - model=model, + model=bare_model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, @@ -264,7 +265,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): elif provider == "openai": # OpenAI imported models use OpenAI Chat Completions format return litellm.AmazonBedrockOpenAIConfig().transform_request( - model=model, + model=bare_model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, @@ -518,62 +519,6 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): """ return False - @staticmethod - def get_bedrock_invoke_provider( - model: str, - ) -> litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None: - """ - Helper function to get the bedrock provider from the model - - handles 4 scenarios: - 1. model=invoke/anthropic.claude-3-5-sonnet-20240620-v1:0 -> Returns `anthropic` - 2. model=anthropic.claude-3-5-sonnet-20240620-v1:0 -> Returns `anthropic` - 3. model=llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n -> Returns `llama` - 4. model=us.amazon.nova-pro-v1:0 -> Returns `nova` - """ - if model.startswith("invoke/"): - model = model.replace("invoke/", "", 1) - - # Special case: Check for "nova" in model name first (before "amazon") - # This handles amazon.nova-* models which would otherwise match "amazon" (Titan) - if "nova" in model.lower(): - if "nova" in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova") - - _split_model: Final = model.split(".")[0] - if _split_model in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model) - - # If not a known provider, check for pattern with two slashes - provider = AmazonInvokeConfig._get_provider_from_model_path(model) - if provider is not None: - return provider - - for provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - if provider in model: - return provider - return None - - @staticmethod - def _get_provider_from_model_path( - model_path: str, - ) -> litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None: - """ - Helper function to get the provider from a model path with format: provider/model-name - - Args: - model_path (str): The model path (e.g., 'llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n' or 'anthropic/model-name') - - Returns: - Optional[str]: The provider name, or None if no valid provider found - """ - parts: Final = model_path.split("/") - if len(parts) >= 1: - provider: Final = parts[0] - if provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): - return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider) - return None - def convert_messages_to_prompt(self, model, messages, provider, custom_prompt_dict) -> tuple[str, list | None]: # handle anthropic prompts and amazon titan prompts prompt = "" diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 4a3748f2c46..901a3d4ed31 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -402,7 +402,6 @@ class AmazonBedrockGlobalConfig: *self.get_sa_regions(), *self.get_me_regions(), *self.get_af_regions(), - *self.get_mx_regions(), ) def get_ap_regions(self) -> list[str]: @@ -434,9 +433,6 @@ class AmazonBedrockGlobalConfig: def get_af_regions(self) -> tuple[str, ...]: return ("af-south-1",) - def get_mx_regions(self) -> tuple[str, ...]: - return ("mx-central-1",) - def get_eu_regions(self) -> list[str]: """ Source: https://www.aws-services.info/bedrock.html @@ -831,19 +827,14 @@ def get_bedrock_base_model(model: str) -> str: elif stripped.startswith("nova/"): return "amazon.nova-custom" - model = strip_bedrock_routing_prefix(model) - model = extract_model_name_from_bedrock_arn(model) - model = strip_bedrock_throughput_suffix(model) - - potential_region: Final = model.split(".", 1)[0] - alt_potential_region: Final = model.split("/", 1)[0] + region_free: Final = split_bedrock_region_prefix(strip_bedrock_routing_prefix(model))[1] + bare: Final = strip_bedrock_throughput_suffix(extract_model_name_from_bedrock_arn(region_free)) + potential_region: Final = bare.split(".", 1)[0] if potential_region in get_bedrock_cross_region_inference_regions(): - return model.split(".", 1)[1] - elif alt_potential_region in _get_all_bedrock_regions() and len(model.split("/", 1)) > 1: - return model.split("/", 1)[1] + return bare.split(".", 1)[1] - return model + return bare def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 9c5211ed072..172e5701d71 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -6,7 +6,7 @@ from typing import Any, Final from litellm._logging import verbose_logger from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.types.utils import LlmProviders, TokenCountResponse @@ -67,15 +67,12 @@ class BedrockTokenCounter(BaseTokenCounter): if system: request_data["system"] = system - # Get the resolved model (strip prefixes like bedrock/, converse/, etc.) - resolved_model: Final = get_bedrock_base_model(model_to_use) - try: handler: Final = BedrockCountTokensHandler() result: Final = await handler.handle_count_tokens_request( request_data=request_data, litellm_params=litellm_params, - resolved_model=resolved_model, + resolved_model=model_to_use, ) # Transform response to TokenCountResponse diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 42fe8941443..968587a16a9 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -22,7 +22,7 @@ from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError +from ..common_utils import BedrockError, split_bedrock_region_prefix from .transformation import BedrockRealtimeConfig _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) @@ -184,7 +184,7 @@ class BedrockRealtime(BaseAWSLLM): async def open_bidirectional_stream() -> BedrockBidirectionalStream: return await bedrock_client.invoke_model_with_bidirectional_stream( - InvokeModelWithBidirectionalStreamOperationInput(model_id=model) + InvokeModelWithBidirectionalStreamOperationInput(model_id=split_bedrock_region_prefix(model)[1]) ) transformation_config: Final = BedrockRealtimeConfig() diff --git a/litellm/utils.py b/litellm/utils.py index 9d20d32d147..7bbb920fbc9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5482,9 +5482,9 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P 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 split_bedrock_region_prefix, strip_bedrock_routing_prefix - split_model = strip_bedrock_routing_prefix(split_model) + split_model = split_bedrock_region_prefix(strip_bedrock_routing_prefix(split_model))[1] return PotentialModelNamesAndCustomLLMProvider( split_model=split_model, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc314..c0892455c38 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -168,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 173 + "limit": 172 }, "RUF012": { "limit": 239 @@ -198,7 +198,7 @@ "limit": 56 }, "SIM102": { - "limit": 310 + "limit": 309 }, "SIM103": { "limit": 119 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1035 + "limit": 1034 }, "TRY002": { "limit": 524 diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 7e5716a7495..969cba9b837 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -16,6 +16,8 @@ we mock; everything else runs for real. from unittest.mock import MagicMock, patch +import json + import httpx import pytest @@ -819,3 +821,25 @@ def test_retrieve_request_accepts_partition_arns(config: BedrockBatchesConfig, a batch_id=arn, optional_params={}, litellm_params={} ) assert result["url"].startswith(expected_prefix) + + +def test_create_request_targets_the_prefix_region_with_the_bare_model_id(config, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + result = config.transform_create_batch_request( + model="eu-central-1/anthropic.claude-3-5-sonnet", + create_batch_data={ + "input_file_id": "s3://in-bucket/path/to/input.jsonl", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + optional_params={}, + litellm_params={ + "s3_output_bucket_name": "out-bucket", + "aws_batch_role_arn": "arn:aws:iam::123:role/my-batch-role", + "aws_access_key_id": "fake-access-key", + "aws_secret_access_key": "fake-secret-key", + }, + ) + + assert result["url"] == "https://bedrock.eu-central-1.amazonaws.com/model-invocation-job" + assert json.loads(result["data"])["modelId"] == "anthropic.claude-3-5-sonnet" diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_openai_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_openai_transformation.py new file mode 100644 index 00000000000..a1539c90560 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_openai_transformation.py @@ -0,0 +1,39 @@ +from typing import Final + +import pytest + +from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, +) + +IMPORTED_MODEL_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123" + + +@pytest.mark.parametrize( + "model", + [ + f"bedrock/openai/us-east-1/{IMPORTED_MODEL_ARN}", + f"openai/us-east-1/{IMPORTED_MODEL_ARN}", + f"bedrock/openai/{IMPORTED_MODEL_ARN}", + ], +) +def test_region_prefixed_imported_model_invokes_the_bare_arn_in_the_prefix_region(monkeypatch, model): + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + config: Final = AmazonBedrockOpenAIConfig() + + url: Final = config.get_complete_url( + api_base=None, api_key=None, model=model, optional_params={}, litellm_params={}, stream=False + ) + body: Final = config.transform_request( + model=model, + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert url == ( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/" + "arn:aws:bedrock:us-east-1:123456789012:imported-model%2Fabc123/invoke" + ) + assert body["model"] == IMPORTED_MODEL_ARN diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index 2ca84806e91..1d5f72aba8f 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -233,3 +233,15 @@ def test_region_prefixed_invoke_model_calls_that_region_with_the_bare_model_id( monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) assert _invoke_with_mocked_post(model, aws_region_name) == expected_url + + +def test_transform_request_builds_the_command_r_chat_body_for_a_region_prefixed_model(): + request = AmazonInvokeConfig().transform_request( + model="us-east-1/cohere.command-r-plus-v1:0", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"max_tokens": 16}, + litellm_params={}, + headers={}, + ) + + assert request == {"message": "Hello", "max_tokens": 16, "chat_history": []} diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_token_counter.py new file mode 100644 index 00000000000..01dbf4e6e78 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -0,0 +1,35 @@ +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter + +COUNT_TOKENS_URL_IN_THE_PREFIX_REGION: Final = ( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-haiku-4-5-20251001-v1%3A0/count-tokens" +) + + +@pytest.mark.asyncio +async def test_count_tokens_calls_the_region_named_in_the_model_prefix_with_the_bare_model_id( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + with respx.mock: + respx.post(COUNT_TOKENS_URL_IN_THE_PREFIX_REGION).mock( + return_value=httpx.Response(200, json={"inputTokens": 7}) + ) + result = await BedrockTokenCounter().count_tokens( + model_to_use="us-east-1/us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"aws_access_key_id": "fake", "aws_secret_access_key": "fake"}}, + ) + + assert result is not None + assert result.total_tokens == 7 diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 0ea5b7ad4a1..74f794841fb 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -492,6 +492,25 @@ class TestBedrockRealtimeAwsAuth: assert stub_aws_sdk_client["operation_input"].model_id == "amazon.nova-sonic-v1:0" assert websocket.closed + @pytest.mark.asyncio + async def test_region_prefixed_model_opens_the_stream_in_that_region_with_the_bare_model_id( + self, stub_aws_sdk_client, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + handler = BedrockRealtime() + websocket = RealtimeClientWS() + + await handler.async_realtime( + model="eu-north-1/amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=FakeLogging(), + aws_access_key_id="litellm-params-access-key", + aws_secret_access_key="litellm-params-secret-key", + ) + + assert stub_aws_sdk_client["config_kwargs"]["region"] == "eu-north-1" + assert stub_aws_sdk_client["operation_input"].model_id == "amazon.nova-sonic-v1:0" + @pytest.mark.asyncio async def test_role_assumption_params_forwarded_to_get_credentials(self, stub_aws_sdk_client): handler = StubCredentialsBedrockRealtime( diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_handler.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_handler.py index 667a4058107..d47d57b06ba 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_handler.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_handler.py @@ -8,7 +8,7 @@ import litellm from litellm.llms.custom_httpx.http_handler import HTTPHandler -def _rerank_with_mocked_post(model: str) -> tuple[str, dict]: +def _rerank_with_mocked_post(model: str) -> tuple[str, dict[str, object]]: body = {"results": [{"index": 0, "relevanceScore": 0.9}], "usage": {"search_units": 1}} mock_response = MagicMock() mock_response.status_code = 200 diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 5be9f4d0676..663353ad46b 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -505,7 +505,6 @@ def test_get_aws_region_name_boto3_fallback(): "us-east-1\\path", "US-EAST-1", # uppercase not allowed "us east 1", # spaces not allowed - "", # empty string not allowed "us-east-1\n", # trailing newline must not slip past $ ], ) @@ -546,13 +545,14 @@ def test_get_aws_region_name_accepts_valid_regions(valid_region): "model,explicit_region,expected_region", [ ("bedrock/invoke/eu-west-1/mistral.mistral-7b-instruct-v0:2", None, "eu-west-1"), + ("eu-west-1/mistral.mistral-7b-instruct-v0:2", "", "eu-west-1"), ("eu-west-1/mistral.mistral-7b-instruct-v0:2", "ap-south-1", "ap-south-1"), - ("eu-west-1/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0", None, "eu-west-1"), + ("eu-west-1/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0", None, "us-east-1"), ("arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0", None, "us-east-1"), ("mistral.mistral-7b-instruct-v0:2", None, "us-west-2"), ], ) -def test_get_aws_region_name_prefers_explicit_region_then_model_prefix_then_arn_then_env( +def test_get_aws_region_name_prefers_explicit_region_then_arn_then_model_prefix_then_env( monkeypatch: pytest.MonkeyPatch, model: str, explicit_region: str | None, expected_region: str ): monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") @@ -561,6 +561,17 @@ def test_get_aws_region_name_prefers_explicit_region_then_model_prefix_then_arn_ assert BaseAWSLLM()._get_aws_region_name(optional_params, model=model) == expected_region +def test_get_aws_region_name_prefers_the_model_id_arn_region_over_the_model_prefix(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_REGION_NAME", "us-west-2") + profile_arn = "arn:aws:bedrock:eu-central-1:123456789012:application-inference-profile/abc123" + + region = BaseAWSLLM()._get_aws_region_name( + {}, model="us-east-1/anthropic.claude-haiku-4-5-20251001-v1:0", model_id=profile_arn + ) + + assert region == "eu-central-1" + + @pytest.mark.parametrize( "model,expected_model_id", [ diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index dfff00405a2..f9695c5c4c2 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -647,3 +647,20 @@ def test_every_region_prefix_the_cost_map_prices_is_a_known_bedrock_region(): assert prefixed_keys unrecognized = sorted({key.split("/")[0] for key in prefixed_keys if split_bedrock_region_prefix(key)[0] is None}) assert unrecognized == [] + + +@pytest.mark.parametrize( + "model,expected", + [ + ("bedrock/us-east-1/us.anthropic.claude-haiku-4-5-20251001-v1:0", "anthropic.claude-haiku-4-5-20251001-v1:0"), + ("eu-west-1/amazon.titan-embed-text-v2:0", "amazon.titan-embed-text-v2:0"), + ( + "bedrock/invoke/ap-south-1/arn:aws:bedrock:ap-south-1::foundation-model/meta.llama3-70b-instruct-v1:0", + "meta.llama3-70b-instruct-v1:0", + ), + ], +) +def test_get_bedrock_base_model_strips_the_region_prefix_before_the_cross_region_prefix(model, expected): + from litellm.llms.bedrock.common_utils import get_bedrock_base_model + + assert get_bedrock_base_model(model) == expected diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 14907e17b1b..5db1ab9a229 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -6067,3 +6067,12 @@ class TestFinalOptionalParamsLineRedaction: assert "'max_tokens': 17" in printed assert "'temperature': 0.25" in printed + + +def test_get_model_info_prices_a_region_prefixed_cross_region_profile_as_that_profile(local_model_cost_map): + prefixed = litellm.get_model_info(model="bedrock/us-east-1/us.anthropic.claude-haiku-4-5-20251001-v1:0") + plain = litellm.get_model_info(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") + + assert prefixed["key"] == plain["key"] == "us.anthropic.claude-haiku-4-5-20251001-v1:0" + assert prefixed["input_cost_per_token"] == plain["input_cost_per_token"] + assert prefixed["output_cost_per_token"] == plain["output_cost_per_token"] diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 6ac76c1ac8b..5577c1fa7bc 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22178 + "limit": 22175 }, "LIT002": { - "limit": 26744 + "limit": 26743 }, "LIT003": { "limit": 261 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1035 + "limit": 1032 }, "LIT007": { "limit": 0 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16454 + "limit": 16449 }, "LIT011": { - "limit": 5506 + "limit": 5499 }, "LIT012": { "limit": 4486