fix: use each chunk's own position for its wider_text window

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

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