Skip to content

Commit ca0bc60

Browse files
committed
feat: add streamDeliveryResources support to memory SDK
1 parent 17e1357 commit ca0bc60

5 files changed

Lines changed: 99 additions & 10 deletions

File tree

pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ classifiers = [
2626
"Topic :: Software Development :: Libraries :: Python Modules",
2727
]
2828
dependencies = [
29-
"boto3>=1.42.54",
30-
"botocore>=1.42.54",
29+
"boto3>=1.42.63",
30+
"botocore>=1.42.63",
3131
"pydantic>=2.0.0,<2.41.3",
3232
"urllib3>=1.26.0",
3333
"starlette>=0.46.2",

src/bedrock_agentcore/memory/client.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ def create_memory(
133133
description: Optional[str] = None,
134134
event_expiry_days: int = 90,
135135
memory_execution_role_arn: Optional[str] = None,
136+
stream_delivery_resources: Optional[Dict[str, Any]] = None,
136137
) -> Dict[str, Any]:
137138
"""Create a memory with simplified configuration."""
138139
if strategies is None:
@@ -154,6 +155,9 @@ def create_memory(
154155
if memory_execution_role_arn is not None:
155156
params["memoryExecutionRoleArn"] = memory_execution_role_arn
156157

158+
if stream_delivery_resources is not None:
159+
params["streamDeliveryResources"] = stream_delivery_resources
160+
157161
response = self.gmcp_client.create_memory(**params)
158162

159163
memory = response["memory"]
@@ -174,6 +178,7 @@ def create_or_get_memory(
174178
description: Optional[str] = None,
175179
event_expiry_days: int = 90,
176180
memory_execution_role_arn: Optional[str] = None,
181+
stream_delivery_resources: Optional[Dict[str, Any]] = None,
177182
) -> Dict[str, Any]:
178183
"""Create a memory resource or fetch the existing memory details if it already exists.
179184
@@ -187,6 +192,7 @@ def create_or_get_memory(
187192
description=description,
188193
event_expiry_days=event_expiry_days,
189194
memory_execution_role_arn=memory_execution_role_arn,
195+
stream_delivery_resources=stream_delivery_resources,
190196
)
191197
return memory
192198
except ClientError as e:
@@ -208,6 +214,7 @@ def create_memory_and_wait(
208214
description: Optional[str] = None,
209215
event_expiry_days: int = 90,
210216
memory_execution_role_arn: Optional[str] = None,
217+
stream_delivery_resources: Optional[Dict[str, Any]] = None,
211218
max_wait: int = 300,
212219
poll_interval: int = 10,
213220
) -> Dict[str, Any]:
@@ -222,6 +229,7 @@ def create_memory_and_wait(
222229
description: Optional description
223230
event_expiry_days: How long to retain events (default: 90 days)
224231
memory_execution_role_arn: IAM role ARN for memory execution
232+
stream_delivery_resources: Optional delivery configuration for streaming memory records
225233
max_wait: Maximum seconds to wait (default: 300)
226234
poll_interval: Seconds between status checks (default: 10)
227235
@@ -239,6 +247,7 @@ def create_memory_and_wait(
239247
description=description,
240248
event_expiry_days=event_expiry_days,
241249
memory_execution_role_arn=memory_execution_role_arn,
250+
stream_delivery_resources=stream_delivery_resources,
242251
)
243252

244253
memory_id = memory.get("memoryId", memory.get("id")) # Handle both field names

src/bedrock_agentcore/memory/controlplane.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ def create_memory(
5151
description: Optional[str] = None,
5252
memory_execution_role_arn: Optional[str] = None,
5353
strategies: Optional[List[Dict[str, Any]]] = None,
54+
stream_delivery_resources: Optional[Dict[str, Any]] = None,
5455
wait_for_active: bool = False,
5556
max_wait: int = 300,
5657
poll_interval: int = 10,
@@ -63,6 +64,7 @@ def create_memory(
6364
description: Optional description
6465
memory_execution_role_arn: IAM role ARN for memory execution
6566
strategies: Optional list of strategy configurations
67+
stream_delivery_resources: Optional delivery configuration for streaming memory records
6668
wait_for_active: Whether to wait for memory to become ACTIVE
6769
max_wait: Maximum seconds to wait if wait_for_active is True
6870
poll_interval: Seconds between status checks if wait_for_active is True
@@ -85,6 +87,9 @@ def create_memory(
8587
if strategies:
8688
params["memoryStrategies"] = strategies
8789

90+
if stream_delivery_resources is not None:
91+
params["streamDeliveryResources"] = stream_delivery_resources
92+
8893
try:
8994
response = self.client.create_memory(**params)
9095
memory = response["memory"]
@@ -174,6 +179,7 @@ def update_memory(
174179
add_strategies: Optional[List[Dict[str, Any]]] = None,
175180
modify_strategies: Optional[List[Dict[str, Any]]] = None,
176181
delete_strategy_ids: Optional[List[str]] = None,
182+
stream_delivery_resources: Optional[Dict[str, Any]] = None,
177183
wait_for_active: bool = False,
178184
max_wait: int = 300,
179185
poll_interval: int = 10,
@@ -188,6 +194,7 @@ def update_memory(
188194
add_strategies: Optional list of strategies to add
189195
modify_strategies: Optional list of strategies to modify
190196
delete_strategy_ids: Optional list of strategy IDs to delete
197+
stream_delivery_resources: Optional delivery configuration for streaming memory records
191198
wait_for_active: Whether to wait for memory to become ACTIVE
192199
max_wait: Maximum seconds to wait if wait_for_active is True
193200
poll_interval: Seconds between status checks if wait_for_active is True
@@ -210,6 +217,9 @@ def update_memory(
210217
if memory_execution_role_arn is not None:
211218
params["memoryExecutionRoleArn"] = memory_execution_role_arn
212219

220+
if stream_delivery_resources is not None:
221+
params["streamDeliveryResources"] = stream_delivery_resources
222+
213223
# Add strategy operations if provided
214224
memory_strategies = {}
215225

tests/bedrock_agentcore/memory/test_controlplane.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,76 @@ def test_update_memory():
128128
assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678"
129129

130130

131+
def test_create_memory_with_stream_delivery():
132+
"""Test create_memory with stream_delivery_resources."""
133+
with patch("boto3.client"):
134+
client = MemoryControlPlaneClient()
135+
136+
mock_client = MagicMock()
137+
client.client = mock_client
138+
139+
delivery_config = {
140+
"resources": [
141+
{
142+
"kinesis": {
143+
"dataStreamArn": "arn:aws:kinesis:us-west-2:123456789012:stream/test",
144+
"contentConfigurations": [{"type": "MEMORY_RECORDS", "level": "FULL_CONTENT"}],
145+
}
146+
}
147+
]
148+
}
149+
150+
mock_client.create_memory.return_value = {
151+
"memory": {
152+
"id": "mem-123",
153+
"name": "Test",
154+
"status": "CREATING",
155+
"streamDeliveryResources": delivery_config,
156+
}
157+
}
158+
159+
result = client.create_memory(name="Test", stream_delivery_resources=delivery_config)
160+
161+
assert result["streamDeliveryResources"] == delivery_config
162+
args, kwargs = mock_client.create_memory.call_args
163+
assert kwargs["streamDeliveryResources"] == delivery_config
164+
165+
166+
def test_update_memory_with_stream_delivery():
167+
"""Test update_memory with stream_delivery_resources."""
168+
with patch("boto3.client"):
169+
client = MemoryControlPlaneClient()
170+
171+
mock_client = MagicMock()
172+
client.client = mock_client
173+
174+
delivery_config = {
175+
"resources": [
176+
{
177+
"kinesis": {
178+
"dataStreamArn": "arn:aws:kinesis:us-west-2:123456789012:stream/test",
179+
"contentConfigurations": [{"type": "MEMORY_RECORDS", "level": "METADATA_ONLY"}],
180+
}
181+
}
182+
]
183+
}
184+
185+
mock_client.update_memory.return_value = {
186+
"memory": {
187+
"id": "mem-123",
188+
"name": "Test",
189+
"status": "UPDATING",
190+
"streamDeliveryResources": delivery_config,
191+
}
192+
}
193+
194+
result = client.update_memory(memory_id="mem-123", stream_delivery_resources=delivery_config)
195+
196+
assert result["streamDeliveryResources"] == delivery_config
197+
args, kwargs = mock_client.update_memory.call_args
198+
assert kwargs["streamDeliveryResources"] == delivery_config
199+
200+
131201
def test_delete_memory():
132202
"""Test delete_memory functionality."""
133203
with patch("boto3.client"):

uv.lock

Lines changed: 8 additions & 8 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)