diff --git a/src/memos/chunkers/simple_chunker.py b/src/memos/chunkers/simple_chunker.py index 58e12e2f1..cebe4fa19 100644 --- a/src/memos/chunkers/simple_chunker.py +++ b/src/memos/chunkers/simple_chunker.py @@ -1,7 +1,25 @@ -class SimpleTextSplitter: - """Simple text splitter wrapper.""" +"""Simple, dependency-free text splitter used as fallback chunker.""" + +from memos.chunkers.base import BaseChunker +from memos.configs.chunker import BaseChunkerConfig + + +class SimpleTextSplitter(BaseChunker): + """Simple text splitter wrapper with URL protection support.""" def __init__(self, chunk_size: int, chunk_overlap: int): + """Initialize the splitter with explicit chunk_size/overlap. + + A minimal ``BaseChunkerConfig`` is synthesized so the base class can + carry the ``protect_urls`` / ``restore_urls`` helpers without + requiring the caller to build a full config. + """ + config = BaseChunkerConfig( + tokenizer_or_token_counter="simple", + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + super().__init__(config) self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap @@ -9,35 +27,37 @@ def chunk(self, text: str, **kwargs) -> list[str]: return self._simple_split_text(text, self.chunk_size, self.chunk_overlap) 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. - - Args: - text: Text to split - chunk_size: Maximum size of chunks - chunk_overlap: Overlap between chunks + """Split text into chunks with URL preservation. - Returns: - List of text chunks + This is used as a fallback when langchain-style splitters are not + available. URLs are protected before splitting and restored after + so that a URL is never truncated in the middle. """ protected_text, url_map = self.protect_urls(text) if not protected_text or len(protected_text) <= chunk_size: - chunks = [protected_text] if protected_text.strip() else [] + chunks = [protected_text] if protected_text and protected_text.strip() else [] return [self.restore_urls(chunk, url_map) for chunk in chunks] - chunks = [] + chunks: list[str] = [] start = 0 text_len = len(protected_text) while start < text_len: - # Calculate end position end = min(start + chunk_size, text_len) - # If not the last chunk, try to break at a good position if end < text_len: - # Try to break at newline, sentence end, or space - for separator in ["\n\n", "\n", "。", "！", "？", ". ", "! ", "? ", " "]: + for separator in [ + "\n\n", + "\n", + "。", + "！", + "？", + ". ", + "! ", + "? ", + " ", + ]: last_sep = protected_text.rfind(separator, start, end) if last_sep != -1: end = last_sep + len(separator) @@ -47,7 +67,6 @@ def _simple_split_text(self, text: str, chunk_size: int, chunk_overlap: int) -> if chunk: chunks.append(chunk) - # Move start position with overlap start = max(start + 1, end - chunk_overlap) return [self.restore_urls(chunk, url_map) for chunk in chunks] diff --git a/tests/chunkers/test_simple_chunker.py b/tests/chunkers/test_simple_chunker.py new file mode 100644 index 000000000..5f66e3505 --- /dev/null +++ b/tests/chunkers/test_simple_chunker.py @@ -0,0 +1,61 @@ +"""Tests for SimpleTextSplitter fallback chunker.""" + +from memos.chunkers.simple_chunker import SimpleTextSplitter + + +class TestSimpleTextSplitterConstruction: + def test_construct_and_expose_helpers(self): + splitter = SimpleTextSplitter(chunk_size=128, chunk_overlap=8) + assert splitter.chunk_size == 128 + assert splitter.chunk_overlap == 8 + protected, url_map = splitter.protect_urls("See https://example.com/x") + assert "https://example.com/x" not in protected + assert "__URL_" in protected + restored = splitter.restore_urls(protected, url_map) + assert restored == "See https://example.com/x" + + +class TestSimpleTextSplitterUrlPreservation: + def test_url_not_truncated_midstring(self): + url = "https://very.long.example.com/path/to/resource?a=1&b=2&c=3&d=4" + text = f"Intro. {url} Conclusion." + # Small chunk size forces a break near the URL; URL must stay whole. + splitter = SimpleTextSplitter(chunk_size=16, chunk_overlap=0) + chunks = splitter.chunk(text) + joined = "".join(chunks) + assert url in joined + for chunk in chunks: + # If chunk contains a portion of URL placeholder, original was intact + if "__URL_" in chunk: + assert splitter.restore_urls(chunk, splitter.protect_urls(text)[1]) + + def test_empty_and_whitespace_input(self): + splitter = SimpleTextSplitter(chunk_size=32, chunk_overlap=0) + assert splitter.chunk("") == [] + assert splitter.chunk(" \n\n ") == [] + + def test_short_text_returns_single_chunk(self): + splitter = SimpleTextSplitter(chunk_size=512, chunk_overlap=0) + text = "Hello world. This is a test." + result = splitter.chunk(text) + assert len(result) == 1 + assert result[0] == text + + def test_long_text_produces_overlapping_chunks(self): + text = "Paragraph one.\n\n" * 20 + splitter = SimpleTextSplitter(chunk_size=64, chunk_overlap=8) + chunks = splitter.chunk(text) + assert len(chunks) > 1 + for c in chunks: + assert isinstance(c, str) + assert len(c) <= 64 + 8 # a little slack for separator inclusion + + +class TestSimpleTextSplitterSplitBoundaries: + def test_prefers_sentence_breaks(self): + text = "First sentence. Second sentence. Third sentence. Fourth sentence." + splitter = SimpleTextSplitter(chunk_size=40, chunk_overlap=0) + chunks = splitter.chunk(text) + assert len(chunks) >= 2 + # First chunk should end at/near a sentence boundary + assert chunks[0].rstrip().endswith((".", "。", "！", "？")) or " " in chunks[0]