diff --git a/deepsearcher/loader/splitter.py b/deepsearcher/loader/splitter.py index ebdf70ae..492ee7a9 100644 --- a/deepsearcher/loader/splitter.py +++ b/deepsearcher/loader/splitter.py @@ -44,7 +44,10 @@ def __init__( def _sentence_window_split( - split_docs: List[Document], original_document: Document, offset: int = 200 + split_docs: List[Document], + original_document: Document, + offset: int = 200, + chunk_overlap: int = 0, ) -> List[Chunk]: """ Create chunks with context windows from split documents. @@ -57,19 +60,38 @@ def _sentence_window_split( split_docs: List of documents that have been split. original_document: The original document before splitting. offset: Number of characters to include before and after each split piece. + chunk_overlap: Maximum number of characters the splitter lets consecutive + split pieces overlap by. Returns: A list of Chunk objects with context windows. """ chunks = [] original_text = original_document.page_content + search_start = 0 + after_previous_start = 0 for doc in split_docs: doc_text = doc.page_content - start_index = original_text.index(doc_text) - end_index = start_index + len(doc_text) - 1 - wider_text = original_text[ - max(0, start_index - offset) : min(len(original_text), end_index + offset) - ] + # Split docs are in document order and each one starts after the previous + # one and at most `chunk_overlap` characters before its end, so search from + # there to locate repeated text at its own position rather than an earlier + # occurrence. Fall back to looser searches if that bound does not hold. + start_index = original_text.find(doc_text, search_start) + if start_index == -1: + start_index = original_text.find(doc_text, after_previous_start) + if start_index == -1: + start_index = original_text.find(doc_text) + if start_index == -1: + # The chunk text is not a substring of the original text (e.g. the + # splitter normalized it), so there is no position to widen around. + wider_text = doc_text + else: + end_index = start_index + len(doc_text) + after_previous_start = start_index + 1 + search_start = max(after_previous_start, end_index - chunk_overlap) + wider_text = original_text[ + max(0, start_index - offset) : min(len(original_text), end_index + offset) + ] reference = doc.metadata.pop("reference", "") doc.metadata["wider_text"] = wider_text chunk = Chunk(text=doc_text, reference=reference, metadata=doc.metadata) @@ -100,6 +122,8 @@ def split_docs_to_chunks( all_chunks = [] for doc in documents: split_docs = text_splitter.split_documents([doc]) - split_chunks = _sentence_window_split(split_docs, doc, offset=300) + split_chunks = _sentence_window_split( + split_docs, doc, offset=300, chunk_overlap=chunk_overlap + ) all_chunks.extend(split_chunks) return all_chunks diff --git a/tests/loader/test_splitter.py b/tests/loader/test_splitter.py index 40c2f9ab..358ca73d 100644 --- a/tests/loader/test_splitter.py +++ b/tests/loader/test_splitter.py @@ -1,5 +1,6 @@ import unittest from langchain_core.documents import Document +from langchain_text_splitters import RecursiveCharacterTextSplitter from deepsearcher.loader.splitter import Chunk, split_docs_to_chunks, _sentence_window_split @@ -59,6 +60,77 @@ def test_sentence_window_split(self): # With smaller offset, wider_text should be shorter than the full original text self.assertLessEqual(len(chunk.metadata["wider_text"]), len(original_text)) + def test_sentence_window_split_repeated_text(self): + """Test that repeated text gets the context window of its own position.""" + original_text = "x" + "." * 20 + "y" + "-" * 20 + "x" + original_doc = Document(page_content=original_text, metadata={"reference": "test_doc"}) + split_docs = [ + Document(page_content="x", metadata={"reference": "test_doc"}), + Document(page_content="y", metadata={"reference": "test_doc"}), + Document(page_content="x", metadata={"reference": "test_doc"}), + ] + + chunks = _sentence_window_split(split_docs, original_doc, offset=3) + + self.assertEqual(chunks[0].metadata["wider_text"], "x...") + self.assertEqual(chunks[1].metadata["wider_text"], "...y---") + self.assertEqual(chunks[2].metadata["wider_text"], "---x") + + def test_sentence_window_split_text_inside_previous_chunk(self): + """Test that a chunk whose text also occurs inside the previous chunk is not placed there.""" + original_doc = Document(page_content="XABCABC", metadata={"reference": "test_doc"}) + split_docs = [ + Document(page_content="XABC", metadata={"reference": "test_doc"}), + Document(page_content="ABC", metadata={"reference": "test_doc"}), + ] + + chunks = _sentence_window_split(split_docs, original_doc, offset=1) + + self.assertEqual(chunks[0].metadata["wider_text"], "XABCA") + self.assertEqual(chunks[1].metadata["wider_text"], "CABC") + + def test_sentence_window_split_matches_splitter_positions(self): + """Test that chunk positions match the splitter's own start indexes.""" + for original_text, chunk_size, chunk_overlap in [ + ("ab ab ab ab ab", 5, 0), + ("x ab ab", 6, 1), + ]: + with self.subTest(original_text=original_text, chunk_overlap=chunk_overlap): + splitter = RecursiveCharacterTextSplitter( + chunk_size=chunk_size, chunk_overlap=chunk_overlap, add_start_index=True + ) + split_docs = splitter.split_documents([Document(page_content=original_text)]) + expected = [] + for d in split_docs: + start = d.metadata["start_index"] + end = start + len(d.page_content) + expected.append(original_text[max(0, start - 1) : end + 1]) + + chunks = _sentence_window_split( + split_docs, + Document(page_content=original_text), + offset=1, + chunk_overlap=chunk_overlap, + ) + + self.assertEqual([c.metadata["wider_text"] for c in chunks], expected) + + def test_sentence_window_split_text_not_in_original(self): + """Test that a chunk missing from the original text falls back to its own text.""" + original_doc = Document(page_content="alpha beta", metadata={"reference": "test_doc"}) + split_docs = [ + Document(page_content="alpha", metadata={"reference": "test_doc"}), + Document(page_content="gamma", metadata={"reference": "test_doc"}), + Document(page_content="beta", metadata={"reference": "test_doc"}), + ] + + chunks = _sentence_window_split(split_docs, original_doc, offset=2) + + self.assertEqual(len(chunks), 3) + self.assertEqual(chunks[0].metadata["wider_text"], "alpha b") + self.assertEqual(chunks[1].metadata["wider_text"], "gamma") + self.assertEqual(chunks[2].metadata["wider_text"], "a beta") + def test_split_docs_to_chunks(self): """Test split_docs_to_chunks function.""" # Create test documents