diff --git a/haystack/components/preprocessors/embedding_based_document_splitter.py b/haystack/components/preprocessors/embedding_based_document_splitter.py index 4d34298785c..8d0f439f46e 100644 --- a/haystack/components/preprocessors/embedding_based_document_splitter.py +++ b/haystack/components/preprocessors/embedding_based_document_splitter.py @@ -178,6 +178,7 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]: - `documents`: List of documents with the split texts. Each document includes: - A metadata field `source_id` to track the original document. - A metadata field `split_id` to track the split number. + - A metadata field `split_idx_start` with the character offset of the chunk in the original document. - A metadata field `page_number` to track the original page number. - All other metadata copied from the original document. @@ -218,6 +219,7 @@ async def run_async(self, documents: list[Document]) -> dict[str, list[Document] - `documents`: List of documents with the split texts. Each document includes: - A metadata field `source_id` to track the original document. - A metadata field `split_id` to track the split number. + - A metadata field `split_idx_start` with the character offset of the chunk in the original document. - A metadata field `page_number` to track the original page number. - All other metadata copied from the original document. @@ -489,12 +491,16 @@ def _create_documents_from_splits(splits: list[str], original_doc: Document) -> metadata = deepcopy(original_doc.meta) metadata["source_id"] = original_doc.id - # Calculate page numbers for each split + # Track page number and character offset across splits. + # Chunks are contiguous substrings of the original text (all joins are pure string concatenation), + # so the character offset is simply accumulated by adding the length of each chunk. current_page = 1 + current_char_pos = 0 for i, split_text in enumerate(splits): split_meta = deepcopy(metadata) split_meta["split_id"] = i + split_meta["split_idx_start"] = current_char_pos # Calculate page number for this split # Count page breaks in the split itself @@ -506,7 +512,8 @@ def _create_documents_from_splits(splits: list[str], original_doc: Document) -> doc = Document(content=split_text, meta=split_meta) documents.append(doc) - # Update page counter for next split + # Advance position and page counter for the next split + current_char_pos += len(split_text) current_page += page_breaks_in_split return documents diff --git a/releasenotes/notes/fix-embedding-splitter-missing-split-idx-start-8727a44789f5256a.yaml b/releasenotes/notes/fix-embedding-splitter-missing-split-idx-start-8727a44789f5256a.yaml new file mode 100644 index 00000000000..b28d67dfdec --- /dev/null +++ b/releasenotes/notes/fix-embedding-splitter-missing-split-idx-start-8727a44789f5256a.yaml @@ -0,0 +1,8 @@ +--- +fixes: + - | + ``EmbeddingBasedDocumentSplitter`` now sets ``split_idx_start`` in the metadata of every + output chunk, consistent with ``DocumentSplitter`` and ``RecursiveDocumentSplitter``. + Previously the field was absent, causing a ``KeyError`` in any downstream code that + reads character offsets from splitter output (e.g. highlighting, provenance tracking, + overlap recovery). diff --git a/test/components/preprocessors/test_embedding_based_document_splitter.py b/test/components/preprocessors/test_embedding_based_document_splitter.py index 96ea195f4ff..470df9ee2b8 100644 --- a/test/components/preprocessors/test_embedding_based_document_splitter.py +++ b/test/components/preprocessors/test_embedding_based_document_splitter.py @@ -266,6 +266,69 @@ def test_create_documents_from_splits_with_consecutive_page_breaks(self): # Should be page 2, not 4, because consecutive page breaks at the end are adjusted assert documents[1].meta["page_number"] == 2 + def test_create_documents_from_splits_split_idx_start(self): + """_create_documents_from_splits must set split_idx_start to the character offset of each chunk.""" + mock_embedder = Mock() + splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder) + + text = "First chunk. Second chunk. Third chunk." + splits = ["First chunk. ", "Second chunk. ", "Third chunk."] + original_doc = Document(content=text) + + documents = splitter._create_documents_from_splits(splits, original_doc) + + assert len(documents) == 3 + assert documents[0].meta["split_idx_start"] == 0 + assert documents[1].meta["split_idx_start"] == len("First chunk. ") + assert documents[2].meta["split_idx_start"] == len("First chunk. ") + len("Second chunk. ") + # Cross-check: split_idx_start correctly points into the original text + for doc in documents: + start = doc.meta["split_idx_start"] + assert text[start : start + len(doc.content)] == doc.content + + def test_run_split_idx_start(self): + """run() must produce chunks with correct split_idx_start character offsets.""" + text = "The sky is blue. The grass is green. The sun is yellow." + + # Mock embedder that returns two distinct embedding clusters so the splitter + # finds at least one split point + def mock_run(documents): + from dataclasses import replace as dc_replace + + embeddings = [] + for i, doc in enumerate(documents): + # Two distinct embedding directions — triggers a split + if i < len(documents) // 2: + embeddings.append(dc_replace(doc, embedding=[1.0, 0.0, 0.0])) + else: + embeddings.append(dc_replace(doc, embedding=[0.0, 1.0, 0.0])) + return {"documents": embeddings} + + mock_embedder = Mock() + mock_embedder.run = Mock(side_effect=mock_run) + + splitter = EmbeddingBasedDocumentSplitter( + document_embedder=mock_embedder, sentences_per_group=1, percentile=0.5, min_length=0, max_length=10000 + ) + splitter.warm_up() + + result = splitter.run(documents=[Document(content=text)]) + chunks = result["documents"] + + # All chunks must have split_idx_start + for chunk in chunks: + assert "split_idx_start" in chunk.meta + + # split_idx_start must point to the correct position in the original text + for chunk in chunks: + start = chunk.meta["split_idx_start"] + assert text[start : start + len(chunk.content)] == chunk.content + + # Offsets must be strictly increasing (chunks are non-empty and contiguous) + starts = [c.meta["split_idx_start"] for c in chunks] + assert starts == sorted(starts) + assert starts[0] == 0 + def test_calculate_embeddings(self): mock_embedder = Mock()