Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@ classifiers = [
"Topic :: Software Development :: Libraries :: Python Modules",
]
dependencies = [
"boto3>=1.42.54",
"botocore>=1.42.54",
"boto3>=1.42.63",
"botocore>=1.42.63",
"pydantic>=2.0.0,<2.41.3",
"urllib3>=1.26.0",
"starlette>=0.46.2",
Expand Down
9 changes: 9 additions & 0 deletions src/bedrock_agentcore/memory/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,7 @@ def create_memory(
description: Optional[str] = None,
event_expiry_days: int = 90,
memory_execution_role_arn: Optional[str] = None,
stream_delivery_resources: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""Create a memory with simplified configuration."""
if strategies is None:
Expand All @@ -154,6 +155,9 @@ def create_memory(
if memory_execution_role_arn is not None:
params["memoryExecutionRoleArn"] = memory_execution_role_arn

if stream_delivery_resources is not None:
params["streamDeliveryResources"] = stream_delivery_resources

response = self.gmcp_client.create_memory(**params)

memory = response["memory"]
Expand All @@ -174,6 +178,7 @@ def create_or_get_memory(
description: Optional[str] = None,
event_expiry_days: int = 90,
memory_execution_role_arn: Optional[str] = None,
stream_delivery_resources: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""Create a memory resource or fetch the existing memory details if it already exists.

Expand All @@ -187,6 +192,7 @@ def create_or_get_memory(
description=description,
event_expiry_days=event_expiry_days,
memory_execution_role_arn=memory_execution_role_arn,
stream_delivery_resources=stream_delivery_resources,
)
return memory
except ClientError as e:
Expand All @@ -208,6 +214,7 @@ def create_memory_and_wait(
description: Optional[str] = None,
event_expiry_days: int = 90,
memory_execution_role_arn: Optional[str] = None,
stream_delivery_resources: Optional[Dict[str, Any]] = None,
max_wait: int = 300,
poll_interval: int = 10,
) -> Dict[str, Any]:
Expand All @@ -222,6 +229,7 @@ def create_memory_and_wait(
description: Optional description
event_expiry_days: How long to retain events (default: 90 days)
memory_execution_role_arn: IAM role ARN for memory execution
stream_delivery_resources: Optional delivery configuration for streaming memory records
max_wait: Maximum seconds to wait (default: 300)
poll_interval: Seconds between status checks (default: 10)

Expand All @@ -239,6 +247,7 @@ def create_memory_and_wait(
description=description,
event_expiry_days=event_expiry_days,
memory_execution_role_arn=memory_execution_role_arn,
stream_delivery_resources=stream_delivery_resources,
)

memory_id = memory.get("memoryId", memory.get("id")) # Handle both field names
Expand Down
10 changes: 10 additions & 0 deletions src/bedrock_agentcore/memory/controlplane.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ def create_memory(
description: Optional[str] = None,
memory_execution_role_arn: Optional[str] = None,
strategies: Optional[List[Dict[str, Any]]] = None,
stream_delivery_resources: Optional[Dict[str, Any]] = None,
wait_for_active: bool = False,
max_wait: int = 300,
poll_interval: int = 10,
Expand All @@ -63,6 +64,7 @@ def create_memory(
description: Optional description
memory_execution_role_arn: IAM role ARN for memory execution
strategies: Optional list of strategy configurations
stream_delivery_resources: Optional delivery configuration for streaming memory records
wait_for_active: Whether to wait for memory to become ACTIVE
max_wait: Maximum seconds to wait if wait_for_active is True
poll_interval: Seconds between status checks if wait_for_active is True
Expand All @@ -85,6 +87,9 @@ def create_memory(
if strategies:
params["memoryStrategies"] = strategies

if stream_delivery_resources is not None:
params["streamDeliveryResources"] = stream_delivery_resources

try:
response = self.client.create_memory(**params)
memory = response["memory"]
Expand Down Expand Up @@ -174,6 +179,7 @@ def update_memory(
add_strategies: Optional[List[Dict[str, Any]]] = None,
modify_strategies: Optional[List[Dict[str, Any]]] = None,
delete_strategy_ids: Optional[List[str]] = None,
stream_delivery_resources: Optional[Dict[str, Any]] = None,
wait_for_active: bool = False,
max_wait: int = 300,
poll_interval: int = 10,
Expand All @@ -188,6 +194,7 @@ def update_memory(
add_strategies: Optional list of strategies to add
modify_strategies: Optional list of strategies to modify
delete_strategy_ids: Optional list of strategy IDs to delete
stream_delivery_resources: Optional delivery configuration for streaming memory records
wait_for_active: Whether to wait for memory to become ACTIVE
max_wait: Maximum seconds to wait if wait_for_active is True
poll_interval: Seconds between status checks if wait_for_active is True
Expand All @@ -210,6 +217,9 @@ def update_memory(
if memory_execution_role_arn is not None:
params["memoryExecutionRoleArn"] = memory_execution_role_arn

if stream_delivery_resources is not None:
params["streamDeliveryResources"] = stream_delivery_resources

# Add strategy operations if provided
memory_strategies = {}

Expand Down
70 changes: 70 additions & 0 deletions tests/bedrock_agentcore/memory/test_controlplane.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,76 @@ def test_update_memory():
assert kwargs["clientToken"] == "12345678-1234-5678-1234-567812345678"


def test_create_memory_with_stream_delivery():
"""Test create_memory with stream_delivery_resources."""
with patch("boto3.client"):
client = MemoryControlPlaneClient()

mock_client = MagicMock()
client.client = mock_client

delivery_config = {
"resources": [
{
"kinesis": {
"dataStreamArn": "arn:aws:kinesis:us-west-2:123456789012:stream/test",
"contentConfigurations": [{"type": "MEMORY_RECORDS", "level": "FULL_CONTENT"}],
}
}
]
}

mock_client.create_memory.return_value = {
"memory": {
"id": "mem-123",
"name": "Test",
"status": "CREATING",
"streamDeliveryResources": delivery_config,
}
}

result = client.create_memory(name="Test", stream_delivery_resources=delivery_config)

assert result["streamDeliveryResources"] == delivery_config
args, kwargs = mock_client.create_memory.call_args
assert kwargs["streamDeliveryResources"] == delivery_config


def test_update_memory_with_stream_delivery():
"""Test update_memory with stream_delivery_resources."""
with patch("boto3.client"):
client = MemoryControlPlaneClient()

mock_client = MagicMock()
client.client = mock_client

delivery_config = {
"resources": [
{
"kinesis": {
"dataStreamArn": "arn:aws:kinesis:us-west-2:123456789012:stream/test",
"contentConfigurations": [{"type": "MEMORY_RECORDS", "level": "METADATA_ONLY"}],
}
}
]
}

mock_client.update_memory.return_value = {
"memory": {
"id": "mem-123",
"name": "Test",
"status": "UPDATING",
"streamDeliveryResources": delivery_config,
}
}

result = client.update_memory(memory_id="mem-123", stream_delivery_resources=delivery_config)

assert result["streamDeliveryResources"] == delivery_config
args, kwargs = mock_client.update_memory.call_args
assert kwargs["streamDeliveryResources"] == delivery_config


def test_delete_memory():
"""Test delete_memory functionality."""
with patch("boto3.client"):
Expand Down
16 changes: 8 additions & 8 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading