diff --git a/src/memos/mem_os/utils/reference_utils.py b/src/memos/mem_os/utils/reference_utils.py index 09b812207..ba4873817 100644 --- a/src/memos/mem_os/utils/reference_utils.py +++ b/src/memos/mem_os/utils/reference_utils.py @@ -1,13 +1,21 @@ +import re + from memos.memories.textual.item import ( TextualMemoryItem, ) +_REFERENCE_SEPARATOR_RE = re.compile(r",\s*") + + def split_continuous_references(text: str) -> str: """ Split continuous reference tags into individual reference tags. - Converts patterns like [1:92ff35fb, 4:bfe6f044] to [1:92ff35fb] [4:bfe6f044] + Converts patterns like [1:92ff35fb, 4:bfe6f044] to [1:92ff35fb][4:bfe6f044]. + Every comma between the brackets is treated as a separator regardless of + the whitespace that follows it, so mixed styles inside one tag (e.g. + ``[1:a,2:b, 3:c]``) all split correctly. Only processes text if: 1. '[' appears exactly once @@ -41,10 +49,11 @@ def split_continuous_references(text: str) -> str: # Check if there's a comma between brackets if "," not in content_between_brackets: return text - text = text.replace(content_between_brackets, content_between_brackets.replace(", ", "][")) - text = text.replace(content_between_brackets, content_between_brackets.replace(",", "][")) - - return text + # Split on every comma (with optional trailing whitespace) in a single pass. + # Sequential str.replace calls fail when the tag mixes ", " and "," styles + # because the first pass mutates the substring the second pass looks for. + split_content = _REFERENCE_SEPARATOR_RE.sub("][", content_between_brackets) + return text.replace(content_between_brackets, split_content) def process_streaming_references_complete(text_buffer: str) -> tuple[str, str]: diff --git a/tests/mem_os/test_reference_utils.py b/tests/mem_os/test_reference_utils.py new file mode 100644 index 000000000..ced5a1224 --- /dev/null +++ b/tests/mem_os/test_reference_utils.py @@ -0,0 +1,72 @@ +""" +Test suite for src/memos/mem_os/utils/reference_utils.py + +Focus: split_continuous_references handles every comma-based +separator style (with or without trailing whitespace, mixed styles +inside a single tag). Related issue: #2417 +""" + +import pytest + +from memos.mem_os.utils.reference_utils import split_continuous_references + + +class TestSplitContinuousReferences: + """Behavioral tests for split_continuous_references.""" + + def test_splits_comma_space_separator(self): + assert split_continuous_references("[1:aaa, 4:bbb, 7:ccc]") == "[1:aaa][4:bbb][7:ccc]" + + def test_splits_bare_comma_separator(self): + assert split_continuous_references("[1:aaa,4:bbb,7:ccc]") == "[1:aaa][4:bbb][7:ccc]" + + def test_splits_mixed_comma_styles_bare_then_space(self): + # Regression for #2417: bare comma followed by ", " style + # previously left the first pair merged. + assert ( + split_continuous_references("[1:92ff35fb,4:bfe6f044, 7:abcd1234]") + == "[1:92ff35fb][4:bfe6f044][7:abcd1234]" + ) + + def test_splits_mixed_comma_styles_space_then_bare(self): + assert split_continuous_references("[1:aaa, 4:bbb,7:ccc]") == "[1:aaa][4:bbb][7:ccc]" + + def test_splits_comma_with_multiple_spaces(self): + assert split_continuous_references("[1:aaa, 4:bbb,\t7:ccc]") == "[1:aaa][4:bbb][7:ccc]" + + def test_single_reference_unchanged(self): + assert split_continuous_references("[1:aaa]") == "[1:aaa]" + + def test_empty_string_returned_as_is(self): + assert split_continuous_references("") == "" + + def test_text_without_brackets_unchanged(self): + assert split_continuous_references("no brackets here") == "no brackets here" + + def test_text_with_multiple_bracket_pairs_unchanged(self): + original = "[1:aaa] and [2:bbb, 3:ccc]" + assert split_continuous_references(original) == original + + def test_text_without_comma_between_brackets_unchanged(self): + assert split_continuous_references("hello [1:aaa] world") == "hello [1:aaa] world" + + def test_prefix_and_suffix_preserved(self): + assert ( + split_continuous_references("prefix [1:aaa, 4:bbb] suffix") + == "prefix [1:aaa][4:bbb] suffix" + ) + + def test_reversed_brackets_returned_unchanged(self): + assert split_continuous_references("]1:aaa,2:bbb[") == "]1:aaa,2:bbb[" + + @pytest.mark.parametrize( + "raw,expected", + [ + ("[1:x,2:y]", "[1:x][2:y]"), + ("[1:x, 2:y]", "[1:x][2:y]"), + ("[1:x, 2:y]", "[1:x][2:y]"), + ("[1:x,\t2:y]", "[1:x][2:y]"), + ], + ) + def test_parametrized_trailing_whitespace_variants(self, raw, expected): + assert split_continuous_references(raw) == expected