Add selective checkpointer

Author: mbklein

chat/src/core/setup.py

@@ -1,4 +1,4 @@
-from persistence.s3_checkpointer import S3Checkpointer
+from persistence.selective_checkpointer import SelectiveCheckpointer
 from search.opensearch_neural_search import OpenSearchNeuralSearch
 from langchain_aws import ChatBedrock
 from langchain_core.language_models.base import BaseModel
@@ -14,7 +14,7 @@ def chat_model(**kwargs) -> BaseModel:
 
 def checkpoint_saver(**kwargs) -> BaseCheckpointSaver:
     checkpoint_bucket: str = os.getenv("CHECKPOINT_BUCKET_NAME")
-    return S3Checkpointer(bucket_name=checkpoint_bucket, **kwargs)
+    return SelectiveCheckpointer(bucket_name=checkpoint_bucket, **kwargs)
 
 def prefix(value):
     env_prefix = os.getenv("ENV_PREFIX")

chat/src/persistence/selective_checkpointer.py

@@ -0,0 +1,69 @@
+import os
+from typing import Optional
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import (
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata
+)
+from persistence.s3_checkpointer import S3Checkpointer
+
+def _is_tool_message(message):
+    if isinstance(message, ToolMessage):
+        return True
+    if isinstance(message, AIMessage) and message.additional_kwargs.get('stop_reason', '') == 'tool_use':
+        return True
+    return False
+
+def _prune_messages(messages):
+    if messages is None:
+        return messages
+
+    last_human_message = None
+    for i, message in reversed(list(enumerate(messages))):
+        if isinstance(message, HumanMessage):
+            last_human_message = i
+            break
+
+    if last_human_message is not None:
+        return [
+            msg
+            for i, msg in enumerate(messages)
+            if not _is_tool_message(msg) or i > last_human_message
+        ]
+
+    return messages
+
+class SelectiveCheckpointer(S3Checkpointer):
+    """S3 Checkpointer that discards ToolMessages from previous checkpoints."""
+  
+    def __init__(
+        self,
+        bucket_name: str,
+        region_name: str = os.getenv("AWS_REGION"),
+        endpoint_url: Optional[str] = None,
+        compression: Optional[str] = None,
+        retain_history: Optional[bool] = False,
+    ) -> None:
+        super().__init__(bucket_name, region_name, endpoint_url, compression)
+        self.retain_history = retain_history
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        # Remove previous checkpoints
+        thread_id = config["configurable"]["thread_id"]
+        if not self.retain_history:
+            self.delete_checkpoints(thread_id)
+        
+        # Remove all ToolMessages except those related to the most 
+        # recent question (HumanMessage)
+        messages = checkpoint.get("channel_values", {}).get("messages", [])
+        checkpoint["channel_values"]["messages"] = _prune_messages(messages)
+
+        return super().put(config, checkpoint, metadata, new_versions)

chat/test/core/test_setup.py

@@ -16,16 +16,16 @@ def test_chat_model_returns_bedrock_instance(self):
 
 class TestCheckpointSaver(unittest.TestCase):
     @patch.dict(os.environ, {"CHECKPOINT_BUCKET_NAME": "test-bucket"})
-    @patch("core.setup.S3Checkpointer")
-    def test_checkpoint_saver_initialization(self, mock_s3_checkpointer):
+    @patch("core.setup.SelectiveCheckpointer")
+    def test_checkpoint_saver_initialization(self, mock_checkpointer):
         kwargs = {"prefix": "test"}
         result = checkpoint_saver(**kwargs)
 
-        mock_s3_checkpointer.assert_called_once_with(
+        mock_checkpointer.assert_called_once_with(
             bucket_name="test-bucket",
             **kwargs
         )
-        self.assertEqual(result, mock_s3_checkpointer.return_value)
+        self.assertEqual(result, mock_checkpointer.return_value)
 
 class TestPrefix(unittest.TestCase):
     def test_prefix_with_env_prefix(self):