fix(memory): paginate delete_all across vector store pages (#6636)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
co-authored by
kartik-mem0
parent
29fa41558c
commit
54328ffd97
+51
-17
@@ -133,6 +133,7 @@ _SENSITIVE_SUFFIXES = (
|
|||||||
|
|
||||||
# Entity parameters that must be passed via filters, not top-level kwargs
|
# Entity parameters that must be passed via filters, not top-level kwargs
|
||||||
ENTITY_PARAMS = frozenset({"user_id", "agent_id", "run_id"})
|
ENTITY_PARAMS = frozenset({"user_id", "agent_id", "run_id"})
|
||||||
|
DELETE_ALL_BATCH_SIZE = 1000
|
||||||
|
|
||||||
# Tenant-scoping fields that update() must never let caller-supplied metadata overwrite (issues #4490, #6277).
|
# Tenant-scoping fields that update() must never let caller-supplied metadata overwrite (issues #4490, #6277).
|
||||||
_IDENTITY_KEYS = ENTITY_PARAMS | {"actor_id"}
|
_IDENTITY_KEYS = ENTITY_PARAMS | {"actor_id"}
|
||||||
@@ -1885,14 +1886,28 @@ class Memory(MemoryBase):
|
|||||||
|
|
||||||
keys, encoded_ids = process_telemetry_filters(filters)
|
keys, encoded_ids = process_telemetry_filters(filters)
|
||||||
capture_event("mem0.delete_all", self, {"keys": keys, "encoded_ids": encoded_ids, "sync_type": "sync"})
|
capture_event("mem0.delete_all", self, {"keys": keys, "encoded_ids": encoded_ids, "sync_type": "sync"})
|
||||||
# delete all vector memories and reset the collections
|
# Keep listing after each batch is deleted. Most vector stores cap
|
||||||
memories = self.vector_store.list(filters=filters)[0]
|
# list() at 100 results by default, which silently truncates deletes.
|
||||||
for memory in memories:
|
deleted_count = 0
|
||||||
self._delete_memory(memory.id)
|
seen_batches = set()
|
||||||
|
while True:
|
||||||
|
memories = self.vector_store.list(
|
||||||
|
filters=filters, top_k=DELETE_ALL_BATCH_SIZE
|
||||||
|
)[0]
|
||||||
|
if not memories:
|
||||||
|
break
|
||||||
|
batch_ids = tuple(sorted(str(memory.id) for memory in memories))
|
||||||
|
if batch_ids in seen_batches:
|
||||||
|
logger.warning("Stopping delete_all after a repeated memory batch")
|
||||||
|
break
|
||||||
|
seen_batches.add(batch_ids)
|
||||||
|
for memory in memories:
|
||||||
|
self._delete_memory(memory.id)
|
||||||
|
deleted_count += len(memories)
|
||||||
|
|
||||||
logger.info(f"Deleted {len(memories)} memories")
|
logger.info(f"Deleted {deleted_count} memories")
|
||||||
|
|
||||||
decay_usage_notice = detect_decay_usage_from_delete_all(len(memories))
|
decay_usage_notice = detect_decay_usage_from_delete_all(deleted_count)
|
||||||
if decay_usage_notice:
|
if decay_usage_notice:
|
||||||
display_decay_usage_notice(self, "sync", "delete_all", *decay_usage_notice)
|
display_decay_usage_notice(self, "sync", "delete_all", *decay_usage_notice)
|
||||||
else:
|
else:
|
||||||
@@ -3515,26 +3530,45 @@ class AsyncMemory(MemoryBase):
|
|||||||
|
|
||||||
keys, encoded_ids = process_telemetry_filters(filters)
|
keys, encoded_ids = process_telemetry_filters(filters)
|
||||||
capture_event("mem0.delete_all", self, {"keys": keys, "encoded_ids": encoded_ids, "sync_type": "async"})
|
capture_event("mem0.delete_all", self, {"keys": keys, "encoded_ids": encoded_ids, "sync_type": "async"})
|
||||||
memories = await asyncio.to_thread(self.vector_store.list, filters=filters)
|
deleted_count = 0
|
||||||
|
errors = []
|
||||||
delete_tasks = []
|
seen_batches = set()
|
||||||
for memory in memories[0]:
|
while True:
|
||||||
delete_tasks.append(self._delete_memory(memory.id, skip_entity_cleanup=True))
|
memories = await asyncio.to_thread(
|
||||||
|
self.vector_store.list,
|
||||||
results = await asyncio.gather(*delete_tasks, return_exceptions=True)
|
filters=filters,
|
||||||
|
top_k=DELETE_ALL_BATCH_SIZE,
|
||||||
|
)
|
||||||
|
batch = memories[0] if memories else []
|
||||||
|
if not batch:
|
||||||
|
break
|
||||||
|
batch_ids = tuple(sorted(str(memory.id) for memory in batch))
|
||||||
|
if batch_ids in seen_batches:
|
||||||
|
logger.warning("Stopping delete_all after a repeated memory batch")
|
||||||
|
break
|
||||||
|
seen_batches.add(batch_ids)
|
||||||
|
delete_tasks = [
|
||||||
|
self._delete_memory(memory.id, skip_entity_cleanup=True)
|
||||||
|
for memory in batch
|
||||||
|
]
|
||||||
|
results = await asyncio.gather(*delete_tasks, return_exceptions=True)
|
||||||
|
batch_errors = [
|
||||||
|
result for result in results if isinstance(result, BaseException)
|
||||||
|
]
|
||||||
|
errors.extend(batch_errors)
|
||||||
|
deleted_count += len(results) - len(batch_errors)
|
||||||
|
|
||||||
if self._entity_store is not None:
|
if self._entity_store is not None:
|
||||||
await self._bulk_clear_entity_store(filters)
|
await self._bulk_clear_entity_store(filters)
|
||||||
|
|
||||||
errors = [r for r in results if isinstance(r, BaseException)]
|
|
||||||
if errors:
|
if errors:
|
||||||
logger.warning("Failed to delete %d out of %d memories", len(errors), len(results))
|
logger.warning("Failed to delete %d memories", len(errors))
|
||||||
for err in errors:
|
for err in errors:
|
||||||
logger.warning("Delete error: %s", err)
|
logger.warning("Delete error: %s", err)
|
||||||
|
|
||||||
logger.info(f"Deleted {len(results) - len(errors)} memories")
|
logger.info(f"Deleted {deleted_count} memories")
|
||||||
|
|
||||||
decay_usage_notice = detect_decay_usage_from_delete_all(len(memories[0]))
|
decay_usage_notice = detect_decay_usage_from_delete_all(deleted_count)
|
||||||
if decay_usage_notice:
|
if decay_usage_notice:
|
||||||
await display_decay_usage_notice_async(self, "async", "delete_all", *decay_usage_notice)
|
await display_decay_usage_notice_async(self, "async", "delete_all", *decay_usage_notice)
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -91,7 +91,7 @@ def test_sync_delete_failure_does_not_trigger_decay_usage_notice(monkeypatch):
|
|||||||
def test_sync_delete_all_decay_usage_runs_after_success(monkeypatch):
|
def test_sync_delete_all_decay_usage_runs_after_success(monkeypatch):
|
||||||
memory = make_sync_memory()
|
memory = make_sync_memory()
|
||||||
memories = [SimpleNamespace(id="memory-1"), SimpleNamespace(id="memory-2")]
|
memories = [SimpleNamespace(id="memory-1"), SimpleNamespace(id="memory-2")]
|
||||||
memory.vector_store.list.return_value = (memories, None)
|
memory.vector_store.list.side_effect = [(memories, None), ([], None)]
|
||||||
decay_notice = MagicMock()
|
decay_notice = MagicMock()
|
||||||
first_run_notice = MagicMock()
|
first_run_notice = MagicMock()
|
||||||
detect_decay = MagicMock(return_value=("delete_all", "bulk_delete", None, 2))
|
detect_decay = MagicMock(return_value=("delete_all", "bulk_delete", None, 2))
|
||||||
@@ -117,6 +117,21 @@ def test_sync_delete_all_decay_usage_runs_after_success(monkeypatch):
|
|||||||
first_run_notice.assert_not_called()
|
first_run_notice.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_sync_delete_all_stops_when_vector_store_repeats_batch(monkeypatch):
|
||||||
|
memory = make_sync_memory()
|
||||||
|
memories = [SimpleNamespace(id="memory-1"), SimpleNamespace(id="memory-2")]
|
||||||
|
memory.vector_store.list.side_effect = [(memories, None), (memories, None)]
|
||||||
|
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||||
|
monkeypatch.setattr(memory_main, "detect_decay_usage_from_delete_all", MagicMock(return_value=None))
|
||||||
|
monkeypatch.setattr(memory_main, "display_first_run_notice", MagicMock())
|
||||||
|
|
||||||
|
result = Memory.delete_all(memory, user_id="u1")
|
||||||
|
|
||||||
|
assert result == {"message": "Memories deleted successfully!"}
|
||||||
|
assert memory.vector_store.list.call_count == 2
|
||||||
|
assert memory._delete_memory.call_count == 2
|
||||||
|
|
||||||
|
|
||||||
def test_sync_delete_all_zero_deletes_uses_first_run_notice(monkeypatch):
|
def test_sync_delete_all_zero_deletes_uses_first_run_notice(monkeypatch):
|
||||||
memory = make_sync_memory()
|
memory = make_sync_memory()
|
||||||
memory.vector_store.list.return_value = ([], None)
|
memory.vector_store.list.return_value = ([], None)
|
||||||
@@ -189,7 +204,7 @@ async def test_async_delete_failure_does_not_trigger_decay_usage_notice(monkeypa
|
|||||||
async def test_async_delete_all_decay_usage_runs_after_success(monkeypatch):
|
async def test_async_delete_all_decay_usage_runs_after_success(monkeypatch):
|
||||||
memory = make_async_memory()
|
memory = make_async_memory()
|
||||||
memories = [SimpleNamespace(id="memory-1"), SimpleNamespace(id="memory-2")]
|
memories = [SimpleNamespace(id="memory-1"), SimpleNamespace(id="memory-2")]
|
||||||
memory.vector_store.list.return_value = (memories, None)
|
memory.vector_store.list.side_effect = [(memories, None), ([], None)]
|
||||||
decay_notice = AsyncMock()
|
decay_notice = AsyncMock()
|
||||||
first_run_notice = AsyncMock()
|
first_run_notice = AsyncMock()
|
||||||
detect_decay = MagicMock(return_value=("delete_all", "bulk_delete", None, 2))
|
detect_decay = MagicMock(return_value=("delete_all", "bulk_delete", None, 2))
|
||||||
@@ -213,3 +228,23 @@ async def test_async_delete_all_decay_usage_runs_after_success(monkeypatch):
|
|||||||
2,
|
2,
|
||||||
)
|
)
|
||||||
first_run_notice.assert_not_awaited()
|
first_run_notice.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_delete_all_stops_when_vector_store_repeats_batch(monkeypatch):
|
||||||
|
memory = make_async_memory()
|
||||||
|
memories = [SimpleNamespace(id="memory-1"), SimpleNamespace(id="memory-2")]
|
||||||
|
memory.vector_store.list.side_effect = [
|
||||||
|
(memories, None),
|
||||||
|
(memories, None),
|
||||||
|
RuntimeError("delete_all should have stopped before a third list() call"),
|
||||||
|
]
|
||||||
|
monkeypatch.setattr(memory_main, "capture_event", MagicMock())
|
||||||
|
monkeypatch.setattr(memory_main, "detect_decay_usage_from_delete_all", MagicMock(return_value=None))
|
||||||
|
monkeypatch.setattr(memory_main, "display_first_run_notice_async", AsyncMock())
|
||||||
|
|
||||||
|
result = await AsyncMemory.delete_all(memory, user_id="u1")
|
||||||
|
|
||||||
|
assert result == {"message": "Memories deleted successfully!"}
|
||||||
|
assert memory.vector_store.list.call_count == 2
|
||||||
|
assert memory._delete_memory.await_count == 2
|
||||||
|
|||||||
+24
-3
@@ -280,19 +280,36 @@ def test_delete(memory_instance):
|
|||||||
|
|
||||||
def test_delete_all(memory_instance):
|
def test_delete_all(memory_instance):
|
||||||
mock_memories = [Mock(id="1"), Mock(id="2")]
|
mock_memories = [Mock(id="1"), Mock(id="2")]
|
||||||
memory_instance.vector_store.list = Mock(return_value=(mock_memories, None))
|
memory_instance.vector_store.list = Mock(side_effect=[(mock_memories, None), ([], None)])
|
||||||
memory_instance.vector_store.reset = Mock()
|
memory_instance.vector_store.reset = Mock()
|
||||||
memory_instance._delete_memory = Mock()
|
memory_instance._delete_memory = Mock()
|
||||||
|
|
||||||
result = memory_instance.delete_all(user_id="test_user")
|
result = memory_instance.delete_all(user_id="test_user")
|
||||||
|
|
||||||
assert memory_instance._delete_memory.call_count == 2
|
assert memory_instance._delete_memory.call_count == 2
|
||||||
|
memory_instance.vector_store.list.assert_called_with(
|
||||||
|
filters={"user_id": "test_user"}, top_k=1000
|
||||||
|
)
|
||||||
# Ensure the collection is NOT dropped — only matched memories should be removed
|
# Ensure the collection is NOT dropped — only matched memories should be removed
|
||||||
memory_instance.vector_store.reset.assert_not_called()
|
memory_instance.vector_store.reset.assert_not_called()
|
||||||
|
|
||||||
assert result["message"] == "Memories deleted successfully!"
|
assert result["message"] == "Memories deleted successfully!"
|
||||||
|
|
||||||
|
|
||||||
|
def test_delete_all_paginates_beyond_vector_store_page_size(memory_instance):
|
||||||
|
first_batch = [Mock(id=str(index)) for index in range(1000)]
|
||||||
|
second_batch = [Mock(id="1000")]
|
||||||
|
memory_instance.vector_store.list = Mock(
|
||||||
|
side_effect=[(first_batch, None), (second_batch, None), ([], None)]
|
||||||
|
)
|
||||||
|
memory_instance._delete_memory = Mock()
|
||||||
|
|
||||||
|
memory_instance.delete_all(user_id="test_user")
|
||||||
|
|
||||||
|
assert memory_instance._delete_memory.call_count == 1001
|
||||||
|
assert memory_instance.vector_store.list.call_count == 3
|
||||||
|
|
||||||
|
|
||||||
def test_get_all(memory_instance):
|
def test_get_all(memory_instance):
|
||||||
mock_memories = [Mock(id="1", payload={"data": "Memory 1", "user_id": "test_user"})]
|
mock_memories = [Mock(id="1", payload={"data": "Memory 1", "user_id": "test_user"})]
|
||||||
memory_instance.vector_store.list = Mock(return_value=(mock_memories, None))
|
memory_instance.vector_store.list = Mock(return_value=(mock_memories, None))
|
||||||
@@ -431,7 +448,9 @@ class TestEntityIdValidation:
|
|||||||
|
|
||||||
memory_instance.delete_all(user_id=" alice ")
|
memory_instance.delete_all(user_id=" alice ")
|
||||||
|
|
||||||
memory_instance.vector_store.list.assert_called_once_with(filters={"user_id": "alice"})
|
memory_instance.vector_store.list.assert_called_once_with(
|
||||||
|
filters={"user_id": "alice"}, top_k=1000
|
||||||
|
)
|
||||||
|
|
||||||
def test_validate_coerces_non_string_entity_id(self):
|
def test_validate_coerces_non_string_entity_id(self):
|
||||||
"""Integer (and other non-string) ids are coerced to str, not crashed on."""
|
"""Integer (and other non-string) ids are coerced to str, not crashed on."""
|
||||||
@@ -444,7 +463,9 @@ class TestEntityIdValidation:
|
|||||||
|
|
||||||
memory_instance.delete_all(user_id=42)
|
memory_instance.delete_all(user_id=42)
|
||||||
|
|
||||||
memory_instance.vector_store.list.assert_called_once_with(filters={"user_id": "42"})
|
memory_instance.vector_store.list.assert_called_once_with(
|
||||||
|
filters={"user_id": "42"}, top_k=1000
|
||||||
|
)
|
||||||
|
|
||||||
def test_get_all_coerces_integer_user_id(self, memory_instance):
|
def test_get_all_coerces_integer_user_id(self, memory_instance):
|
||||||
"""get_all should accept an integer user_id in filters and scope by its str form."""
|
"""get_all should accept an integer user_id in filters and scope by its str form."""
|
||||||
|
|||||||
@@ -983,7 +983,7 @@ async def test_async_delete_all_continues_on_partial_failure(mock_sqlite, mock_l
|
|||||||
mem3.id = "mem-3"
|
mem3.id = "mem-3"
|
||||||
mem3.payload = {"data": "three", "created_at": "2024-01-01T00:00:00+00:00", "actor_id": None, "role": None}
|
mem3.payload = {"data": "three", "created_at": "2024-01-01T00:00:00+00:00", "actor_id": None, "role": None}
|
||||||
|
|
||||||
mock_vector_store.list.return_value = ([mem1, mem2, mem3],)
|
mock_vector_store.list.side_effect = [([mem1, mem2, mem3],), ([],)]
|
||||||
|
|
||||||
def _get_side_effect(vector_id):
|
def _get_side_effect(vector_id):
|
||||||
if vector_id == "mem-2":
|
if vector_id == "mem-2":
|
||||||
@@ -999,6 +999,10 @@ async def test_async_delete_all_continues_on_partial_failure(mock_sqlite, mock_l
|
|||||||
|
|
||||||
assert result == {"message": "Memories deleted successfully!"}
|
assert result == {"message": "Memories deleted successfully!"}
|
||||||
assert mock_vector_store.delete.call_count == 2
|
assert mock_vector_store.delete.call_count == 2
|
||||||
|
assert mock_vector_store.list.call_args_list[0].kwargs == {
|
||||||
|
"filters": {"user_id": "test-user"},
|
||||||
|
"top_k": 1000,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
@patch('mem0.utils.factory.EmbedderFactory.create')
|
||||||
@@ -1473,7 +1477,7 @@ class TestAsyncDeleteAllEntityRace:
|
|||||||
mem_b = MagicMock()
|
mem_b = MagicMock()
|
||||||
mem_b.id = "mem-b"
|
mem_b.id = "mem-b"
|
||||||
mem_b.payload = {"data": "Alice works at Acme", "user_id": "alice"}
|
mem_b.payload = {"data": "Alice works at Acme", "user_id": "alice"}
|
||||||
mock_vector_store.list.return_value = ([mem_a, mem_b],)
|
mock_vector_store.list.side_effect = [([mem_a, mem_b],), ([],)]
|
||||||
mock_vector_store.get.side_effect = lambda vector_id: {"mem-a": mem_a, "mem-b": mem_b}[vector_id]
|
mock_vector_store.get.side_effect = lambda vector_id: {"mem-a": mem_a, "mem-b": mem_b}[vector_id]
|
||||||
mock_vector_factory.return_value = mock_vector_store
|
mock_vector_factory.return_value = mock_vector_store
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user