Bug: AgentCoreMemoryStore `_handle_put` not using `op.key`

[ltoniazzi]

Error

I cannot retrieve something that I store running this basic example:

from langgraph_checkpoint_aws import AgentCoreMemoryStore
from langchain_core.messages import HumanMessage

MEMORY_ID = "a-memory-id"
ACTOR_ID = "a-actor-id"
THREAD_ID = "a-thread-id"
REGION = "us-west-1"

store = AgentCoreMemoryStore(memory_id=MEMORY_ID, region_name=REGION)

namespace = (ACTOR_ID, THREAD_ID)
key = "mem-" + "1"*40
# Put
store.put(namespace, key, {"message": HumanMessage(content="User preference")})     
# Get
fetched_message = store.get(namespace, key)

assert fetched_message, "No message found!"  # This gets raised

Bug here?

It seems to me that there might be an issue here as events are created as:

self.client.create_event(
            memoryId=self.memory_id,
            actorId=actor_id,
            sessionId=session_id,
            eventTimestamp=datetime.now(timezone.utc),
            payload=conversational_payloads,
        )
        logger.debug(f"Created event for message in namespace {op.namespace}")

where actor_id, session_id = op.namespace.

But op.key does not seem to be passed in any way to self.client.create_event.

Then, store.get, tries to do:

response = self.client.get_memory_record(
                memoryId=self.memory_id, memoryRecordId=op.key
            )

But op.key seem to have never being passed.

(Also event and memory record seem different storage locations from the aws docs)

Am I getting confused? Or if it's a bug I can help fix it, let me know 😃

Versions

langgraph-checkpoint-aws = 0.2.0

[ps06756]

I think the key can be added in the metadata , since create_event doesn't directly support adding a key.

[ltoniazzi]

@ps06756 I am a bit confused as events (short term memory) seem different from memory_records (long-term memory).

Also, in the docs, I cannot find a connection between the parameters namespace-memoryRecordId of *_records within *_events methods .

I was looking for this connection because the strategy of store.put seems to force the namespace to be a 2-dim tuple assigned as:

actorId=op.namespace[0]
sessionId=op.namespace[1]

and then I was expecting that store.get would list all events for that namespace and fetch the one for the op.key (assuming op.key was passed in the payload for create_event)

So can you explain me what is the strategy behind mixing events and memory_records?

[ltoniazzi]

Also for both create_event and batch_create_memory_records, the corresponding id (eventId and memoryRecordId) seems to be assigned by AWS and provided in the response, and I do not see in the documentation how one can pass their own id.

[ltoniazzi]

Something that works for me is the AgentCoreMemoryStoreCustom below.

Main design:

• Only use events (though this means it's unclear how to use the op.query param for search)
• Assign a storeEventId=op.key as part of metadata for each event (though now it's ugly to guarantee uniqueness).
• Store: anything (not only BaseMessages)
• Search: by listing events for namespace
• Get: by listing events for namespace and filtering the desired storeEventId

Click to see AgentCoreMemoryStoreCustom

import base64
import json
import logging
from collections.abc import Iterable
from datetime import UTC, datetime
from typing import Any
from botocore.exceptions import ClientError

from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.store.base import (
    GetOp,
    Item,
    Op,
    PutOp,
    Result,
    SearchItem,
    SearchOp,
)
from langgraph_checkpoint_aws.agentcore import AgentCoreMemoryStore
from langgraph_checkpoint_aws.agentcore.constants import (
    EventDecodingError,
)
from langgraph_checkpoint_aws.agentcore.helpers import AgentCoreEventClient, EventSerializer


logger = logging.getLogger(__name__)


STORE_ID_VAR = "storeEventId"


class AgentCoreEventClientCustom(AgentCoreEventClient):
    def get_events(self, session_id: str, actor_id: str, limit: int = None) -> Any:
        """Retrieve events from AgentCore Memory."""

        if limit is not None and limit <= 0:
            return []

        all_events = []
        next_token = None

        while True:
            params = {
                "memoryId": self.memory_id,
                "actorId": actor_id,
                "sessionId": session_id,
                "maxResults": 100,
                "includePayloads": True,
            }

            if next_token:
                params["nextToken"] = next_token

            response = self.client.list_events(**params)

            for event in response.get("events", []):
                for payload_item in event.get("payload", []):
                    blob = payload_item.get("blob")
                    if blob:
                        try:
                            parsed_event = self.serializer.deserialize_value(
                                blob
                            )  # Main change as we do not store events, just values
                            all_events.append(parsed_event)
                            payload_item["blob"] = parsed_event
                        except EventDecodingError as e:
                            logger.warning(f"Failed to decode event: {e}")

            next_token = response.get("nextToken")
            if not next_token or (limit is not None and len(all_events) >= limit):
                break

        return response["events"]

class EventSerializerCustom(EventSerializer):
    def deserialize_value(self, serialized: str) -> Any:
        """Deserialize a value using the serde protocol."""
        try:
            serialized = json.loads(serialized)
            type_tag = serialized["type"]
            binary_data = base64.b64decode(serialized["data"])
            return self.serde.loads_typed((type_tag, binary_data))
        except Exception as e:
            raise EventDecodingError(f"Failed to deserialize value: {e}") from e  # noqa: EM102

class AgentCoreMemoryStoreCustom(AgentCoreMemoryStore):
    def __init__(self, *, memory_id: str, **boto3_kwargs: Any) -> None:
        super().__init__(memory_id=memory_id, **boto3_kwargs)

        serializer = EventSerializerCustom(serde=JsonPlusSerializer())
        self.event_client = AgentCoreEventClientCustom(
            memory_id=self.memory_id, serializer=serializer, **boto3_kwargs
        )

    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
        """Execute multiple operations asynchronously (not implemented)."""
        logger.warning(
            "Asynchronous operations not supported for AgentCoreMemoryStore. " "Falling back to synchronous batch."
        )
        return self.batch(ops)

    def _handle_put(self, op: PutOp) -> None:
        """Handle PutOp by creating conversational events in AgentCore Memory."""
        if op.value is None:
            # TODO: Delete operation support - need to figure out if we are deleting events or records # noqa: E501
            logger.warning("Delete operations not supported in AgentCore Memory")
            return

        # Convert namespace tuple to actor_id and session_id
        if len(op.namespace) != 2:
            raise ValueError("Namespace must be a tuple of (actor_id, session_id)")

        actor_id, session_id = op.namespace

        payload = []
        serialized = self.event_client.serializer.serialize_value(op.value)
        payload.append({"blob": json.dumps(serialized)})

        self.client.create_event(
            memoryId=self.memory_id,
            actorId=actor_id,
            sessionId=session_id,
            eventTimestamp=datetime.now(UTC),
            payload=payload,
            metadata={STORE_ID_VAR: {"stringValue": op.key}},
        )

        logger.debug(f"Created event for message in namespace {op.namespace}")

    def _handle_get(self, op: GetOp) -> Item | None:
        """Handle GetOp by retrieving a specific memory record from AgentCore Memory."""
        try:
            actor_id, session_id = op.namespace
            events = self.event_client.get_events(session_id=session_id, actor_id=actor_id)

            store_item = self.get_store_item_from_events(events, op)

            if not store_item:
                logger.warning(f"Store item not found for key: {op.key}")
                return None

            return store_item

        except ClientError as e:
            error_code = e.response["Error"]["Code"]
            if error_code == "ResourceNotFoundException":
                # Memory record not found
                return None
            # Re-raise other client errors
            logger.error(f"Failed to get memory record: {e}")
            raise
        except Exception as e:
            logger.error(f"Failed to get memory record: {e}")
            raise

    def _handle_search(self, op: SearchOp) -> list[SearchItem]:
        """Handle SearchOp by retrieving memory records from AgentCore Memory."""
        if op.query:
            logger.warning("Search requires not currently supported for AgentCore Memory Custom.")

        try:
            actor_id, session_id = op.namespace_prefix
            events = self.event_client.get_events(session_id=session_id, actor_id=actor_id)
            return self._convert_events_to_search_items(events, op.namespace_prefix)
        except ClientError as e:
            error_code = e.response["Error"]["Code"]
            if error_code == "ResourceNotFoundException":
                # Memory record not found
                return []
            # Re-raise other client errors
            logger.error(f"Failed to get memory records: {e}")
            raise

    def _convert_events_to_search_items(self, events: list[dict], namespace: tuple[str, ...]) -> list[SearchItem]:
        """Convert AgentCore memory records to SearchItem objects."""
        results = []

        for event in events:
            search_item = SearchItem(
                namespace=namespace,
                key=self._get_store_event_id(event),
                value=event["payload"][0]["blob"],
                created_at=event["eventTimestamp"],
                updated_at=event["eventTimestamp"],  # TODO what to do with this field?
                score=None,
            )
            results.append(search_item)

        return results

    def _get_store_event_id(self, event: dict) -> str | None:
        """Extract the store event ID from an AgentCore event's metadata."""
        metadata = event.get("metadata", {})
        store_id_entry = metadata.get(STORE_ID_VAR)
        if store_id_entry:
            return store_id_entry.get("stringValue")
        return None

    def get_store_item_from_events(self, events: list[dict[str, Any]], op: GetOp) -> dict[str, Any] | None:
        matched_items = []
        for event in events:
            if event["metadata"].get(STORE_ID_VAR):
                key_candidate = event["metadata"].get(STORE_ID_VAR).get("stringValue")
                if key_candidate == op.key:
                    matched_items.append(
                        Item(
                            **{
                                "key": op.key,
                                "value": event["payload"][0]["blob"],
                                "namespace": op.namespace,
                                "created_at": event["eventTimestamp"],
                                "updated_at": event["eventTimestamp"],  # TODO what to do with this field?
                            }
                        )
                    )

        if not matched_items:
            return None
        if len(matched_items) > 1:
            logger.warning(f"Multiple items found for key: {op.key}, returning the latest one.")
        # Return the last item TODO select based on created_at timestamp?
        return matched_items[-1]

[michaelnchin]

cc: @jgordley @bergjaak