Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 31 additions & 14 deletions src/memos/chunkers/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]:
"""
Expand All @@ -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:
Expand All @@ -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."""
71 changes: 66 additions & 5 deletions src/memos/chunkers/simple_chunker.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -8,6 +20,45 @@ def __init__(self, chunk_size: int, chunk_overlap: int):
def chunk(self, text: str, **kwargs) -> list[str]:
return self._simple_split_text(text, self.chunk_size, self.chunk_overlap)

@staticmethod
def _placeholder_spans(protected_text: str, url_map: dict[str, str]) -> list[tuple[int, int]]:
"""Return the protected-text ranges occupied by URL placeholders."""
spans = []
for placeholder in url_map:
placeholder_start = protected_text.find(placeholder)
if placeholder_start >= 0:
spans.append((placeholder_start, placeholder_start + len(placeholder)))
return sorted(spans)

@staticmethod
def _align_end_to_placeholder(
end: int,
start: int,
chunk_overlap: int,
placeholder_spans: list[tuple[int, int]],
) -> int:
"""Move a chunk end away from the middle of a URL placeholder."""
for placeholder_start, placeholder_end in placeholder_spans:
if placeholder_start < end < placeholder_end:
if placeholder_start - start > chunk_overlap:
return placeholder_start
return placeholder_end
return end

@staticmethod
def _align_start_to_placeholder(
next_start: int,
previous_start: int,
placeholder_spans: list[tuple[int, int]],
) -> int:
"""Keep overlap starts from landing in the middle of a URL placeholder."""
for placeholder_start, placeholder_end in placeholder_spans:
if placeholder_start < next_start < placeholder_end:
if placeholder_start > previous_start:
return placeholder_start
return placeholder_end
return next_start

def _simple_split_text(self, text: str, chunk_size: int, chunk_overlap: int) -> list[str]:
"""
Simple text splitter as fallback when langchain is not available.
Expand All @@ -29,6 +80,7 @@ def _simple_split_text(self, text: str, chunk_size: int, chunk_overlap: int) ->
chunks = []
start = 0
text_len = len(protected_text)
placeholder_spans = self._placeholder_spans(protected_text, url_map)

while start < text_len:
# Calculate end position
Expand All @@ -39,15 +91,24 @@ def _simple_split_text(self, text: str, chunk_size: int, chunk_overlap: int) ->
# Try to break at newline, sentence end, or space
for separator in ["\n\n", "\n", "。", "!", "?", ". ", "! ", "? ", " "]:
last_sep = protected_text.rfind(separator, start, end)
if last_sep != -1:
end = last_sep + len(separator)
if last_sep == -1:
continue
split_end = last_sep + len(separator)
if split_end - start > chunk_overlap:
end = split_end
break

end = self._align_end_to_placeholder(end, start, chunk_overlap, placeholder_spans)

chunk = protected_text[start:end].strip()
if chunk:
chunks.append(chunk)

if end >= text_len:
break

# Move start position with overlap
start = max(start + 1, end - chunk_overlap)
next_start = max(start + 1, end - chunk_overlap)
start = self._align_start_to_placeholder(next_start, start, placeholder_spans)

return [self.restore_urls(chunk, url_map) for chunk in chunks]
110 changes: 110 additions & 0 deletions tests/chunkers/test_simple_chunker.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
"""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_does_not_cut_placeholder_at_chunk_boundary():
"""A raw chunk boundary inside a placeholder must be moved to a safe edge."""
url = "https://example.com/path"
text = "A" * 95 + url + " " + "B" * 120
splitter = SimpleTextSplitter(chunk_size=100, chunk_overlap=20)

chunks = splitter.chunk(text)

assert any(url in chunk for chunk in chunks)
assert all(URLProtectionMixin._URL_PLACEHOLDER_PREFIX not in chunk for chunk in chunks)
assert all("https://" not in chunk or url in chunk for chunk in chunks)
assert len(chunks) < 10, "splitter made insufficient progress around the URL boundary"


def test_simple_text_splitter_stops_after_emitting_final_chunk():
"""The final overlap must not be emitted repeatedly as shrinking suffixes."""
splitter = SimpleTextSplitter(chunk_size=100, chunk_overlap=20)

chunks = splitter.chunk("X" * 150)

assert chunks == ["X" * 100, "X" * 70]


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)
Loading