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_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).
|
||||
_IDENTITY_KEYS = ENTITY_PARAMS | {"actor_id"}
|
||||
@@ -1885,14 +1886,28 @@ class Memory(MemoryBase):
|
||||
|
||||
keys, encoded_ids = process_telemetry_filters(filters)
|
||||
capture_event("mem0.delete_all", self, {"keys": keys, "encoded_ids": encoded_ids, "sync_type": "sync"})
|
||||
# delete all vector memories and reset the collections
|
||||
memories = self.vector_store.list(filters=filters)[0]
|
||||
for memory in memories:
|
||||
self._delete_memory(memory.id)
|
||||
# Keep listing after each batch is deleted. Most vector stores cap
|
||||
# list() at 100 results by default, which silently truncates deletes.
|
||||
deleted_count = 0
|
||||
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:
|
||||
display_decay_usage_notice(self, "sync", "delete_all", *decay_usage_notice)
|
||||
else:
|
||||
@@ -3515,26 +3530,45 @@ class AsyncMemory(MemoryBase):
|
||||
|
||||
keys, encoded_ids = process_telemetry_filters(filters)
|
||||
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)
|
||||
|
||||
delete_tasks = []
|
||||
for memory in memories[0]:
|
||||
delete_tasks.append(self._delete_memory(memory.id, skip_entity_cleanup=True))
|
||||
|
||||
results = await asyncio.gather(*delete_tasks, return_exceptions=True)
|
||||
deleted_count = 0
|
||||
errors = []
|
||||
seen_batches = set()
|
||||
while True:
|
||||
memories = await asyncio.to_thread(
|
||||
self.vector_store.list,
|
||||
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:
|
||||
await self._bulk_clear_entity_store(filters)
|
||||
|
||||
errors = [r for r in results if isinstance(r, BaseException)]
|
||||
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:
|
||||
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:
|
||||
await display_decay_usage_notice_async(self, "async", "delete_all", *decay_usage_notice)
|
||||
else:
|
||||
|
||||
Reference in New Issue
Block a user