fix: handle missing 'data' key in memory payload during search operations (#3524)

Author: lsvishaal

mem0/memory/main.py

@@ -382,7 +382,7 @@ def _add_to_vector_store(self, messages, metadata, filters, infer):
                 filters=filters,
             )
             for mem in existing_memories:
-                retrieved_old_memory.append({"id": mem.id, "text": mem.payload["data"]})
+                retrieved_old_memory.append({"id": mem.id, "text": mem.payload.get("data", "")})
 
         unique_data = {}
         for item in retrieved_old_memory:
@@ -518,7 +518,7 @@ def get(self, memory_id):
 
         result_item = MemoryItem(
             id=memory.id,
-            memory=memory.payload["data"],
+            memory=memory.payload.get("data", ""),
             hash=memory.payload.get("hash"),
             created_at=memory.payload.get("created_at"),
             updated_at=memory.payload.get("updated_at"),
@@ -623,7 +623,7 @@ def _get_all_from_vector_store(self, filters, limit):
         for mem in actual_memories:
             memory_item_dict = MemoryItem(
                 id=mem.id,
-                memory=mem.payload["data"],
+                memory=mem.payload.get("data", ""),
                 hash=mem.payload.get("hash"),
                 created_at=mem.payload.get("created_at"),
                 updated_at=mem.payload.get("updated_at"),
@@ -735,7 +735,7 @@ def _search_vector_store(self, query, filters, limit, threshold: Optional[float]
         for mem in memories:
             memory_item_dict = MemoryItem(
                 id=mem.id,
-                memory=mem.payload["data"],
+                memory=mem.payload.get("data", ""),
                 hash=mem.payload.get("hash"),
                 created_at=mem.payload.get("created_at"),
                 updated_at=mem.payload.get("updated_at"),
@@ -962,7 +962,7 @@ def _update_memory(self, memory_id, data, existing_embeddings, metadata=None):
     def _delete_memory(self, memory_id):
         logger.info(f"Deleting memory with {memory_id=}")
         existing_memory = self.vector_store.get(vector_id=memory_id)
-        prev_value = existing_memory.payload["data"]
+        prev_value = existing_memory.payload.get("data", "")
         self.vector_store.delete(vector_id=memory_id)
         self.db.add_history(
             memory_id,
@@ -1234,7 +1234,7 @@ async def process_fact_for_search(new_mem_content):
                 limit=5,
                 filters=effective_filters,  # 'filters' is query_filters_for_inference
             )
-            return [{"id": mem.id, "text": mem.payload["data"]} for mem in existing_mems]
+            return [{"id": mem.id, "text": mem.payload.get("data", "")} for mem in existing_mems]
 
         search_tasks = [process_fact_for_search(fact) for fact in new_retrieved_facts]
         search_results_list = await asyncio.gather(*search_tasks)
@@ -1382,7 +1382,7 @@ async def get(self, memory_id):
 
         result_item = MemoryItem(
             id=memory.id,
-            memory=memory.payload["data"],
+            memory=memory.payload.get("data", ""),
             hash=memory.payload.get("hash"),
             created_at=memory.payload.get("created_at"),
             updated_at=memory.payload.get("updated_at"),
@@ -1492,7 +1492,7 @@ async def _get_all_from_vector_store(self, filters, limit):
         for mem in actual_memories:
             memory_item_dict = MemoryItem(
                 id=mem.id,
-                memory=mem.payload["data"],
+                memory=mem.payload.get("data", ""),
                 hash=mem.payload.get("hash"),
                 created_at=mem.payload.get("created_at"),
                 updated_at=mem.payload.get("updated_at"),
@@ -1609,7 +1609,7 @@ async def _search_vector_store(self, query, filters, limit, threshold: Optional[
         for mem in memories:
             memory_item_dict = MemoryItem(
                 id=mem.id,
-                memory=mem.payload["data"],
+                memory=mem.payload.get("data", ""),
                 hash=mem.payload.get("hash"),
                 created_at=mem.payload.get("created_at"),
                 updated_at=mem.payload.get("updated_at"),
@@ -1860,7 +1860,7 @@ async def _update_memory(self, memory_id, data, existing_embeddings, metadata=No
     async def _delete_memory(self, memory_id):
         logger.info(f"Deleting memory with {memory_id=}")
         existing_memory = await asyncio.to_thread(self.vector_store.get, vector_id=memory_id)
-        prev_value = existing_memory.payload["data"]
+        prev_value = existing_memory.payload.get("data", "")
 
         await asyncio.to_thread(self.vector_store.delete, vector_id=memory_id)
         await asyncio.to_thread(

tests/test_memory.py

@@ -6,6 +6,15 @@
 from mem0.configs.base import MemoryConfig
 
 
+class MockVectorMemory:
+    """Mock memory object for testing incomplete payloads."""
+    
+    def __init__(self, memory_id: str, payload: dict, score: float = 0.8):
+        self.id = memory_id
+        self.payload = payload
+        self.score = score
+
+
 @pytest.fixture
 def memory_client():
     with patch.object(Memory, "__init__", return_value=None):
@@ -95,3 +104,40 @@ def test_collection_name_preserved_after_reset(mock_sqlite, mock_llm_factory, mo
     if reset_calls:
         reset_config = reset_calls[-1][0][1]  
         assert reset_config.collection_name == test_collection_name, f"Reset used wrong collection name: {reset_config.collection_name}"
+
+
+@patch('mem0.utils.factory.EmbedderFactory.create')
+@patch('mem0.utils.factory.VectorStoreFactory.create')
+@patch('mem0.utils.factory.LlmFactory.create')
+@patch('mem0.memory.storage.SQLiteManager')
+def test_search_handles_incomplete_payloads(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
+    """Test that search operations handle memory objects with missing 'data' key gracefully."""
+    mock_embedder_factory.return_value = MagicMock()
+    mock_vector_store = MagicMock()
+    mock_vector_factory.return_value = mock_vector_store
+    mock_llm_factory.return_value = MagicMock()
+    mock_sqlite.return_value = MagicMock()
+
+    from mem0.memory.main import Memory as MemoryClass
+    config = MemoryConfig()
+    memory = MemoryClass(config)
+
+    # Create test data with both complete and incomplete payloads
+    incomplete_memory = MockVectorMemory("mem_1", {"hash": "abc123"})
+    complete_memory = MockVectorMemory("mem_2", {"data": "content", "hash": "def456"})
+
+    mock_vector_store.search.return_value = [incomplete_memory, complete_memory]
+    
+    mock_embedder = MagicMock()
+    mock_embedder.embed.return_value = [0.1, 0.2, 0.3]
+    memory.embedding_model = mock_embedder
+
+    # This should not raise KeyError even with incomplete payloads
+    result = memory._search_vector_store("test", {"user_id": "test"}, 10)
+    
+    assert len(result) == 2
+    memories_by_id = {mem["id"]: mem for mem in result}
+    
+    # Verify defensive programming works correctly
+    assert memories_by_id["mem_1"]["memory"] == ""  # Missing data gets empty string
+    assert memories_by_id["mem_2"]["memory"] == "content"  # Normal data preserved