prepare_reference_data: handle dict entries with missing or non-string id

src/memos/mem_os/utils/reference_utils.py

@@ -152,11 +152,16 @@ def prepare_reference_data(memories_list: list[TextualMemoryItem]) -> list[dict]
             reference.append({"metadata": memories_json["metadata"]})
         else:
             memories_json = memories
-            memories_json["metadata"]["ref_id"] = f"{memories_json['id'].split('-')[0]}"
+            # already-serialized entries may lack keys or carry a non-string
+            # id (e.g. an int from a JSON payload); derive the ref_id from
+            # the string form when an id exists at all
+            memory_id = memories_json.get("id")
+            if memory_id is not None:
+                memories_json["metadata"]["ref_id"] = f"{str(memory_id).split('-')[0]}"
+                memories_json["metadata"]["id"] = memory_id
             memories_json["metadata"]["embedding"] = []
             memories_json["metadata"]["sources"] = []
             memories_json["metadata"]["memory"] = memories_json["memory"]
-            memories_json["metadata"]["id"] = memories_json["id"]
             reference.append({"metadata": memories_json["metadata"]})
 
     return reference

tests/mem_os/test_prepare_reference_data.py

@@ -0,0 +1,35 @@
+from memos.mem_os.utils.reference_utils import prepare_reference_data
+
+
+def test_textual_item_still_gets_ref_id():
+    from memos.memories.textual.item import TextualMemoryItem, TextualMemoryMetadata
+
+    item = TextualMemoryItem(
+        memory="hello world",
+        metadata=TextualMemoryMetadata(user_id="u1"),
+    )
+    (reference,) = prepare_reference_data([item])
+    assert reference["metadata"]["ref_id"] == str(item.id).split("-")[0]
+    assert reference["metadata"]["memory"] == "hello world"
+
+
+def test_dict_entry_with_string_id():
+    (reference,) = prepare_reference_data(
+        [{"id": "92ff35fb-1234", "memory": "m", "metadata": {}}]
+    )
+    assert reference["metadata"]["ref_id"] == "92ff35fb"
+    assert reference["metadata"]["id"] == "92ff35fb-1234"
+
+
+def test_dict_entry_with_int_id():
+    # an int id used to crash with AttributeError: 'int' object has no
+    # attribute 'split'
+    (reference,) = prepare_reference_data([{"id": 12345, "memory": "m", "metadata": {}}])
+    assert reference["metadata"]["ref_id"] == "12345"
+
+
+def test_dict_entry_without_id():
+    # a missing id used to crash with KeyError: 'id'
+    (reference,) = prepare_reference_data([{"memory": "m", "metadata": {}}])
+    assert reference["metadata"]["memory"] == "m"
+    assert "ref_id" not in reference["metadata"]