fix(langgraph-checkpoint-aws): Add max_results and limit attributes to AgentCoreMemorySaver (#728)

Addresses #719.

This PR updates the AgentCore memory checkpointer integration to:
- Add an `max_results` attribute to `AgentCoreMemorySaver`, allowing
control of the `maxResults` value passed to [ListEvents
API](https://docs.aws.amazon.com/bedrock-agentcore/latest/APIReference/API_ListEvents.html).
- Effective default will remain at the previously hardcoded limit of 100
results.
- Add another `limit` attribute to control the maximum events parsed
when paginating on the ListEvents response.
  - Effective default will move from 100 -> `None` (unlimited).
- In `AgentCoreEventClient,get_events()`, exit from the response
processing loop early if the number of events found exceeds the `limit`
attribute value. If this occurs, a warning will be returned to the user.

Author: michaelnchin

libs/langgraph-checkpoint-aws/langgraph_checkpoint_aws/agentcore/helpers.py

@@ -8,6 +8,7 @@
 import datetime
 import json
 import logging
+import warnings
 from collections import defaultdict
 from typing import Any, cast
 
@@ -184,22 +185,37 @@ def store_blob_events_batch(
         )
 
     def get_events(
-        self, session_id: str, actor_id: str, limit: int = 100
+        self,
+        session_id: str,
+        actor_id: str,
+        limit: int | None = None,
+        max_results: int | None = 100,
     ) -> list[EventType]:
-        """Retrieve events from AgentCore Memory."""
+        """Retrieve events from AgentCore Memory.
 
-        if limit is not None and limit <= 0:
+        Args:
+            session_id: The session ID to retrieve events for
+            actor_id: The actor ID to retrieve events for
+            limit: The maximum number of events to parse from ListEvents
+            max_results: Maximum number of results to retrieve. Defaults to 100.
+
+        Returns:
+            List of retrieved events
+        """
+
+        if max_results is not None and max_results <= 0:
             return []
 
         all_events = []
         next_token = None
+        limit_reached = False
 
         while True:
             params = {
                 "memoryId": self.memory_id,
                 "actorId": actor_id,
                 "sessionId": session_id,
-                "maxResults": 100,
+                "maxResults": max_results,
                 "includePayloads": True,
             }
 
@@ -218,8 +234,26 @@ def get_events(
                         except EventDecodingError as e:
                             logger.warning(f"Failed to decode event: {e}")
 
+                        if limit is not None and len(all_events) >= limit:
+                            limit_reached = True
+                            break
+
+                if limit_reached:
+                    break
+
             next_token = response.get("nextToken")
-            if not next_token or (limit is not None and len(all_events) >= limit):
+
+            if limit_reached and next_token:
+                warnings.warn(
+                    f"Stopped retrieving events at limit of {limit}. "
+                    f"There may be additional checkpoints that were not retrieved. "
+                    f"Consider increasing the limit parameter, or set None for no "
+                    f"limit.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+            if limit_reached or not next_token:
                 break
 
         return all_events

libs/langgraph-checkpoint-aws/langgraph_checkpoint_aws/agentcore/saver.py

@@ -49,26 +49,43 @@ class AgentCoreMemorySaver(BaseCheckpointSaver[str]):
     Args:
         memory_id: the ID of the memory resource created in AgentCore Memory
         serde: serialization protocol to be used. Defaults to JSONPlusSerializer
+        limit: maximum number of events to parse from ListEvents.
+        max_results: maximum number of results to retrieve from AgentCore Memory.
+                    Defaults to 100
     """
 
     def __init__(
         self,
         memory_id: str,
         *,
         serde: SerializerProtocol | None = None,
+        limit: int | None = None,
+        max_results: int | None = 100,
         **boto3_kwargs: Any,
     ) -> None:
         super().__init__(serde=serde)
 
         self.memory_id = memory_id
+        self.limit = limit
+        self.max_results = max_results
         self.serializer = EventSerializer(self.serde)
         self.checkpoint_event_client = AgentCoreEventClient(
             memory_id, self.serializer, **boto3_kwargs
         )
         self.processor = EventProcessor()
 
-    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
-        """Get a checkpoint tuple from Bedrock AgentCore Memory."""
+    def get_tuple(
+        self,
+        config: RunnableConfig,
+    ) -> CheckpointTuple | None:
+        """Get a checkpoint tuple from Bedrock AgentCore Memory.
+
+        Args:
+            config: The runnable config containing checkpoint information
+
+        Returns:
+            CheckpointTuple if found, None otherwise
+        """
 
         # TODO: There is room for caching here on the client side
 
@@ -77,7 +94,10 @@ def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
         )
 
         events = self.checkpoint_event_client.get_events(
-            checkpoint_config.session_id, checkpoint_config.actor_id
+            checkpoint_config.session_id,
+            checkpoint_config.actor_id,
+            self.limit,
+            self.max_results,
         )
 
         checkpoints, writes_by_checkpoint, channel_data = self.processor.process_events(
@@ -122,7 +142,8 @@ def list(
         events = self.checkpoint_event_client.get_events(
             checkpoint_config.session_id,
             checkpoint_config.actor_id,
-            100 if limit is None else limit,
+            limit,
+            self.max_results,
         )
 
         checkpoints, writes_by_checkpoint, channel_data = self.processor.process_events(