Batch indexing, fix tests

Author: abrookins

agent-memory-client/tests/test_client.py

@@ -546,9 +546,9 @@ async def test_append_messages_to_working_memory(self, enhanced_test_client):
             # Check that messages were appended
             working_memory_arg = mock_put.call_args[0][1]
             assert len(working_memory_arg.messages) == 3
-            assert working_memory_arg.messages[0]["content"] == "First message"
-            assert working_memory_arg.messages[1]["content"] == "Second message"
-            assert working_memory_arg.messages[2]["content"] == "Third message"
+            assert working_memory_arg.messages[0].content == "First message"
+            assert working_memory_arg.messages[1].content == "Second message"
+            assert working_memory_arg.messages[2].content == "Third message"
 
     def test_deep_merge_dicts(self, enhanced_test_client):
         """Test the deep merge dictionary utility method."""

agent_memory_server/api.py

@@ -327,7 +327,9 @@ async def put_working_memory(
     )
 
     # Background tasks for long-term memory promotion and indexing (if enabled)
-    if settings.long_term_memory and updated_memory.memories:
+    if settings.long_term_memory and (
+        updated_memory.memories or updated_memory.messages
+    ):
         # Promote structured memories from working memory to long-term storage
         await background_tasks.add_task(
             long_term_memory.promote_working_memory_to_long_term,

agent_memory_server/long_term_memory.py

@@ -1326,6 +1326,7 @@ async def promote_working_memory_to_long_term(
 
     # Process unpersisted messages
     updated_messages = []
+    memory_records_to_index = []
     for msg in current_working_memory.messages:
         if msg.persisted_at is None:
             # Generate ID if not present (backward compatibility)
@@ -1352,12 +1353,8 @@ async def promote_working_memory_to_long_term(
             current_memory = deduped_memory or memory_record
             current_memory.persisted_at = datetime.now(UTC)
 
-            # Index in long-term storage
-            await index_long_term_memories(
-                [current_memory],
-                redis_client=redis,
-                deduplicate=False,  # Already deduplicated by ID
-            )
+            # Collect memory record for batch indexing
+            memory_records_to_index.append(current_memory)
 
             # Update message with persisted_at timestamp
             msg.persisted_at = current_memory.persisted_at
@@ -1370,6 +1367,14 @@ async def promote_working_memory_to_long_term(
 
         updated_messages.append(msg)
 
+    # Batch index all new memory records for messages
+    if memory_records_to_index:
+        await index_long_term_memories(
+            memory_records_to_index,
+            redis_client=redis,
+            deduplicate=False,  # Already deduplicated by ID
+        )
+
     # Update working memory with the new persisted_at timestamps and extracted memories
     if promoted_count > 0 or extracted_memories:
         updated_working_memory = current_working_memory.model_copy()

tests/test_api.py

@@ -6,7 +6,6 @@
 
 from agent_memory_server.config import Settings
 from agent_memory_server.long_term_memory import (
-    index_long_term_memories,
     promote_working_memory_to_long_term,
 )
 from agent_memory_server.models import (
@@ -153,7 +152,11 @@ async def test_put_memory(self, client):
             "/v1/working-memory/test-session?namespace=test-namespace"
         )
         assert updated_session.status_code == 200
-        assert updated_session.json()["messages"] == payload["messages"]
+        retrieved_messages = updated_session.json()["messages"]
+        assert len(retrieved_messages) == len(payload["messages"])
+        for i, msg in enumerate(retrieved_messages):
+            assert msg["role"] == payload["messages"][i]["role"]
+            assert msg["content"] == payload["messages"][i]["content"]
 
     @pytest.mark.requires_api_keys
     @pytest.mark.asyncio
@@ -188,10 +191,10 @@ async def test_put_memory_stores_messages_in_long_term_memory(
         # Check that background tasks were called
         assert mock_background_tasks.add_task.call_count == 1
 
-        # Check that the last call was for long-term memory indexing
+        # Check that the last call was for long-term memory promotion
         assert (
             mock_background_tasks.add_task.call_args_list[-1][0][0]
-            == index_long_term_memories
+            == promote_working_memory_to_long_term
         )
 
     @pytest.mark.requires_api_keys