diff --git a/src/bedrock_agentcore/memory/client.py b/src/bedrock_agentcore/memory/client.py index 421dfcff..8505d076 100644 --- a/src/bedrock_agentcore/memory/client.py +++ b/src/bedrock_agentcore/memory/client.py @@ -442,6 +442,7 @@ def create_event( event_timestamp: Optional[datetime] = None, branch: Optional[Dict[str, str]] = None, metadata: Optional[Dict[str, MetadataValue]] = None, + extraction_mode: Optional[str] = None, ) -> Dict[str, Any]: """Save an event of an agent interaction or conversation with a user. @@ -461,6 +462,8 @@ def create_event( metadata: Optional custom key-value metadata to attach to the event. Maximum 15 key-value pairs. Keys must be 1-128 characters. Example: {"location": {"stringValue": "NYC"}} + extraction_mode: Controls long-term memory extraction. Set to "SKIP" to store + the event in short-term memory without triggering extraction. Returns: Created event @@ -543,6 +546,9 @@ def create_event( if metadata: params["metadata"] = metadata + if extraction_mode: + params["extractionMode"] = extraction_mode + response = self.gmdp_client.create_event(**params) event = response["event"] diff --git a/tests/bedrock_agentcore/memory/test_client.py b/tests/bedrock_agentcore/memory/test_client.py index f7df4c4a..bcb6258b 100644 --- a/tests/bedrock_agentcore/memory/test_client.py +++ b/tests/bedrock_agentcore/memory/test_client.py @@ -2148,6 +2148,69 @@ def test_create_event_with_multiple_metadata_keys(): assert len(kwargs["metadata"]) == 3 +def test_create_event_with_extraction_mode(): + """Test create_event with extraction_mode parameter.""" + with patch("boto3.Session"): + client = MemoryClient() + + # Mock the client + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # Mock create_event response + mock_gmdp.create_event.return_value = { + "event": { + "eventId": "event-skip-123", + "memoryId": "mem-123", + } + } + + # Test create_event with extraction_mode=SKIP + result = client.create_event( + memory_id="mem-123", + actor_id="user-123", + session_id="session-456", + messages=[("Sensitive data here", "USER")], + extraction_mode="SKIP", + ) + + assert result["eventId"] == "event-skip-123" + + # Verify extractionMode was passed correctly + args, kwargs = mock_gmdp.create_event.call_args + assert kwargs["extractionMode"] == "SKIP" + + +def test_create_event_without_extraction_mode(): + """Test create_event without extraction_mode does not include the field.""" + with patch("boto3.Session"): + client = MemoryClient() + + # Mock the client + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # Mock create_event response + mock_gmdp.create_event.return_value = { + "event": { + "eventId": "event-normal-123", + "memoryId": "mem-123", + } + } + + # Test create_event without extraction_mode + client.create_event( + memory_id="mem-123", + actor_id="user-123", + session_id="session-456", + messages=[("Normal data here", "USER")], + ) + + # Verify extractionMode was NOT passed + args, kwargs = mock_gmdp.create_event.call_args + assert "extractionMode" not in kwargs + + def test_create_memory_and_wait_client_error(): """Test create_memory_and_wait with ClientError during status check.""" with patch("boto3.Session"):