From f5c942f9e89f130506a34ef0b8a95c506d964f09 Mon Sep 17 00:00:00 2001 From: Krishnatejaswi S Date: Tue, 27 Oct 2026 22:04:07 +0530 Subject: [PATCH] fix(bedrock): resolve embedding region the same way as chat --- litellm/llms/bedrock/embed/embedding.py | 41 +++++++----------- .../bedrock/embed/test_bedrock_embedding.py | 42 +++++++++++++++++++ 2 files changed, 58 insertions(+), 25 deletions(-) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 4d27d94664f..f1e980d336b 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -20,7 +20,6 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, ) -from litellm.secret_managers.main import get_secret from litellm.types.llms.bedrock import ( AmazonEmbeddingRequest, CohereEmbeddingRequest, @@ -70,6 +69,9 @@ class BedrockEmbedding(BaseAWSLLM): self, optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place bearer_token: None = None, + *, + model: str | None = None, + model_id: str | None = None, ) -> tuple[Credentials, str]: ... @overload @@ -77,31 +79,25 @@ class BedrockEmbedding(BaseAWSLLM): self, optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place bearer_token: str, + *, + model: str | None = None, + model_id: str | None = None, ) -> tuple[None, str]: ... def _load_credentials( self, optional_params: dict, bearer_token: str | None = None, + *, + model: str | None = None, + model_id: str | None = None, ) -> tuple[Credentials | None, str]: auth_params: Final = pop_aws_auth_params(optional_params) - aws_region_name = optional_params.pop("aws_region_name", None) - - ### SET REGION NAME ### - if aws_region_name is None: - # check env # - litellm_aws_region_name: Final = get_secret("AWS_REGION_NAME", None) - - if litellm_aws_region_name is not None and isinstance(litellm_aws_region_name, str): - aws_region_name = litellm_aws_region_name - - standard_aws_region_name: Final = get_secret("AWS_REGION", None) - if standard_aws_region_name is not None and isinstance(standard_aws_region_name, str): - aws_region_name = standard_aws_region_name - - if aws_region_name is None: - aws_region_name = "us-west-2" - + aws_region_name: Final = self._get_aws_region_name( + optional_params={"aws_region_name": optional_params.pop("aws_region_name", None)}, + model=model, + model_id=model_id, + ) credentials: Final[Credentials | None] = ( None if bearer_token is not None else self.resolve_credentials(auth_params, aws_region_name) ) @@ -409,18 +405,13 @@ class BedrockEmbedding(BaseAWSLLM): litellm_params: dict, api_key: str | None = None, ) -> EmbeddingResponse: + unencoded_model_id: Final = optional_params.pop("model_id", None) or model # default to model if not passed credentials, aws_region_name = self._load_credentials( - optional_params, bearer_token=bedrock_bearer_token(api_key) + optional_params, bearer_token=bedrock_bearer_token(api_key), model=model, model_id=unencoded_model_id ) ### TRANSFORMATION ### - unencoded_model_id: Final = optional_params.pop("model_id", None) or model # default to model if not passed modelId: Final = urllib.parse.quote(unencoded_model_id, safe="") - aws_region_name = self._get_aws_region_name( - optional_params={"aws_region_name": aws_region_name}, - model=model, - model_id=unencoded_model_id, - ) # Check async invoke needs to be used has_async_invoke: Final = "async_invoke/" in model if has_async_invoke: diff --git a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py index c98844fef55..472a9608486 100644 --- a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py @@ -1,6 +1,7 @@ import json import asyncio import os +from pathlib import Path from typing import Final from unittest.mock import Mock, patch from unittest.mock import AsyncMock, MagicMock @@ -189,6 +190,47 @@ def test_bedrock_cohere_embedding_returns_the_scripted_vector(respx_mock: respx. assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.3, 0.2, 0.1]}] +@pytest.mark.parametrize( + "env,model_id,expected_region", + [ + pytest.param({"AWS_DEFAULT_REGION": "ap-south-1"}, None, "ap-south-1", id="boto3-session-region"), + pytest.param( + {"AWS_REGION": "us-east-1"}, + "arn:aws:bedrock:eu-central-1:123456789012:application-inference-profile/abc123", + "eu-central-1", + id="model-arn-region-beats-env", + ), + ], +) +def test_bedrock_embedding_sends_to_the_same_region_as_chat( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + env: dict[str, str], + model_id: str | None, + expected_region: str, +) -> None: + for name in ("AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_PROFILE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "no-aws-config")) + for name, value in env.items(): + monkeypatch.setenv(name, value) + route: Final = respx_mock.post(url__regex=r"^https://bedrock-runtime\.[a-z0-9-]+\.amazonaws\.com/").mock( + return_value=httpx.Response(200, json=titan_embedding_response) + ) + model_id_param: Final = {} if model_id is None else {"model_id": model_id} + + litellm.embedding( + model="bedrock/amazon.titan-embed-text-v2:0", + input=["hello"], + api_key="test-bearer-token", + **model_id_param, + ) + + sent_to: Final = route.calls.last.request.url + assert sent_to.host == f"bedrock-runtime.{expected_region}.amazonaws.com", str(sent_to) + + @pytest.mark.asyncio async def test_async_bedrock_embedding_with_bearer_token(): """Test async embedding functionality with bearer token authentication"""