From 45420b3c680db646d2f60c364f89ff1fbcf76077 Mon Sep 17 00:00:00 2001 From: Shanto Date: Tue, 13 Oct 2026 22:28:39 +0600 Subject: [PATCH] Resolve the BedrockEmbeddingClient region like the chat client --- .../_embedding_client.py | 14 ++-- .../bedrock/test_bedrock_embedding_client.py | 73 ++++++++++++++++++- 2 files changed, 80 insertions(+), 7 deletions(-) diff --git a/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py b/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py index 9d55d4a22f8..562b4229f5a 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py @@ -85,8 +85,9 @@ class RawBedrockEmbeddingClient( model: The Bedrock embedding model ID (e.g. "amazon.titan-embed-text-v2:0"). Can also be set via environment variable BEDROCK_EMBEDDING_MODEL. region: AWS region. Will try to load from BEDROCK_REGION env var, - if not set, the regular Boto3 configuration/loading applies - (which may include other env vars, config files, or instance metadata). + if not set, the region of ``boto3_session`` is used, or without a session + the regular Boto3 configuration/loading applies (which may include other + env vars, config files, or instance metadata). Falls back to us-east-1. access_key: AWS access key for manual credential injection. secret_key: AWS secret key paired with access_key. session_token: AWS session token for temporary credentials. @@ -138,10 +139,10 @@ def __init__( if session_token := settings.get("session_token"): session_kwargs["aws_session_token"] = session_token.get_secret_value() boto3_session = Boto3Session(**session_kwargs) - region_name = boto3_session.region_name + resolved_region = settings.get("region") or boto3_session.region_name or DEFAULT_REGION self._bedrock_client = boto3_session.client( "bedrock-runtime", - region_name=region_name or resolved_region, + region_name=resolved_region, config=BotoConfig(user_agent_extra=get_user_agent()), ) @@ -239,8 +240,9 @@ class BedrockEmbeddingClient( Keyword Args: model: The Bedrock embedding model ID (e.g. "amazon.titan-embed-text-v2:0"). Can also be set via environment variable BEDROCK_EMBEDDING_MODEL. - region: AWS region. Defaults to "us-east-1". - Can also be set via environment variable BEDROCK_REGION. + region: AWS region. Can also be set via environment variable BEDROCK_REGION. + If not set, the region of ``boto3_session`` is used, or without a session + the regular Boto3 configuration. Defaults to "us-east-1". access_key: AWS access key for manual credential injection. secret_key: AWS secret key paired with access_key. session_token: AWS session token for temporary credentials. diff --git a/python/packages/bedrock/tests/bedrock/test_bedrock_embedding_client.py b/python/packages/bedrock/tests/bedrock/test_bedrock_embedding_client.py index b16ec44790a..e7ec102509a 100644 --- a/python/packages/bedrock/tests/bedrock/test_bedrock_embedding_client.py +++ b/python/packages/bedrock/tests/bedrock/test_bedrock_embedding_client.py @@ -4,11 +4,12 @@ import json import os -from typing import Any +from typing import Any, cast from unittest.mock import MagicMock import pytest from agent_framework import Embedding, GeneratedEmbeddings +from boto3.session import Session as Boto3Session from agent_framework_bedrock import BedrockEmbeddingClient, BedrockEmbeddingOptions @@ -143,6 +144,76 @@ async def test_bedrock_embedding_default_region() -> None: assert client.region == "us-east-1" +class _FakeSession: + def __init__(self, region_name: str | None = None) -> None: + self.calls: list[dict[str, Any]] = [] + self.region_name = region_name + + def client(self, service_name: str, *, region_name: str, config: Any) -> _StubBedrockEmbeddingRuntime: + self.calls.append({"service_name": service_name, "region_name": region_name}) + return _StubBedrockEmbeddingRuntime() + + +@pytest.mark.parametrize( + ("region", "env_region", "session_region", "expected_region"), + [ + (None, None, "eu-west-1", "eu-west-1"), + ("us-west-2", None, "eu-west-1", "us-west-2"), + (None, "ap-south-1", "eu-west-1", "ap-south-1"), + (None, None, None, "us-east-1"), + ], + ids=["session-region", "region-argument-wins", "env-region-wins", "default-region"], +) +async def test_bedrock_embedding_resolves_region_with_boto3_session( + monkeypatch: pytest.MonkeyPatch, + region: str | None, + env_region: str | None, + session_region: str | None, + expected_region: str, +) -> None: + """A configured region wins, then the provided session's region, then the default region.""" + if env_region: + monkeypatch.setenv("BEDROCK_REGION", env_region) + else: + monkeypatch.delenv("BEDROCK_REGION", raising=False) + session = _FakeSession(region_name=session_region) + + client = BedrockEmbeddingClient( + model="amazon.titan-embed-text-v2:0", region=region, boto3_session=cast(Boto3Session, session) + ) + + assert session.calls == [{"service_name": "bedrock-runtime", "region_name": expected_region}] + assert client.region == expected_region + + +async def test_bedrock_embedding_reports_boto3_default_region(monkeypatch: pytest.MonkeyPatch) -> None: + """Without a region or session, the region boto3 resolves is used and reported.""" + monkeypatch.delenv("BEDROCK_REGION", raising=False) + monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-central-1") + + client = BedrockEmbeddingClient(model="amazon.titan-embed-text-v2:0") + + assert client.service_url() == "https://bedrock-runtime.eu-central-1.amazonaws.com" + assert client.region == "eu-central-1" + + +async def test_bedrock_embedding_ignores_boto3_session_when_runtime_client_supplied( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A preconfigured runtime client wins, so the unused session should not change the reported region.""" + monkeypatch.delenv("BEDROCK_REGION", raising=False) + session = _FakeSession(region_name="eu-west-1") + + client = BedrockEmbeddingClient( + model="amazon.titan-embed-text-v2:0", + client=_StubBedrockEmbeddingRuntime(), # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType] + boto3_session=cast(Boto3Session, session), + ) + + assert session.calls == [] + assert client.region == "us-east-1" + + # region: Integration Tests skip_if_bedrock_embedding_integration_tests_disabled = pytest.mark.skipif(