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
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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).
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
Loading