Skip to content

Commit 2b81b1c

Browse files
Merge branch 'strands-agents:main' into stop
2 parents c9b184b + ae9d5ad commit 2b81b1c

File tree

2 files changed

+28
-8
lines changed

2 files changed

+28
-8
lines changed

src/strands/models/bedrock.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ def __init__(
103103
boto_session: Optional[boto3.Session] = None,
104104
boto_client_config: Optional[BotocoreConfig] = None,
105105
region_name: Optional[str] = None,
106+
endpoint_url: Optional[str] = None,
106107
**model_config: Unpack[BedrockConfig],
107108
):
108109
"""Initialize provider instance.
@@ -112,6 +113,7 @@ def __init__(
112113
boto_client_config: Configuration to use when creating the Bedrock-Runtime Boto Client.
113114
region_name: AWS region to use for the Bedrock service.
114115
Defaults to the AWS_REGION environment variable if set, or "us-west-2" if not set.
116+
endpoint_url: Custom endpoint URL for VPC endpoints (PrivateLink)
115117
**model_config: Configuration options for the Bedrock model.
116118
"""
117119
if region_name and boto_session:
@@ -143,6 +145,7 @@ def __init__(
143145
self.client = session.client(
144146
service_name="bedrock-runtime",
145147
config=client_config,
148+
endpoint_url=endpoint_url,
146149
region_name=resolved_region,
147150
)
148151

tests/strands/models/test_bedrock.py

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ def test__init__with_default_region(session_cls, mock_client_method):
129129
with unittest.mock.patch.object(os, "environ", {}):
130130
BedrockModel()
131131
session_cls.return_value.client.assert_called_with(
132-
region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY
132+
region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY, endpoint_url=None
133133
)
134134

135135

@@ -139,22 +139,22 @@ def test__init__with_session_region(session_cls, mock_client_method):
139139

140140
BedrockModel()
141141

142-
mock_client_method.assert_called_with(region_name="eu-blah-1", config=ANY, service_name=ANY)
142+
mock_client_method.assert_called_with(region_name="eu-blah-1", config=ANY, service_name=ANY, endpoint_url=None)
143143

144144

145145
def test__init__with_custom_region(mock_client_method):
146146
"""Test that BedrockModel uses the provided region."""
147147
custom_region = "us-east-1"
148148
BedrockModel(region_name=custom_region)
149-
mock_client_method.assert_called_with(region_name=custom_region, config=ANY, service_name=ANY)
149+
mock_client_method.assert_called_with(region_name=custom_region, config=ANY, service_name=ANY, endpoint_url=None)
150150

151151

152152
def test__init__with_default_environment_variable_region(mock_client_method):
153153
"""Test that BedrockModel uses the AWS_REGION since we code that in."""
154154
with unittest.mock.patch.object(os, "environ", {"AWS_REGION": "eu-west-2"}):
155155
BedrockModel()
156156

157-
mock_client_method.assert_called_with(region_name="eu-west-2", config=ANY, service_name=ANY)
157+
mock_client_method.assert_called_with(region_name="eu-west-2", config=ANY, service_name=ANY, endpoint_url=None)
158158

159159

160160
def test__init__region_precedence(mock_client_method, session_cls):
@@ -164,21 +164,38 @@ def test__init__region_precedence(mock_client_method, session_cls):
164164

165165
# specifying a region always wins out
166166
BedrockModel(region_name="us-specified-1")
167-
mock_client_method.assert_called_with(region_name="us-specified-1", config=ANY, service_name=ANY)
167+
mock_client_method.assert_called_with(
168+
region_name="us-specified-1", config=ANY, service_name=ANY, endpoint_url=None
169+
)
168170

169171
# other-wise uses the session's
170172
BedrockModel()
171-
mock_client_method.assert_called_with(region_name="us-session-1", config=ANY, service_name=ANY)
173+
mock_client_method.assert_called_with(
174+
region_name="us-session-1", config=ANY, service_name=ANY, endpoint_url=None
175+
)
172176

173177
# environment variable next
174178
session_cls.return_value.region_name = None
175179
BedrockModel()
176-
mock_client_method.assert_called_with(region_name="us-environment-1", config=ANY, service_name=ANY)
180+
mock_client_method.assert_called_with(
181+
region_name="us-environment-1", config=ANY, service_name=ANY, endpoint_url=None
182+
)
177183

178184
mock_os_environ.pop("AWS_REGION")
179185
session_cls.return_value.region_name = None # No session region
180186
BedrockModel()
181-
mock_client_method.assert_called_with(region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY)
187+
mock_client_method.assert_called_with(
188+
region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY, endpoint_url=None
189+
)
190+
191+
192+
def test__init__with_endpoint_url(mock_client_method):
193+
"""Test that BedrockModel uses the provided endpoint_url for VPC endpoints."""
194+
custom_endpoint = "https://vpce-12345-abcde.bedrock-runtime.us-west-2.vpce.amazonaws.com"
195+
BedrockModel(endpoint_url=custom_endpoint)
196+
mock_client_method.assert_called_with(
197+
region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY, endpoint_url=custom_endpoint
198+
)
182199

183200

184201
def test__init__with_region_and_session_raises_value_error():

0 commit comments

Comments
 (0)