From 368dec358bb233c7d649e6541c28802ca50916ba Mon Sep 17 00:00:00 2001 From: "T.J Ariyawansa" Date: Thu, 29 Jan 2026 21:26:14 -0500 Subject: [PATCH] feat(memory): add metadata support to MemoryClient events Add metadata parameter to create_event, create_blob_event, fork_conversation, and event_metadata filter to list_events in MemoryClient. - create_event: Added optional metadata parameter for attaching key-value metadata - create_blob_event: Added optional metadata parameter - list_events: Added event_metadata filter for querying events by metadata - fork_conversation: Added optional metadata parameter (passes through to create_event) Metadata supports up to 15 key-value pairs with keys 1-128 characters. Filtering supports EQUALS_TO, EXISTS, and NOT_EXISTS operators. --- src/bedrock_agentcore/memory/client.py | 53 +++- tests/bedrock_agentcore/memory/test_client.py | 254 ++++++++++++++++++ 2 files changed, 297 insertions(+), 10 deletions(-) diff --git a/src/bedrock_agentcore/memory/client.py b/src/bedrock_agentcore/memory/client.py index afd8415..eb0d6d4 100644 --- a/src/bedrock_agentcore/memory/client.py +++ b/src/bedrock_agentcore/memory/client.py @@ -35,6 +35,7 @@ Role, StrategyType, ) +from .models.filters import EventMetadataFilter, MetadataValue logger = logging.getLogger(__name__) @@ -344,6 +345,7 @@ def create_event( messages: List[Tuple[str, str]], event_timestamp: Optional[datetime] = None, branch: Optional[Dict[str, str]] = None, + metadata: Optional[Dict[str, MetadataValue]] = None, ) -> Dict[str, Any]: """Save an event of an agent interaction or conversation with a user. @@ -360,6 +362,9 @@ def create_event( branch: Optional branch info. For new branches: {"rootEventId": "...", "name": "..."} For continuing existing branch: {"name": "..."} or {"name": "...", "rootEventId": "..."} A branch is used when you want to have a different history of events. + 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"}} Returns: Created event @@ -439,6 +444,9 @@ def create_event( if branch: params["branch"] = branch + if metadata: + params["metadata"] = metadata + response = self.gmdp_client.create_event(**params) event = response["event"] @@ -458,6 +466,7 @@ def create_blob_event( blob_data: Any, event_timestamp: Optional[datetime] = None, branch: Optional[Dict[str, str]] = None, + metadata: Optional[Dict[str, MetadataValue]] = None, ) -> Dict[str, Any]: """Save a blob event to AgentCore Memory. @@ -468,17 +477,20 @@ def create_blob_event( blob_data: Binary or structured data to store event_timestamp: Optional timestamp for the event branch: Optional branch info + 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"}} Returns: Created event Example: - # Store binary data event = client.create_blob_event( memory_id="mem-xyz", actor_id="user-123", session_id="session-456", - blob_data={"file_content": "base64_encoded_data", "metadata": {"type": "image"}} + blob_data={"file_content": "base64_encoded_data"}, + metadata={"type": {"stringValue": "image"}} ) """ try: @@ -498,6 +510,9 @@ def create_blob_event( if branch: params["branch"] = branch + if metadata: + params["metadata"] = metadata + response = self.gmdp_client.create_event(**params) event = response["event"] @@ -771,6 +786,7 @@ def list_events( session_id: str, branch_name: Optional[str] = None, include_parent_branches: bool = False, + event_metadata: Optional[List[EventMetadataFilter]] = None, max_results: int = 100, include_payload: bool = True, ) -> List[Dict[str, Any]]: @@ -785,6 +801,9 @@ def list_events( session_id: Session identifier branch_name: Optional branch name to filter events (None for all branches) include_parent_branches: Whether to include parent branch events (only applies with branch_name) + event_metadata: Optional list of event metadata filters to apply. + Example: [{"left": {"metadataKey": "location"}, "operator": "EQUALS_TO", + "right": {"metadataValue": {"stringValue": "NYC"}}}] max_results: Maximum number of events to return include_payload: Whether to include event payloads in response @@ -795,11 +814,15 @@ def list_events( # Get all events events = client.list_events(memory_id, actor_id, session_id) - # Get only main branch events - main_events = client.list_events(memory_id, actor_id, session_id, branch_name="main") - - # Get events from a specific branch - branch_events = client.list_events(memory_id, actor_id, session_id, branch_name="test-branch") + # Get events filtered by metadata + events = client.list_events( + memory_id, actor_id, session_id, + event_metadata=[{ + "left": {"metadataKey": "location"}, + "operator": "EQUALS_TO", + "right": {"metadataValue": {"stringValue": "NYC"}} + }] + ) """ try: all_events = [] @@ -817,11 +840,19 @@ def list_events( if next_token: params["nextToken"] = next_token + # Build filter map + filter_map = {} + # Add branch filter if specified (but not for "main") if branch_name and branch_name != "main": - params["filter"] = { - "branch": {"name": branch_name, "includeParentBranches": include_parent_branches} - } + filter_map["branch"] = {"name": branch_name, "includeParentBranches": include_parent_branches} + + # Add event metadata filter if specified + if event_metadata: + filter_map["eventMetadata"] = event_metadata + + if filter_map: + params["filter"] = filter_map response = self.gmdp_client.list_events(**params) @@ -1178,6 +1209,7 @@ def fork_conversation( branch_name: str, new_messages: List[Tuple[str, str]], event_timestamp: Optional[datetime] = None, + metadata: Optional[Dict[str, MetadataValue]] = None, ) -> Dict[str, Any]: """Fork a conversation from a specific event to create a new branch.""" try: @@ -1190,6 +1222,7 @@ def fork_conversation( messages=new_messages, branch=branch, event_timestamp=event_timestamp, + metadata=metadata, ) logger.info("Created branch '%s' from event %s", branch_name, root_event_id) diff --git a/tests/bedrock_agentcore/memory/test_client.py b/tests/bedrock_agentcore/memory/test_client.py index 6c20a08..2dca1e2 100644 --- a/tests/bedrock_agentcore/memory/test_client.py +++ b/tests/bedrock_agentcore/memory/test_client.py @@ -573,6 +573,104 @@ def test_list_events_with_branch_filter(): assert kwargs["filter"]["branch"]["includeParentBranches"] is True +def test_list_events_with_event_metadata_filter(): + """Test list_events with event metadata filtering.""" + with patch("boto3.client"): + client = MemoryClient() + + # Mock the client + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # Mock response with events containing metadata + mock_events = [ + { + "eventId": "event-nyc-1", + "eventTimestamp": datetime(2023, 1, 1, 10, 0, 0), + "metadata": {"location": {"stringValue": "NYC"}}, + "payload": [{"conversational": {"role": "USER", "content": {"text": "NYC message"}}}], + }, + { + "eventId": "event-nyc-2", + "eventTimestamp": datetime(2023, 1, 1, 10, 1, 0), + "metadata": {"location": {"stringValue": "NYC"}}, + "payload": [{"conversational": {"role": "USER", "content": {"text": "Another NYC message"}}}], + }, + ] + mock_gmdp.list_events.return_value = {"events": mock_events, "nextToken": None} + + # Test with event metadata filter + event_metadata_filter = [ + { + "left": {"metadataKey": "location"}, + "operator": "EQUALS_TO", + "right": {"metadataValue": {"stringValue": "NYC"}}, + } + ] + events = client.list_events( + memory_id="mem-123", + actor_id="user-123", + session_id="session-456", + event_metadata=event_metadata_filter, + ) + + assert len(events) == 2 + assert events[0]["eventId"] == "event-nyc-1" + assert events[1]["eventId"] == "event-nyc-2" + + # Verify filter was applied correctly + args, kwargs = mock_gmdp.list_events.call_args + assert "filter" in kwargs + assert "eventMetadata" in kwargs["filter"] + assert kwargs["filter"]["eventMetadata"] == event_metadata_filter + + +def test_list_events_with_branch_and_event_metadata_filter(): + """Test list_events with both branch and event metadata filtering.""" + with patch("boto3.client"): + client = MemoryClient() + + # Mock the client + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # Mock response + mock_events = [ + { + "eventId": "event-branch-meta-1", + "eventTimestamp": datetime(2023, 1, 1, 10, 0, 0), + "branch": {"name": "test-branch", "rootEventId": "event-0"}, + "metadata": {"location": {"stringValue": "NYC"}}, + } + ] + mock_gmdp.list_events.return_value = {"events": mock_events, "nextToken": None} + + # Test with both filters + event_metadata_filter = [ + { + "left": {"metadataKey": "location"}, + "operator": "EQUALS_TO", + "right": {"metadataValue": {"stringValue": "NYC"}}, + } + ] + events = client.list_events( + memory_id="mem-123", + actor_id="user-123", + session_id="session-456", + branch_name="test-branch", + include_parent_branches=True, + event_metadata=event_metadata_filter, + ) + + assert len(events) == 1 + + # Verify both filters were applied + args, kwargs = mock_gmdp.list_events.call_args + assert "filter" in kwargs + assert kwargs["filter"]["branch"]["name"] == "test-branch" + assert kwargs["filter"]["eventMetadata"] == event_metadata_filter + + def test_list_events_max_results_limit(): """Test list_events respects max_results limit.""" with patch("boto3.client"): @@ -1758,6 +1856,46 @@ def test_fork_conversation(): assert len(kwargs["payload"]) == 2 +def test_fork_conversation_with_metadata(): + """Test fork_conversation with metadata parameter.""" + with patch("boto3.client"): + client = MemoryClient() + + # Mock the client + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # Mock create_event response with metadata + metadata = {"fork_reason": {"stringValue": "alternative_response"}} + mock_gmdp.create_event.return_value = { + "event": { + "eventId": "event-fork-meta-123", + "memoryId": "mem-123", + "metadata": metadata, + } + } + + # Test fork_conversation with metadata + result = client.fork_conversation( + memory_id="mem-123", + actor_id="user-123", + session_id="session-456", + root_event_id="event-root-456", + branch_name="test-branch", + new_messages=[("Forked message", "USER")], + metadata=metadata, + ) + + assert result["eventId"] == "event-fork-meta-123" + assert result["metadata"] == metadata + + # Verify metadata was passed correctly along with branch info + args, kwargs = mock_gmdp.create_event.call_args + assert kwargs["metadata"] == metadata + assert kwargs["branch"]["name"] == "test-branch" + assert kwargs["branch"]["rootEventId"] == "event-root-456" + + def test_delete_strategy(): """Test delete_strategy functionality.""" with patch("boto3.client"): @@ -1878,6 +2016,84 @@ def test_create_event_with_branch(): assert kwargs["branch"] == branch +def test_create_event_with_metadata(): + """Test create_event with metadata parameter.""" + with patch("boto3.client"): + client = MemoryClient() + + # Mock the client + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # Mock create_event response with metadata + mock_gmdp.create_event.return_value = { + "event": { + "eventId": "event-meta-123", + "memoryId": "mem-123", + "metadata": {"location": {"stringValue": "NYC"}}, + } + } + + # Test create_event with metadata + metadata = {"location": {"stringValue": "NYC"}} + result = client.create_event( + memory_id="mem-123", + actor_id="user-123", + session_id="session-456", + messages=[("Hello from NYC", "USER")], + metadata=metadata, + ) + + assert result["eventId"] == "event-meta-123" + assert result["metadata"] == metadata + + # Verify metadata was passed correctly + args, kwargs = mock_gmdp.create_event.call_args + assert kwargs["metadata"] == metadata + + +def test_create_event_with_multiple_metadata_keys(): + """Test create_event with multiple metadata key-value pairs.""" + with patch("boto3.client"): + client = MemoryClient() + + # Mock the client + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # Mock create_event response + metadata = { + "location": {"stringValue": "NYC"}, + "category": {"stringValue": "weather"}, + "priority": {"stringValue": "high"}, + } + mock_gmdp.create_event.return_value = { + "event": { + "eventId": "event-multi-meta-123", + "memoryId": "mem-123", + "metadata": metadata, + } + } + + # Test create_event with multiple metadata keys + result = client.create_event( + memory_id="mem-123", + actor_id="user-123", + session_id="session-456", + messages=[("Weather check", "USER")], + metadata=metadata, + ) + + assert result["eventId"] == "event-multi-meta-123" + assert result["metadata"]["location"]["stringValue"] == "NYC" + assert result["metadata"]["category"]["stringValue"] == "weather" + assert result["metadata"]["priority"]["stringValue"] == "high" + + # Verify all metadata keys were passed + args, kwargs = mock_gmdp.create_event.call_args + assert len(kwargs["metadata"]) == 3 + + def test_create_memory_and_wait_client_error(): """Test create_memory_and_wait with ClientError during status check.""" with patch("boto3.client"): @@ -2736,6 +2952,44 @@ def test_create_blob_event_with_branch(): assert kwargs["branch"] == branch +def test_create_blob_event_with_metadata(): + """Test create_blob_event with metadata parameter.""" + with patch("boto3.client"): + client = MemoryClient() + + # Mock the client + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + # Mock create_event response with metadata + metadata = {"file_type": {"stringValue": "pdf"}} + mock_gmdp.create_event.return_value = { + "event": { + "eventId": "event-blob-meta-123", + "memoryId": "mem-123", + "metadata": metadata, + } + } + + # Test create_blob_event with metadata + blob_data = {"file_content": "base64_encoded_data"} + result = client.create_blob_event( + memory_id="mem-123", + actor_id="user-123", + session_id="session-456", + blob_data=blob_data, + metadata=metadata, + ) + + assert result["eventId"] == "event-blob-meta-123" + assert result["metadata"] == metadata + + # Verify metadata was passed correctly + args, kwargs = mock_gmdp.create_event.call_args + assert kwargs["metadata"] == metadata + assert "blob" in kwargs["payload"][0] + + def test_create_blob_event_client_error(): """Test create_blob_event with ClientError.""" with patch("boto3.client"):