From 58cc44ff38b0637b203bf8369c50597bad00001f Mon Sep 17 00:00:00 2001 From: Vishaal LS <62366204+lsvishaal@users.noreply.github.com> Date: Wed, 1 Oct 2025 03:40:43 +0530 Subject: [PATCH] fix: handle missing 'data' key in memory payload during search operations (#3524) --- mem0/memory/main.py | 20 +++++++++---------- tests/test_memory.py | 46 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 10 deletions(-) diff --git a/mem0/memory/main.py b/mem0/memory/main.py index f4b2f6a3c..efcbc52d6 100644 --- a/mem0/memory/main.py +++ b/mem0/memory/main.py @@ -382,7 +382,7 @@ class Memory(MemoryBase): 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 @@ class Memory(MemoryBase): 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 @@ class Memory(MemoryBase): 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 @@ class Memory(MemoryBase): 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 @@ class Memory(MemoryBase): 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 @@ class AsyncMemory(MemoryBase): 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 @@ class AsyncMemory(MemoryBase): 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 @@ class AsyncMemory(MemoryBase): 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 @@ class AsyncMemory(MemoryBase): 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 @@ class AsyncMemory(MemoryBase): 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( diff --git a/tests/test_memory.py b/tests/test_memory.py index afd36e0fd..b43e5e815 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -6,6 +6,15 @@ from mem0 import Memory 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