diff --git a/src/memos/chunkers/base.py b/src/memos/chunkers/base.py index e858132e1..b78fad128 100644 --- a/src/memos/chunkers/base.py +++ b/src/memos/chunkers/base.py @@ -14,16 +14,22 @@ def __init__(self, text: str, token_count: int, sentences: list[str]): self.sentences = sentences -class BaseChunker(ABC): - """Base class for all text chunkers.""" - - @abstractmethod - def __init__(self, config: BaseChunkerConfig): - """Initialize the chunker with the given configuration.""" - - @abstractmethod - def chunk(self, text: str) -> list[Chunk]: - """Chunk the given text into smaller chunks.""" +class URLProtectionMixin: + """Shared URL protect/restore helpers used across chunkers. + + Extracted so that lightweight fallbacks such as + :class:`memos.chunkers.simple_chunker.SimpleTextSplitter` can reuse the + same URL-aware splitting logic as :class:`BaseChunker` without inheriting + the full chunker contract (see issue #2115). + """ + + _URL_PATTERN = r'https?://[^\s<>"{}|\\^`\[\]]+' + # Prefix used for the placeholders emitted by :meth:`protect_urls`. Exposed + # as a class-level constant so tests (and any other consumer that needs to + # detect leaked placeholders) can reference it without hardcoding the + # literal — keeping the assertion in sync with the implementation if the + # placeholder format ever changes. + _URL_PLACEHOLDER_PREFIX = "__URL_" def protect_urls(self, text: str) -> tuple[str, dict[str, str]]: """ @@ -35,16 +41,15 @@ def protect_urls(self, text: str) -> tuple[str, dict[str, str]]: Returns: tuple: (Text with URLs replaced by placeholders, URL mapping dictionary) """ - url_pattern = r'https?://[^\s<>"{}|\\^`\[\]]+' - url_map = {} + url_map: dict[str, str] = {} def replace_url(match): url = match.group(0) - placeholder = f"__URL_{len(url_map)}__" + placeholder = f"{self._URL_PLACEHOLDER_PREFIX}{len(url_map)}__" url_map[placeholder] = url return placeholder - protected_text = re.sub(url_pattern, replace_url, text) + protected_text = re.sub(self._URL_PATTERN, replace_url, text) return protected_text, url_map def restore_urls(self, text: str, url_map: dict[str, str]) -> str: @@ -63,3 +68,15 @@ def restore_urls(self, text: str, url_map: dict[str, str]) -> str: restored_text = restored_text.replace(placeholder, url) return restored_text + + +class BaseChunker(URLProtectionMixin, ABC): + """Base class for all text chunkers.""" + + @abstractmethod + def __init__(self, config: BaseChunkerConfig): + """Initialize the chunker with the given configuration.""" + + @abstractmethod + def chunk(self, text: str) -> list[Chunk]: + """Chunk the given text into smaller chunks.""" diff --git a/src/memos/chunkers/simple_chunker.py b/src/memos/chunkers/simple_chunker.py index 58e12e2f1..c4bccd667 100644 --- a/src/memos/chunkers/simple_chunker.py +++ b/src/memos/chunkers/simple_chunker.py @@ -1,5 +1,17 @@ -class SimpleTextSplitter: - """Simple text splitter wrapper.""" +from memos.chunkers.base import URLProtectionMixin + + +class SimpleTextSplitter(URLProtectionMixin): + """Simple text splitter wrapper. + + Fallback used by :mod:`memos.mem_reader.read_multi_modal.utils` when the + optional ``langchain_text_splitters``-backed chunkers (``CharacterTextChunker`` + / ``MarkdownChunker``) cannot be constructed at import time. + + Inherits URL protect/restore helpers from :class:`URLProtectionMixin` + (see issue #2115: without the mixin, ``chunk()`` raised ``AttributeError`` + on every call that reached the fallback path). + """ def __init__(self, chunk_size: int, chunk_overlap: int): self.chunk_size = chunk_size diff --git a/tests/chunkers/test_simple_chunker.py b/tests/chunkers/test_simple_chunker.py new file mode 100644 index 000000000..405995a7e --- /dev/null +++ b/tests/chunkers/test_simple_chunker.py @@ -0,0 +1,87 @@ +"""Regression tests for `SimpleTextSplitter` fallback (issue #2115). + +The fallback is exercised in production when `langchain_text_splitters` is +missing (ACK image drift). Prior to the fix, `SimpleTextSplitter.chunk()` +raised `AttributeError: 'SimpleTextSplitter' object has no attribute +'protect_urls'` because `_simple_split_text` referenced `self.protect_urls` +/ `self.restore_urls`, which are only defined on `BaseChunker`. +""" + +import pytest + +from memos.chunkers.base import URLProtectionMixin +from memos.chunkers.simple_chunker import SimpleTextSplitter + + +def test_simple_text_splitter_short_text_with_url_returns_single_chunk(): + """Short text below chunk_size should return one chunk with the URL intact.""" + splitter = SimpleTextSplitter(chunk_size=512, chunk_overlap=128) + text = "This is a test document with a URL: https://example.com/path/to/resource" + + chunks = splitter.chunk(text) + + assert chunks == [text] + + +def test_simple_text_splitter_long_text_preserves_url(): + """A URL must never be split across chunks — it either appears whole or not at all in a chunk. + + The critical property (issue #2115): even after fallback splitting, we + must never see a chunk that contains only part of a URL. Overlap MAY + cause the same URL to appear in more than one chunk; that is by design + for retrieval quality and is not what the issue asks us to change. + """ + url = "https://example.com/very/long/path/segment?query=one&other=two#fragment" + prefix = "A" * 400 + suffix = "B" * 400 + text = f"{prefix} {url} {suffix}" + + splitter = SimpleTextSplitter(chunk_size=200, chunk_overlap=50) + chunks = splitter.chunk(text) + + assert len(chunks) > 1, "text should be split into multiple chunks" + # The URL must appear whole at least once. + assert any(url in c for c in chunks), ( + f"URL was fully lost after splitting; chunks (first 5)={chunks[:5]}" + ) + # No chunk should contain the placeholder marker leftover. + for c in chunks: + assert URLProtectionMixin._URL_PLACEHOLDER_PREFIX not in c, ( + f"unresolved URL placeholder leaked into chunk: {c!r}" + ) + # No chunk should contain a *partial* URL — i.e., if the chunk mentions + # "https://" it must contain the URL in full. + for c in chunks: + if "https://" in c: + assert url in c, f"chunk contains a partial URL: {c!r}" + + +def test_simple_text_splitter_empty_input_returns_empty_list(): + splitter = SimpleTextSplitter(chunk_size=100, chunk_overlap=20) + assert splitter.chunk("") == [] + assert splitter.chunk(" \n \t ") == [] + + +def test_simple_text_splitter_no_url_still_chunks(): + splitter = SimpleTextSplitter(chunk_size=50, chunk_overlap=10) + text = "Hello world. " * 20 # > 50 chars, no URL + chunks = splitter.chunk(text) + assert len(chunks) >= 2 + # Reassembling should recover all non-whitespace content. + joined = "".join(chunks) + for word in ["Hello", "world"]: + assert word in joined + + +@pytest.mark.parametrize( + ("chunk_size", "overlap"), + [(100, 20), (256, 64), (1024, 128)], +) +def test_simple_text_splitter_various_sizes_do_not_raise(chunk_size, overlap): + """The fallback used to raise AttributeError for *any* input containing a URL.""" + splitter = SimpleTextSplitter(chunk_size=chunk_size, chunk_overlap=overlap) + text = "prefix " + ("word " * 200) + "https://example.com/x " + ("tail " * 200) + # Must not raise. + chunks = splitter.chunk(text) + assert isinstance(chunks, list) + assert all(isinstance(c, str) for c in chunks)