Merge f5c942f9e8 into 295d79fc7d

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:

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"""