Skip to content

Commit c3f2324

Browse files
authored
Merge branch 'strands-agents:main' into main
2 parents 4db15c8 + ae9d5ad commit c3f2324

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:
@@ -151,6 +153,7 @@ def __init__(
151153
self.client = session.client(
152154
service_name="bedrock-runtime",
153155
config=client_config,
156+
endpoint_url=endpoint_url,
154157
region_name=resolved_region,
155158
)
156159

tests/strands/models/test_bedrock.py

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ def test__init__with_default_region(session_cls, mock_client_method, mock_bedroc
140140
with unittest.mock.patch.object(os, "environ", {}):
141141
BedrockModel()
142142
session_cls.return_value.client.assert_called_with(
143-
region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY
143+
region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY, endpoint_url=None
144144
)
145145

146146

@@ -150,22 +150,22 @@ def test__init__with_session_region(session_cls, mock_client_method, mock_bedroc
150150

151151
BedrockModel()
152152

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

155155

156156
def test__init__with_custom_region(mock_client_method, mock_bedrock_inference_profiles):
157157
"""Test that BedrockModel uses the provided region."""
158158
custom_region = "us-east-1"
159159
BedrockModel(region_name=custom_region)
160-
mock_client_method.assert_called_with(region_name=custom_region, config=ANY, service_name=ANY)
160+
mock_client_method.assert_called_with(region_name=custom_region, config=ANY, service_name=ANY, endpoint_url=None)
161161

162162

163163
def test__init__with_default_environment_variable_region(mock_client_method, mock_bedrock_inference_profiles):
164164
"""Test that BedrockModel uses the AWS_REGION since we code that in."""
165165
with unittest.mock.patch.object(os, "environ", {"AWS_REGION": "eu-west-2"}):
166166
BedrockModel()
167167

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

170170

171171
def test__init__region_precedence(mock_client_method, session_cls, mock_bedrock_inference_profiles):
@@ -175,21 +175,38 @@ def test__init__region_precedence(mock_client_method, session_cls, mock_bedrock_
175175

176176
# specifying a region always wins out
177177
BedrockModel(region_name="us-specified-1")
178-
mock_client_method.assert_called_with(region_name="us-specified-1", config=ANY, service_name=ANY)
178+
mock_client_method.assert_called_with(
179+
region_name="us-specified-1", config=ANY, service_name=ANY, endpoint_url=None
180+
)
179181

180182
# other-wise uses the session's
181183
BedrockModel()
182-
mock_client_method.assert_called_with(region_name="us-session-1", config=ANY, service_name=ANY)
184+
mock_client_method.assert_called_with(
185+
region_name="us-session-1", config=ANY, service_name=ANY, endpoint_url=None
186+
)
183187

184188
# environment variable next
185189
session_cls.return_value.region_name = None
186190
BedrockModel()
187-
mock_client_method.assert_called_with(region_name="us-environment-1", config=ANY, service_name=ANY)
191+
mock_client_method.assert_called_with(
192+
region_name="us-environment-1", config=ANY, service_name=ANY, endpoint_url=None
193+
)
188194

189195
mock_os_environ.pop("AWS_REGION")
190196
session_cls.return_value.region_name = None # No session region
191197
BedrockModel()
192-
mock_client_method.assert_called_with(region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY)
198+
mock_client_method.assert_called_with(
199+
region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY, endpoint_url=None
200+
)
201+
202+
203+
def test__init__with_endpoint_url(mock_client_method):
204+
"""Test that BedrockModel uses the provided endpoint_url for VPC endpoints."""
205+
custom_endpoint = "https://vpce-12345-abcde.bedrock-runtime.us-west-2.vpce.amazonaws.com"
206+
BedrockModel(endpoint_url=custom_endpoint)
207+
mock_client_method.assert_called_with(
208+
region_name=DEFAULT_BEDROCK_REGION, config=ANY, service_name=ANY, endpoint_url=custom_endpoint
209+
)
193210

194211

195212
def test__init__with_region_and_session_raises_value_error():

0 commit comments

Comments
 (0)