diff --git a/mem0/memory/main.py b/mem0/memory/main.py index 1804e7831..582e56d88 100644 --- a/mem0/memory/main.py +++ b/mem0/memory/main.py @@ -63,6 +63,19 @@ def _normalize_iso_timestamp_to_utc(timestamp: Optional[str]) -> Optional[str]: return parsed.astimezone(timezone.utc).isoformat() +def _resolve_mapped_id(temp_uuid_mapping, resp, event_type): + """Resolve a temp integer ID from the LLM response to a real UUID. + + Returns the UUID if found, or None (with a warning log) if the LLM + hallucinated an ID that doesn't exist in the mapping. + """ + raw_id = resp.get("id") + memory_id = temp_uuid_mapping.get(raw_id) + if memory_id is None: + logger.warning(f"{event_type} skipped: LLM returned unknown id {raw_id!r}") + return memory_id + + # Fields that hold runtime auth/connection objects and must be preserved. # These are non-serializable objects (e.g. AWSV4SignerAuth, RequestsHttpConnection) # needed by clients like OpenSearch — not sensitive strings to redact. @@ -639,28 +652,34 @@ class Memory(MemoryBase): ) returned_memories.append({"id": memory_id, "memory": action_text, "event": event_type}) elif event_type == "UPDATE": + memory_id = _resolve_mapped_id(temp_uuid_mapping, resp, "UPDATE") + if memory_id is None: + continue # Ensure action_text has an embedding cached to avoid redundant API calls if action_text not in new_message_embeddings: new_message_embeddings[action_text] = self.embedding_model.embed(action_text, "update") self._update_memory( - memory_id=temp_uuid_mapping[resp.get("id")], + memory_id=memory_id, data=action_text, existing_embeddings=new_message_embeddings, metadata=deepcopy(metadata), ) returned_memories.append( { - "id": temp_uuid_mapping[resp.get("id")], + "id": memory_id, "memory": action_text, "event": event_type, "previous_memory": resp.get("old_memory"), } ) elif event_type == "DELETE": - self._delete_memory(memory_id=temp_uuid_mapping[resp.get("id")]) + memory_id = _resolve_mapped_id(temp_uuid_mapping, resp, "DELETE") + if memory_id is None: + continue + self._delete_memory(memory_id=memory_id) returned_memories.append( { - "id": temp_uuid_mapping[resp.get("id")], + "id": memory_id, "memory": action_text, "event": event_type, } @@ -1755,6 +1774,9 @@ class AsyncMemory(MemoryBase): ) memory_tasks.append((task, resp, "ADD", None)) elif event_type == "UPDATE": + memory_id = _resolve_mapped_id(temp_uuid_mapping, resp, "UPDATE") + if memory_id is None: + continue # Ensure action_text has an embedding cached to avoid redundant API calls if action_text not in new_message_embeddings: new_message_embeddings[action_text] = await asyncio.to_thread( @@ -1762,16 +1784,19 @@ class AsyncMemory(MemoryBase): ) task = asyncio.create_task( self._update_memory( - memory_id=temp_uuid_mapping[resp["id"]], + memory_id=memory_id, data=action_text, existing_embeddings=new_message_embeddings, metadata=deepcopy(metadata), ) ) - memory_tasks.append((task, resp, "UPDATE", temp_uuid_mapping[resp["id"]])) + memory_tasks.append((task, resp, "UPDATE", memory_id)) elif event_type == "DELETE": - task = asyncio.create_task(self._delete_memory(memory_id=temp_uuid_mapping[resp.get("id")])) - memory_tasks.append((task, resp, "DELETE", temp_uuid_mapping[resp.get("id")])) + memory_id = _resolve_mapped_id(temp_uuid_mapping, resp, "DELETE") + if memory_id is None: + continue + task = asyncio.create_task(self._delete_memory(memory_id=memory_id)) + memory_tasks.append((task, resp, "DELETE", memory_id)) elif event_type == "NONE": # Even if content doesn't need updating, update session IDs if provided memory_id = temp_uuid_mapping.get(resp.get("id")) diff --git a/tests/memory/test_main.py b/tests/memory/test_main.py index 9938baf38..15507e2be 100644 --- a/tests/memory/test_main.py +++ b/tests/memory/test_main.py @@ -627,6 +627,228 @@ async def test_async_update_preserves_actor_id_when_different_actor_updates(mock assert stored["actor_id"] == "Alice" +class TestHallucinatedIdGuard: + """Tests for temp_uuid_mapping guard against LLM-hallucinated IDs (issue #3931). + + When the LLM returns an UPDATE or DELETE with an ID that doesn't exist in + temp_uuid_mapping, the code should skip gracefully instead of raising KeyError. + """ + + def test_sync_update_with_hallucinated_id_skips_gracefully(self, mocker, caplog): + """Sync UPDATE with an out-of-range ID should be skipped with a warning.""" + memory = _build_memory_instance(mocker, Memory) + memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3] + mocker.patch("mem0.memory.main.capture_event") + + # Simulate: 2 existing memories (IDs "0" and "1"), but LLM returns UPDATE for ID "12" + existing_mem = MagicMock() + existing_mem.id = "uuid-aaa" + existing_mem.payload = {"data": "User likes coffee"} + memory.vector_store.search.return_value = [existing_mem] + + # First LLM call: fact extraction → returns one fact + # Second LLM call: memory update actions → returns UPDATE with hallucinated ID "12" + memory.llm.generate_response.side_effect = [ + '{"facts": ["User likes tea"]}', + '{"memory": [{"id": "12", "text": "User likes tea", "event": "UPDATE", "old_memory": "User likes coffee"}]}', + ] + + with caplog.at_level(logging.WARNING): + result = memory._add_to_vector_store( + messages=[{"role": "user", "content": "I like tea"}], + metadata={}, + filters={}, + infer=True, + ) + + # Should not crash, should return empty (the hallucinated UPDATE was skipped) + assert result == [] + assert "UPDATE skipped: LLM returned unknown id" in caplog.text + # _update_memory should NOT have been called + memory.vector_store.update.assert_not_called() + + def test_sync_delete_with_hallucinated_id_skips_gracefully(self, mocker, caplog): + """Sync DELETE with an out-of-range ID should be skipped with a warning.""" + memory = _build_memory_instance(mocker, Memory) + memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3] + mocker.patch("mem0.memory.main.capture_event") + + existing_mem = MagicMock() + existing_mem.id = "uuid-aaa" + existing_mem.payload = {"data": "User likes coffee"} + memory.vector_store.search.return_value = [existing_mem] + + memory.llm.generate_response.side_effect = [ + '{"facts": ["Remove coffee preference"]}', + '{"memory": [{"id": "9", "text": "User likes coffee", "event": "DELETE"}]}', + ] + + with caplog.at_level(logging.WARNING): + result = memory._add_to_vector_store( + messages=[{"role": "user", "content": "I no longer like coffee"}], + metadata={}, + filters={}, + infer=True, + ) + + assert result == [] + assert "DELETE skipped: LLM returned unknown id" in caplog.text + memory.vector_store.delete.assert_not_called() + + def test_sync_valid_id_still_processes_normally(self, mocker, caplog): + """A valid ID should still be processed — the guard must not block legitimate operations.""" + memory = _build_memory_instance(mocker, Memory) + memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3] + mocker.patch("mem0.memory.main.capture_event") + + existing_mem = MagicMock() + existing_mem.id = "uuid-aaa" + existing_mem.payload = {"data": "User likes coffee"} + memory.vector_store.search.return_value = [existing_mem] + memory.vector_store.get.return_value = MagicMock( + payload={"data": "User likes coffee", "created_at": "2026-01-01T00:00:00+00:00"} + ) + + # ID "0" is valid since there's exactly 1 existing memory + memory.llm.generate_response.side_effect = [ + '{"facts": ["User likes tea now"]}', + '{"memory": [{"id": "0", "text": "User likes tea now", "event": "UPDATE", "old_memory": "User likes coffee"}]}', + ] + + with caplog.at_level(logging.WARNING): + result = memory._add_to_vector_store( + messages=[{"role": "user", "content": "I like tea now"}], + metadata={}, + filters={}, + infer=True, + ) + + assert len(result) == 1 + assert result[0]["event"] == "UPDATE" + assert result[0]["memory"] == "User likes tea now" + assert result[0]["id"] == "uuid-aaa" + assert "skipped" not in caplog.text + + @pytest.mark.asyncio + async def test_async_update_with_hallucinated_id_skips_gracefully(self, mocker, caplog): + """Async UPDATE with an out-of-range ID should be skipped with a warning.""" + memory = _build_memory_instance(mocker, AsyncMemory) + memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3] + mocker.patch("mem0.memory.main.capture_event") + + existing_mem = MagicMock() + existing_mem.id = "uuid-bbb" + existing_mem.payload = {"data": "User works at Acme"} + memory.vector_store.search.return_value = [existing_mem] + + memory.llm.generate_response.side_effect = [ + '{"facts": ["User works at Globex"]}', + '{"memory": [{"id": "7", "text": "User works at Globex", "event": "UPDATE", "old_memory": "User works at Acme"}]}', + ] + + with caplog.at_level(logging.WARNING): + result = await memory._add_to_vector_store( + messages=[{"role": "user", "content": "I now work at Globex"}], + metadata={}, + effective_filters={}, + infer=True, + ) + + assert result == [] + assert "UPDATE skipped: LLM returned unknown id" in caplog.text + + @pytest.mark.asyncio + async def test_async_delete_with_hallucinated_id_skips_gracefully(self, mocker, caplog): + """Async DELETE with an out-of-range ID should be skipped with a warning.""" + memory = _build_memory_instance(mocker, AsyncMemory) + memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3] + mocker.patch("mem0.memory.main.capture_event") + + existing_mem = MagicMock() + existing_mem.id = "uuid-ccc" + existing_mem.payload = {"data": "User lives in SF"} + memory.vector_store.search.return_value = [existing_mem] + + memory.llm.generate_response.side_effect = [ + '{"facts": ["Remove SF reference"]}', + '{"memory": [{"id": "16", "text": "User lives in SF", "event": "DELETE"}]}', + ] + + with caplog.at_level(logging.WARNING): + result = await memory._add_to_vector_store( + messages=[{"role": "user", "content": "I moved away from SF"}], + metadata={}, + effective_filters={}, + infer=True, + ) + + assert result == [] + assert "DELETE skipped: LLM returned unknown id" in caplog.text + + def test_sync_update_with_missing_id_key_skips_gracefully(self, mocker, caplog): + """UPDATE where the LLM omits the 'id' field entirely should be skipped.""" + memory = _build_memory_instance(mocker, Memory) + memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3] + mocker.patch("mem0.memory.main.capture_event") + + existing_mem = MagicMock() + existing_mem.id = "uuid-aaa" + existing_mem.payload = {"data": "User likes coffee"} + memory.vector_store.search.return_value = [existing_mem] + + # LLM response has no "id" key at all + memory.llm.generate_response.side_effect = [ + '{"facts": ["User likes tea"]}', + '{"memory": [{"text": "User likes tea", "event": "UPDATE", "old_memory": "User likes coffee"}]}', + ] + + with caplog.at_level(logging.WARNING): + result = memory._add_to_vector_store( + messages=[{"role": "user", "content": "I like tea"}], + metadata={}, + filters={}, + infer=True, + ) + + assert result == [] + assert "UPDATE skipped: LLM returned unknown id" in caplog.text + + @pytest.mark.asyncio + async def test_async_valid_id_still_processes_normally(self, mocker, caplog): + """Async path: a valid ID should process normally — no false positives from the guard.""" + memory = _build_memory_instance(mocker, AsyncMemory) + memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3] + mocker.patch("mem0.memory.main.capture_event") + + existing_mem = MagicMock() + existing_mem.id = "uuid-bbb" + existing_mem.payload = {"data": "User works at Acme"} + memory.vector_store.search.return_value = [existing_mem] + memory.vector_store.get.return_value = MagicMock( + payload={"data": "User works at Acme", "created_at": "2026-01-01T00:00:00+00:00"} + ) + + # ID "0" is valid since there's exactly 1 existing memory + memory.llm.generate_response.side_effect = [ + '{"facts": ["User works at Globex now"]}', + '{"memory": [{"id": "0", "text": "User works at Globex now", "event": "UPDATE", "old_memory": "User works at Acme"}]}', + ] + + with caplog.at_level(logging.WARNING): + result = await memory._add_to_vector_store( + messages=[{"role": "user", "content": "I now work at Globex"}], + metadata={}, + effective_filters={}, + infer=True, + ) + + assert len(result) == 1 + assert result[0]["event"] == "UPDATE" + assert result[0]["memory"] == "User works at Globex now" + assert result[0]["id"] == "uuid-bbb" + assert "skipped" not in caplog.text + + def test_normalize_iso_timestamp_to_utc_preserves_naive_values(): assert _normalize_iso_timestamp_to_utc("2026-03-18T00:00:00") == "2026-03-18T00:00:00"