From ca0bc60fba9e0c102ec0c18a40d2dbb96b868225 Mon Sep 17 00:00:00 2001 From: hkobew Date: Mon, 9 Mar 2026 15:14:03 -0400 Subject: [PATCH] feat: add streamDeliveryResources support to memory SDK --- pyproject.toml | 4 +- src/bedrock_agentcore/memory/client.py | 9 +++ src/bedrock_agentcore/memory/controlplane.py | 10 +++ .../memory/test_controlplane.py | 70 +++++++++++++++++++ uv.lock | 16 ++--- 5 files changed, 99 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 03487097..849ba1b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/src/bedrock_agentcore/memory/client.py b/src/bedrock_agentcore/memory/client.py index 8c8649ac..ec2e9b2c 100644 --- a/src/bedrock_agentcore/memory/client.py +++ b/src/bedrock_agentcore/memory/client.py @@ -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: @@ -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"] @@ -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. @@ -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: @@ -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]: @@ -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) @@ -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 diff --git a/src/bedrock_agentcore/memory/controlplane.py b/src/bedrock_agentcore/memory/controlplane.py index 25e75d09..64fb77ee 100644 --- a/src/bedrock_agentcore/memory/controlplane.py +++ b/src/bedrock_agentcore/memory/controlplane.py @@ -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, @@ -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 @@ -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"] @@ -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, @@ -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 @@ -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 = {} diff --git a/tests/bedrock_agentcore/memory/test_controlplane.py b/tests/bedrock_agentcore/memory/test_controlplane.py index 38b691c0..8510840b 100644 --- a/tests/bedrock_agentcore/memory/test_controlplane.py +++ b/tests/bedrock_agentcore/memory/test_controlplane.py @@ -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"): diff --git a/uv.lock b/uv.lock index c2894e2b..0921d93f 100644 --- a/uv.lock +++ b/uv.lock @@ -261,8 +261,8 @@ dev = [ [package.metadata] requires-dist = [ - { name = "boto3", specifier = ">=1.42.54" }, - { name = "botocore", specifier = ">=1.42.54" }, + { name = "boto3", specifier = ">=1.42.63" }, + { name = "botocore", specifier = ">=1.42.63" }, { name = "pydantic", specifier = ">=2.0.0,<2.41.3" }, { name = "starlette", specifier = ">=0.46.2" }, { name = "strands-agents", marker = "extra == 'strands-agents'", specifier = ">=1.1.0" }, @@ -292,30 +292,30 @@ dev = [ [[package]] name = "boto3" -version = "1.42.55" +version = "1.42.63" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore" }, { name = "jmespath" }, { name = "s3transfer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3d/74/9e97b4ca692ed3ee2c5cb616790c3b00290b73babc63b85c7ed392ed74b1/boto3-1.42.55.tar.gz", hash = "sha256:e7b8fcc123da442449da8a2be65b3e60a3d8cfb2b26a52f7b3c6f9f8e84cbdf0", size = 112771, upload-time = "2026-02-23T20:29:29.933Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7f/2a/33d5d4b16fd97dfd629421ebed2456392eae1553cc401d9f86010c18065e/boto3-1.42.63.tar.gz", hash = "sha256:cd008cfd0d7ea30f1c5e22daf0998c55b7c6c68cb68eea05110e33fe641173d5", size = 112778, upload-time = "2026-03-06T22:47:55.96Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/19/04/ca0b37dbb980fc59e9414f7bcb5d6209b1d3a03da433784e21fdd7282269/boto3-1.42.55-py3-none-any.whl", hash = "sha256:cb4bc94c0ba522242e291d16b4f631e139f525fbc9772229f3e84f5d834fd88e", size = 140556, upload-time = "2026-02-23T20:29:27.402Z" }, + { url = "https://files.pythonhosted.org/packages/f5/19/f1d8d2b24871d3d0ccb2cbd0b0cb64a3396d439384bd9643d2c25c641b84/boto3-1.42.63-py3-none-any.whl", hash = "sha256:d502a89a0acc701692ae020d15981f2a82e9eb3485acc651cfd0cf1a3afe79ee", size = 140554, upload-time = "2026-03-06T22:47:53.463Z" }, ] [[package]] name = "botocore" -version = "1.42.55" +version = "1.42.63" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jmespath" }, { name = "python-dateutil" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/5a/b9/958d53c0e0b843c25d93d7593364b3e92913dfac381c82fa2b8a470fdf78/botocore-1.42.55.tar.gz", hash = "sha256:af22a7d7881883bcb475a627d0750ec6f8ee3d7b2f673e9ff342ebaa498447ee", size = 14927543, upload-time = "2026-02-23T20:29:17.923Z" } +sdist = { url = "https://files.pythonhosted.org/packages/af/eb/a1c042f6638ada552399a9977335a6de2668a85bf80bece193c953531236/botocore-1.42.63.tar.gz", hash = "sha256:1fdfc33cff58d21e8622cf620ba2bba3cff324557932aaf935b5374e4610f059", size = 14965362, upload-time = "2026-03-06T22:47:44.158Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e5/64/fe72b409660b8da44a8763f9165d36650e41e4e591dd7d3ad708397496c7/botocore-1.42.55-py3-none-any.whl", hash = "sha256:c092eb99d17b653af3ec9242061a7cde1c7b1940ed4abddfada68a9e1a3492d6", size = 14598862, upload-time = "2026-02-23T20:29:11.589Z" }, + { url = "https://files.pythonhosted.org/packages/9a/60/17a2d3b94658bb999c6aee7bba6c76b271905debf0c8c8e6ac63ca8491bc/botocore-1.42.63-py3-none-any.whl", hash = "sha256:83f39d04f2b316bdfc59a3cac2d12238bde7126ac99d9a57d910dbd86d58c528", size = 14639889, upload-time = "2026-03-06T22:47:39.347Z" }, ] [[package]]