JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: guard temp_uuid_mapping lookups against LLM-hallucinated IDs (fixes #3931) (#4674)

Browse Source 
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
Ignazio De Santis
2026-04-08 21:55:52 +05:30
committed by GitHub
co-authored by kartik-mem0
parent 2434b9d550
commit 081eca6d8f
2 changed files with 255 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/memory/main.py
+33 -8
View File
@@ -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"))