diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 2d822fcfe6a..302e5e30981 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -274,7 +274,8 @@ def __init__( """Create a Bedrock chat client and load AWS credentials. Args: - region: Region to send Bedrock requests to; falls back to BEDROCK_REGION. + region: Region to send Bedrock requests to; falls back to BEDROCK_REGION, then to the region + of ``boto3_session``, then to us-east-1. model: Default model identifier; falls back to BEDROCK_CHAT_MODEL. access_key: Optional AWS access key for manual credential injection. secret_key: Optional AWS secret key paired with ``access_key``. @@ -318,7 +319,8 @@ class MyOptions(BedrockChatOptions, total=False): env_file_path=env_file_path, env_file_encoding=env_file_encoding, ) - region = settings.get("region") or DEFAULT_REGION + session_region = boto3_session.region_name if boto3_session and not client else None + region = settings.get("region") or session_region or DEFAULT_REGION chat_model = settings.get("chat_model") if client: diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index 3d51255eca4..209cbaa9315 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -526,18 +526,18 @@ def test_parse_usage_returns_none_when_no_recognized_keys() -> None: assert client._parse_usage(None) is None -def test_init_uses_boto3_session_when_runtime_client_not_supplied() -> None: - """BedrockChatClient should build a runtime client from a provided boto3 session.""" +class _FakeSession: + def __init__(self, region_name: str | None = None) -> None: + self.calls: list[dict[str, Any]] = [] + self.region_name = region_name - class _FakeSession: - def __init__(self) -> None: - self.calls: list[dict[str, Any]] = [] - self.region_name: str | None = None + def client(self, service_name: str, *, region_name: str, config: Any) -> _StubBedrockRuntime: + self.calls.append({"service_name": service_name, "region_name": region_name, "config": config}) + return _StubBedrockRuntime() - def client(self, service_name: str, *, region_name: str, config: Any) -> _StubBedrockRuntime: - self.calls.append({"service_name": service_name, "region_name": region_name, "config": config}) - return _StubBedrockRuntime() +def test_init_uses_boto3_session_when_runtime_client_not_supplied() -> None: + """BedrockChatClient should build a runtime client from a provided boto3 session.""" session = _FakeSession() client = BedrockChatClient( @@ -556,6 +556,53 @@ def client(self, service_name: str, *, region_name: str, config: Any) -> _StubBe ] +@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"], +) +def test_init_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 = BedrockChatClient(model="amazon.titan-text", region=region, boto3_session=cast(Boto3Session, session)) + + assert session.calls[0]["region_name"] == expected_region + assert client.region == expected_region + assert client.service_url() == f"https://bedrock-runtime.{expected_region}.amazonaws.com" + + +def test_init_ignores_boto3_session_region_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 = BedrockChatClient( + model="amazon.titan-text", + client=cast(BaseClient, _StubBedrockRuntime()), + boto3_session=cast(Boto3Session, session), + ) + + assert session.calls == [] + assert client.region == "us-east-1" + assert client.service_url() == "https://bedrock-runtime.us-east-1.amazonaws.com" + + @pytest.mark.parametrize("secret_type", [str, SecretString], ids=["str", "secret"]) @pytest.mark.parametrize( ("client_type", "module"),