Change how message filtering happens

Split checkpoint into interactions and filter all but the last one
Identify tool-related AIMessages by response_metadata

Author: mbklein

chat/src/core/setup.py

@@ -14,7 +14,7 @@ def chat_model(**kwargs) -> BaseModel:
 
 def checkpoint_saver(**kwargs) -> BaseCheckpointSaver:
     checkpoint_bucket: str = os.getenv("CHECKPOINT_BUCKET_NAME")
-    return SelectiveCheckpointer(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

@@ -9,31 +9,41 @@
 )
 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.additional_kwargs.get('stop_reason', '') == 'tool_use':
+    if isinstance(message, AIMessage) and message.response_metadata.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
-        ]
+    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 messages
+    # 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."""
@@ -44,7 +54,7 @@ def __init__(
         region_name: str = os.getenv("AWS_REGION"),
         endpoint_url: Optional[str] = None,
         compression: Optional[str] = None,
-        retain_history: Optional[bool] = False,
+        retain_history: Optional[bool] = True,
     ) -> None:
         super().__init__(bucket_name, region_name, endpoint_url, compression)
         self.retain_history = retain_history

chat/test/core/test_setup.py

@@ -23,6 +23,7 @@ def test_checkpoint_saver_initialization(self, mock_checkpointer):
 
         mock_checkpointer.assert_called_once_with(
             bucket_name="test-bucket",
+            retain_history=False,
             **kwargs
         )
         self.assertEqual(result, mock_checkpointer.return_value)