feat: integrate metadata support for short-term-memory (STM) (#114)

Co-authored-by: Sukruth Gowdru Lingaraju <[email protected]>

Author: glsukki

src/bedrock_agentcore/memory/models/__init__.py

@@ -3,7 +3,15 @@
 from typing import Any, Dict
 
 from .DictWrapper import DictWrapper
-
+from .filters import (
+    StringValue,
+    MetadataValue,
+    MetadataKey,
+    LeftExpression,
+    OperatorType,
+    RightExpression,
+    EventMetadataFilter,
+)
 
 class ActorSummary(DictWrapper):
     """A class representing an actor summary."""
@@ -75,3 +83,20 @@ def __init__(self, session_summary: Dict[str, Any]):
             session_summary: Dictionary containing session summary data.
         """
         super().__init__(session_summary)
+
+__all__ = [
+    "DictWrapper",
+    "ActorSummary",
+    "Branch",
+    "Event",
+    "EventMessage",
+    "MemoryRecord",
+    "SessionSummary",
+    "StringValue",
+    "MetadataValue",
+    "MetadataKey",
+    "LeftExpression",
+    "OperatorType",
+    "RightExpression",
+    "EventMetadataFilter",
+]

src/bedrock_agentcore/memory/models/filters.py

@@ -0,0 +1,118 @@
+from enum import Enum
+from typing import Optional, TypedDict, Union, NotRequired
+
+class StringValue(TypedDict):
+    """Value associated with the `eventMetadata` key."""
+    stringValue: str
+    
+    @staticmethod
+    def build(value: str) -> 'StringValue':
+        return {
+            "stringValue": value
+        }
+
+MetadataValue = Union[StringValue]
+"""
+Union type representing metadata values.
+
+Variants:
+- StringValue: {"stringValue": str} - String metadata value
+"""
+
+MetadataKey = Union[str]
+"""
+Union type representing metadata key.
+"""
+
+class LeftExpression(TypedDict):
+    """
+    Left operand of the event metadata filter expression.
+    """
+    metadataKey: MetadataKey
+    
+    @staticmethod
+    def build(key: str) -> 'LeftExpression':
+        """Builds the `metadataKey` for `LeftExpression`"""
+        return {
+            "metadataKey": key
+        }
+
+class OperatorType(Enum):
+    """
+    Operator applied to the event metadata filter expression.
+    
+    Currently supports:
+    - `EQUALS_TO`
+    - `EXISTS`
+    - `NOT_EXISTS`
+    """
+    EQUALS_TO = "EQUALS_TO"
+    EXISTS = "EXISTS"
+    NOT_EXISTS = "NOT_EXISTS"
+
+class RightExpression(TypedDict):
+    """
+    Right operand of the event metadata filter expression.
+    
+    Variants:
+    - StringValue: {"metadataValue": {"stringValue": str}}
+    """
+    metadataValue: MetadataValue
+
+    @staticmethod
+    def build(value: str) -> 'RightExpression':
+        """Builds the `RightExpression` for `stringValue` type"""
+        return {"metadataValue": StringValue.build(value)}
+
+class EventMetadataFilter(TypedDict):
+    """
+    Filter expression for retrieving events based on metadata associated with an event.
+    
+    Args:
+        left: `LeftExpression` of the event metadata filter expression.
+        operator: `OperatorType` applied to the event metadata filter expression.
+        right: Optional `RightExpression` of the event metadata filter expression.
+    """
+    left: LeftExpression
+    operator: OperatorType
+    right: NotRequired[RightExpression]
+    
+    def build_expression(left_operand: LeftExpression, operator: OperatorType, right_operand: Optional[RightExpression] = None) -> 'EventMetadataFilter':
+        """
+        This method builds the required event metadata filter expression into the `EventMetadataFilterExpression` type when querying listEvents.
+        
+        Args:  
+            left_operand: Left operand of the event metadata filter expression
+            operator: Operator applied to the event metadata filter expression
+            right_operand: Optional right_operand of the event metadata filter expression.
+        
+        Example:
+        ```
+            left_operand = LeftExpression.build_key(key='location')
+            operator = OperatorType.EQUALS_TO
+            right_operand = RightExpression.build_string_value(value='NYC')
+        ```
+
+        #### Response Object:
+        ```
+            {
+                'left': {
+                    'metadataKey': 'location'
+                },
+                'operator': 'EQUALS_TO',
+                'right': {
+                    'metadataValue': {
+                        'stringValue': 'NYC'
+                    }
+                }
+            }
+        ```
+        """
+        filter = {
+            'left': left_operand,
+            'operator': operator.value
+        }
+        
+        if right_operand:
+            filter['right'] = right_operand
+        return filter

src/bedrock_agentcore/memory/session.py

@@ -18,6 +18,8 @@
     EventMessage,
     MemoryRecord,
     SessionSummary,
+    MetadataValue,
+    EventMetadataFilter
 )
 
 logger = logging.getLogger(__name__)
@@ -246,6 +248,7 @@ def process_turn_with_llm(
         user_input: str,
         llm_callback: Callable[[str, List[Dict[str, Any]]], str],
         retrieval_config: Optional[Dict[str, RetrievalConfig]],
+        metadata: Optional[Dict[str, MetadataValue]] = None,
         event_timestamp: Optional[datetime] = None,
     ) -> Tuple[List[Dict[str, Any]], str, Dict[str, Any]]:
         r"""Complete conversation turn with LLM callback integration.
@@ -263,6 +266,7 @@ def process_turn_with_llm(
             retrieval_config: Optional dictionary mapping namespaces to RetrievalConfig objects.
                             Each namespace can contain template variables like {actorId}, {sessionId},
                             {memoryStrategyId} that will be resolved at runtime.
+            metadata: Optional custom key-value metadata to attach to an event.
             event_timestamp: Optional timestamp for the event
 
         Returns:
@@ -340,6 +344,7 @@ def my_llm(user_input: str, memories: List[Dict]) -> str:
                 ConversationalMessage(user_input, MessageRole.USER),
                 ConversationalMessage(agent_response, MessageRole.ASSISTANT),
             ],
+            metadata=metadata,
             event_timestamp=event_timestamp,
         )
 
@@ -352,6 +357,7 @@ def add_turns(
         session_id: str,
         messages: List[Union[ConversationalMessage, BlobMessage]],
         branch: Optional[Dict[str, str]] = None,
+        metadata: Optional[Dict[str, MetadataValue]] = None,
         event_timestamp: Optional[datetime] = None,
     ) -> Event:
         """Adds conversational turns or blob objects to short-term memory.
@@ -365,21 +371,31 @@ def add_turns(
                 - ConversationalMessage objects for conversational messages
                 - BlobMessage objects for blob data
             branch: Optional branch info
+            metadata: Optional custom key-value metadata to attach to an event.
             event_timestamp: Optional timestamp for the event
 
         Returns:
             Created event
 
         Example:
+        ```
             manager.add_turns(
                 actor_id="user-123",
                 session_id="session-456",
                 messages=[
                     ConversationalMessage("Hello", USER),
                     BlobMessage({"file_data": "base64_content"}),
                     ConversationalMessage("How can I help?", ASSISTANT)
+                ],
+                metadata=[
+                    {
+                        'location': {
+                            'stringValue': 'NYC'
+                        }
+                    }
                 ]
             )
+        ```
         """
         logger.info("  -> Storing %d messages in short-term memory...", len(messages))
 
@@ -412,6 +428,10 @@ def add_turns(
 
         if branch:
             params["branch"] = branch
+        
+        if metadata:
+            params["metadata"] = metadata
+
         try:
             response = self._data_plane_client.create_event(**params)
             logger.info("     ✅ Turn stored successfully with Event ID: %s", response.get("eventId"))
@@ -427,6 +447,7 @@ def fork_conversation(
         root_event_id: str,
         branch_name: str,
         messages: List[Union[ConversationalMessage, BlobMessage]],
+        metadata: Optional[Dict[str, MetadataValue]] = None,
         event_timestamp: Optional[datetime] = None,
     ) -> Dict[str, Any]:
         """Fork a conversation from a specific event to create a new branch."""
@@ -439,6 +460,7 @@ def fork_conversation(
                 messages=messages,
                 event_timestamp=event_timestamp,
                 branch=branch,
+                metadata=metadata,
             )
 
             logger.info("Created branch '%s' from event %s", branch_name, root_event_id)
@@ -454,6 +476,7 @@ def list_events(
         session_id: str,
         branch_name: Optional[str] = None,
         include_parent_branches: bool = False,
+        eventMetadata: Optional[List[EventMetadataFilter]] = None,
         max_results: int = 100,
         include_payload: bool = True,
     ) -> List[Event]:
@@ -482,6 +505,49 @@ def list_events(
 
             # Get events from a specific branch
             branch_events = client.list_events(actor_id, session_id, branch_name="test-branch")
+            
+            #### Get events with event metadata filter
+            ```
+            filtered_events_with_metadata = client.list_events(
+                actor_id=actor_id,
+                session_id=session_id, 
+                eventMetadata=[
+                    {
+                        'left': {
+                            'metadataKey': 'location'
+                        },
+                        'operator': 'EQUALS_TO',
+                        'right': {
+                            'metadataValue': {
+                                'stringValue': 'NYC'
+                            }
+                        }
+                    }       
+                ]
+            )
+            ```
+
+            #### Get events with event metadata filter + specific branch filter
+            ```
+            branch_with_metadata_filtered_events = client.list_events(
+                actor_id=actor_id,
+                session_id=session_id,
+                branch_name="test-branch",
+                eventMetadata=[
+                    {
+                        'left': {
+                            'metadataKey': 'location'
+                        },
+                        'operator': 'EQUALS_TO',
+                        'right': {
+                            'metadataValue': {
+                                'stringValue': 'NYC'
+                            }
+                        }
+                    }       
+                ]
+            )
+            ```
         """
         try:
             all_events: List[Event] = []
@@ -509,6 +575,12 @@ def list_events(
                         "branch": {"name": branch_name, "includeParentBranches": include_parent_branches}
                     }
 
+                # Add eventMetadata filter if specified
+                if eventMetadata:
+                    params["filter"] = {
+                        "eventMetadata": eventMetadata
+                    }
+
                 response = self._data_plane_client.list_events(**params)
 
                 events = response.get("events", [])
@@ -888,28 +960,31 @@ def add_turns(
         self,
         messages: List[Union[ConversationalMessage, BlobMessage]],
         branch: Optional[Dict[str, str]] = None,
+        metadata: Optional[Dict[str, MetadataValue]] = None,
         event_timestamp: Optional[datetime] = None,
     ) -> Event:
         """Delegates to manager.add_turns."""
-        return self._manager.add_turns(self._actor_id, self._session_id, messages, branch, event_timestamp)
+        return self._manager.add_turns(self._actor_id, self._session_id, messages, branch, metadata, event_timestamp)
 
     def fork_conversation(
         self,
         messages: List[Union[ConversationalMessage, BlobMessage]],
         root_event_id: str,
         branch_name: str,
+        metadata: Optional[Dict[str, MetadataValue]] = None,
         event_timestamp: Optional[datetime] = None,
     ) -> Event:
         """Delegates to manager.fork_conversation."""
         return self._manager.fork_conversation(
-            self._actor_id, self._session_id, root_event_id, branch_name, messages, event_timestamp
+            self._actor_id, self._session_id, root_event_id, branch_name, messages, metadata, event_timestamp
         )
 
     def process_turn_with_llm(
         self,
         user_input: str,
         llm_callback: Callable[[str, List[Dict[str, Any]]], str],
         retrieval_config: Optional[Dict[str, RetrievalConfig]],
+        metadata: Optional[Dict[str, MetadataValue]] = None,
         event_timestamp: Optional[datetime] = None,
     ) -> Tuple[List[Dict[str, Any]], str, Dict[str, Any]]:
         """Delegates to manager.process_turn_with_llm."""
@@ -919,6 +994,7 @@ def process_turn_with_llm(
             user_input,
             llm_callback,
             retrieval_config,
+            metadata,
             event_timestamp,
         )
 
@@ -975,6 +1051,7 @@ def list_events(
         self,
         branch_name: Optional[str] = None,
         include_parent_branches: bool = False,
+        eventMetadata: Optional[List[EventMetadataFilter]] = None,
         max_results: int = 100,
         include_payload: bool = True,
     ) -> List[Event]:
@@ -984,6 +1061,7 @@ def list_events(
             session_id=self._session_id,
             branch_name=branch_name,
             include_parent_branches=include_parent_branches,
+            eventMetadata=eventMetadata,
             include_payload=include_payload,
             max_results=max_results,
         )