Ali Zarinkolah
5cdfb70085
Adds src/application/ingestion/ -- Persian normalization, DOCX body walk with structural data/layout table classification, CSV/XLSX row rendering, and fixed-size token chunking (cl100k_base, 400/60/512) -- as pure functions per ADR-0015, tested against real production documents (asia_data_sample, kept out of the repo). ADR-0018 records where this diverges from ADR-0004 (fixed-size default, no invented headings/tree, structural table classification, header-provable labeling only). Plan 001's scope line is corrected from CSV-only to DOCX/XLSX/CSV, and CLAUDE.md's stale project-status paragraph is updated to match current implementation state.
271 lines
9.2 KiB
Python
271 lines
9.2 KiB
Python
"""Fixed-size chunking, chunk identity, and neighbor linking (ADR-0001, ADR-0018)."""
|
|
|
|
import uuid
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from src.application.ingestion import (
|
|
ContentType,
|
|
ParsedDocument,
|
|
StructuralUnit,
|
|
chunk_document,
|
|
chunk_id_for,
|
|
count_tokens,
|
|
split_by_tokens,
|
|
)
|
|
from src.config import ChunkingSettings
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
ENCODING = "cl100k_base"
|
|
|
|
FILE_ID = uuid.UUID("11111111-1111-1111-1111-111111111111")
|
|
OTHER_FILE_ID = uuid.UUID("22222222-2222-2222-2222-222222222222")
|
|
|
|
|
|
@pytest.fixture
|
|
def settings() -> ChunkingSettings:
|
|
return ChunkingSettings()
|
|
|
|
|
|
def _long_text(tokens: int) -> str:
|
|
"""Text of roughly `tokens` cl100k tokens: one common word per token."""
|
|
return " ".join(["word"] * tokens)
|
|
|
|
|
|
def _rows(rows: list[str]) -> ParsedDocument:
|
|
"""A parsed spreadsheet: every row an atomic table_row unit."""
|
|
return ParsedDocument(
|
|
units=[StructuralUnit(text=row, content_type=ContentType.TABLE_ROW) for row in rows],
|
|
block_count=len(rows),
|
|
)
|
|
|
|
|
|
def _prose(text: str) -> ParsedDocument:
|
|
"""A parsed document consisting of one flowing prose block."""
|
|
return ParsedDocument(
|
|
units=[StructuralUnit(text=text, content_type=ContentType.PARAGRAPH)],
|
|
block_count=1,
|
|
)
|
|
|
|
|
|
# ── split_by_tokens ───────────────────────────────────────────────
|
|
|
|
|
|
def test_split_by_tokens_text_under_limit_returns_single_chunk() -> None:
|
|
text = _long_text(50)
|
|
|
|
assert split_by_tokens(text, chunk_size=400, overlap=60, encoding_name=ENCODING) == [text]
|
|
|
|
|
|
def test_split_by_tokens_long_text_splits_into_multiple_windows() -> None:
|
|
windows = split_by_tokens(_long_text(1000), chunk_size=400, overlap=60, encoding_name=ENCODING)
|
|
|
|
assert len(windows) > 1
|
|
assert all(count_tokens(window, ENCODING) <= 400 for window in windows)
|
|
|
|
|
|
def test_split_by_tokens_consecutive_windows_share_overlapping_tokens() -> None:
|
|
windows = split_by_tokens(
|
|
" ".join(str(number) for number in range(1000)),
|
|
chunk_size=100,
|
|
overlap=20,
|
|
encoding_name=ENCODING,
|
|
)
|
|
|
|
# The tail of one window must reappear at the head of the next, or a
|
|
# sentence cut at the boundary is unrecoverable.
|
|
tail_tokens = windows[0].split()[-5:]
|
|
assert " ".join(tail_tokens) in windows[1]
|
|
|
|
|
|
def test_split_by_tokens_overlap_at_least_chunk_size_raises() -> None:
|
|
"""Guards an infinite loop: `start = end - overlap` never advances."""
|
|
with pytest.raises(ValueError, match="smaller than chunk_size"):
|
|
split_by_tokens("text", chunk_size=100, overlap=100, encoding_name=ENCODING)
|
|
|
|
|
|
# ── chunk identity ────────────────────────────────────────────────
|
|
|
|
|
|
def test_chunk_id_for_is_deterministic_across_calls() -> None:
|
|
assert chunk_id_for(FILE_ID, 7) == chunk_id_for(FILE_ID, 7)
|
|
|
|
|
|
def test_chunk_id_for_differs_across_files_and_indices() -> None:
|
|
assert chunk_id_for(FILE_ID, 0) != chunk_id_for(OTHER_FILE_ID, 0)
|
|
assert chunk_id_for(FILE_ID, 0) != chunk_id_for(FILE_ID, 1)
|
|
|
|
|
|
def test_chunk_document_identical_content_in_different_files_gets_different_ids(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
"""The evaluation repo derived ids from the filename with non-ASCII stripped,
|
|
so every Farsi filename collapsed to underscores and collided."""
|
|
parsed = _rows(["a: 1"])
|
|
|
|
first = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
second = chunk_document(parsed, file_id=OTHER_FILE_ID, settings=settings)
|
|
|
|
assert first[0].chunk_id != second[0].chunk_id
|
|
|
|
|
|
def test_chunk_document_reingesting_same_file_yields_same_ids(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
"""Deterministic ids are what make a re-upload upsert instead of duplicate."""
|
|
parsed = _rows(["a: 1", "a: 2"])
|
|
|
|
first = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
second = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert [chunk.chunk_id for chunk in first] == [chunk.chunk_id for chunk in second]
|
|
|
|
|
|
# ── ordering and neighbor links ───────────────────────────────────
|
|
|
|
|
|
def test_chunk_document_assigns_contiguous_indices_and_one_based_order(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
parsed = _rows(["a: 1", "a: 2", "a: 3"])
|
|
|
|
chunks = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert [chunk.chunk_index for chunk in chunks] == [0, 1, 2]
|
|
assert [chunk.order_id for chunk in chunks] == [1.0, 2.0, 3.0]
|
|
|
|
|
|
def test_chunk_document_blank_rows_do_not_leave_index_gaps(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
"""An index gap breaks the previous/next chain retrieval expansion walks."""
|
|
parsed = _rows(["a: 1", " ", "a: 2"])
|
|
|
|
chunks = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert [chunk.chunk_index for chunk in chunks] == [0, 1]
|
|
|
|
|
|
def test_chunk_document_links_neighbors_with_null_at_both_ends(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
parsed = _rows(["a: 1", "a: 2", "a: 3"])
|
|
|
|
chunks = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert chunks[0].previous_chunk_id is None
|
|
assert chunks[-1].next_chunk_id is None
|
|
for position, chunk in enumerate(chunks[:-1]):
|
|
assert chunk.next_chunk_id == chunks[position + 1].chunk_id
|
|
assert chunks[position + 1].previous_chunk_id == chunk.chunk_id
|
|
|
|
|
|
def test_chunk_document_single_chunk_has_no_neighbors(settings: ChunkingSettings) -> None:
|
|
parsed = _rows(["a: 1"])
|
|
|
|
(chunk,) = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert chunk.previous_chunk_id is None
|
|
assert chunk.next_chunk_id is None
|
|
|
|
|
|
# ── content types and the token cap ───────────────────────────────
|
|
|
|
|
|
def test_chunk_document_spreadsheet_row_becomes_one_atomic_chunk(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
rows = ["branch: one\ncity: two", "branch: three\ncity: four"]
|
|
parsed = _rows(rows)
|
|
|
|
chunks = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert [chunk.content for chunk in chunks] == rows
|
|
assert all(chunk.content_type is ContentType.TABLE_ROW for chunk in chunks)
|
|
|
|
|
|
def test_chunk_document_oversized_row_splits_rather_than_truncating(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
"""Silent truncation is the failure mode: the model drops the tail, not raises."""
|
|
oversized = _long_text(settings.max_chunk_tokens * 2)
|
|
parsed = _rows([oversized])
|
|
|
|
chunks = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert len(chunks) > 1
|
|
assert all(chunk.token_count <= settings.max_chunk_tokens for chunk in chunks)
|
|
|
|
|
|
def test_chunk_document_markdown_produces_paragraph_chunks(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
parsed = _prose("# Heading\n\nBody text.")
|
|
|
|
chunks = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert all(chunk.content_type is ContentType.PARAGRAPH for chunk in chunks)
|
|
|
|
|
|
def test_chunk_document_no_chunk_exceeds_the_model_sequence_length(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
parsed = _prose(_long_text(5000))
|
|
|
|
chunks = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert chunks
|
|
assert all(chunk.token_count <= settings.max_chunk_tokens for chunk in chunks)
|
|
|
|
|
|
def test_chunk_document_persian_text_never_exceeds_the_cap(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
"""Persian is the case where the cap can actually be breached.
|
|
|
|
Slicing a token list can cut a multi-byte character in half; decoding then
|
|
re-encoding that window does not always round-trip to the same token count.
|
|
`ChunkTooLargeError` is the guard for exactly this, and it must not fire.
|
|
"""
|
|
persian_word = "".join(chr(code) for code in (0x0628, 0x06CC, 0x0645, 0x0647))
|
|
parsed = _prose(" ".join([persian_word] * 4000))
|
|
|
|
chunks = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert len(chunks) > 1
|
|
assert all(chunk.token_count <= settings.max_chunk_tokens for chunk in chunks)
|
|
|
|
|
|
def test_chunk_document_records_token_and_character_counts(
|
|
settings: ChunkingSettings,
|
|
) -> None:
|
|
parsed = _rows(["branch: one"])
|
|
|
|
(chunk,) = chunk_document(parsed, file_id=FILE_ID, settings=settings)
|
|
|
|
assert chunk.character_count == len("branch: one")
|
|
assert chunk.token_count == count_tokens("branch: one", ENCODING)
|
|
|
|
|
|
# ── settings validation ───────────────────────────────────────────
|
|
|
|
|
|
def test_chunking_settings_overlap_not_smaller_than_size_is_rejected() -> None:
|
|
with pytest.raises(ValidationError, match="must be smaller than"):
|
|
ChunkingSettings(chunk_size=100, chunk_overlap=100)
|
|
|
|
|
|
def test_chunking_settings_size_over_model_limit_is_rejected() -> None:
|
|
with pytest.raises(ValidationError, match="must not exceed"):
|
|
ChunkingSettings(chunk_size=600, max_chunk_tokens=512)
|
|
|
|
|
|
def test_chunking_settings_defaults_match_adr_0018() -> None:
|
|
settings = ChunkingSettings()
|
|
|
|
assert (settings.chunk_size, settings.chunk_overlap) == (400, 60)
|
|
assert settings.max_chunk_tokens == 512
|
|
assert settings.strategy == "fixed_size"
|