diff --git a/src/memos/mem_os/utils/reference_utils.py b/src/memos/mem_os/utils/reference_utils.py index 09b812207..7768b8aef 100644 --- a/src/memos/mem_os/utils/reference_utils.py +++ b/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 diff --git a/tests/mem_os/test_prepare_reference_data.py b/tests/mem_os/test_prepare_reference_data.py new file mode 100644 index 000000000..866fbc4ae --- /dev/null +++ b/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"]