Merge pull request #296 from nulib/5373-checkpoint-bloat

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, retain_history=False, **kwargs)
 
 def prefix(value):
     env_prefix = os.getenv("ENV_PREFIX")

chat/src/persistence/selective_checkpointer.py

@@ -0,0 +1,79 @@
+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
+
+
+# Split messages into interactions, each one starting with a HumanMessage
+def _split_interactions(messages):
+    if messages is None:
+        return []
+
+    interactions = []
+    current_interaction = []
+
+    for message in messages:
+        if isinstance(message, HumanMessage) and current_interaction:
+            interactions.append(current_interaction)
+            current_interaction = []
+        current_interaction.append(message)
+
+    if current_interaction:
+        interactions.append(current_interaction)
+
+    return interactions
+    
+def _is_tool_message(message):
+    if isinstance(message, ToolMessage):
+        return True
+    if isinstance(message, AIMessage) and message.response_metadata.get('stop_reason', '') == 'tool_use':
+        return True
+    return False
+
+def _prune_messages(messages):
+    interactions = _split_interactions(messages)
+    # Remove all tool-related messages except those related to the most recent interaction
+    for i, interaction in enumerate(interactions[:-1]):
+        interactions[i] = [message for message in interaction if not _is_tool_message(message)]
+
+    # Return the flattened list of messages
+    return [message for interaction in interactions for message in interaction]
+
+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] = True,
+    ) -> 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,17 @@ 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",
+            retain_history=False,
             **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):