117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
"""Unit coverage for deterministic structure-aware chunking."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from uuid import uuid4
|
|
|
|
from ragcore.domain.sources import SourceFingerprint
|
|
from ragcore.ingestion.chunking import ChunkProfile, StructureAwareChunker
|
|
from ragcore.ingestion.parsing import (
|
|
CanonicalArtifact,
|
|
CanonicalBlock,
|
|
CanonicalBlockKind,
|
|
ExtractionReport,
|
|
)
|
|
|
|
|
|
def artifact() -> CanonicalArtifact:
|
|
return CanonicalArtifact(
|
|
document_version_id=uuid4(),
|
|
source_fingerprint=SourceFingerprint("c" * 64),
|
|
blocks=(
|
|
CanonicalBlock(
|
|
CanonicalBlockKind.HEADING,
|
|
"/blocks/0000",
|
|
"Steps",
|
|
{"line_start": 1, "line_end": 1},
|
|
),
|
|
CanonicalBlock(
|
|
CanonicalBlockKind.PARAGRAPH,
|
|
"/blocks/0001",
|
|
"Do the thing.",
|
|
{"line_start": 2, "line_end": 2},
|
|
),
|
|
CanonicalBlock(
|
|
CanonicalBlockKind.HEADING,
|
|
"/blocks/0002",
|
|
"Code",
|
|
{"line_start": 4, "line_end": 4},
|
|
),
|
|
CanonicalBlock(
|
|
CanonicalBlockKind.CODE,
|
|
"/blocks/0003",
|
|
"def first():\n return 1\n\nclass Second:\n pass",
|
|
{"line_start": 5, "line_end": 9},
|
|
),
|
|
CanonicalBlock(
|
|
CanonicalBlockKind.TABLE,
|
|
"/blocks/0004",
|
|
"| A | B |\n|---|---|\n| 1 | 2 |\n| 3 | 4 |",
|
|
{"line_start": 11, "line_end": 14},
|
|
),
|
|
),
|
|
extraction_report=ExtractionReport("text-parser", "1", "text-v1", False),
|
|
)
|
|
|
|
|
|
def test_same_canonical_input_and_profile_has_identical_ids() -> None:
|
|
source = artifact()
|
|
first = StructureAwareChunker().chunk(source, ChunkProfile.PROSE)
|
|
second = StructureAwareChunker().chunk(source, ChunkProfile.PROSE)
|
|
|
|
assert [chunk.id for chunk in first.chunks] == [chunk.id for chunk in second.chunks]
|
|
assert [chunk.content for chunk in first.chunks] == [
|
|
chunk.content for chunk in second.chunks
|
|
]
|
|
|
|
|
|
def test_section_symbol_and_table_boundaries_have_parent_links() -> None:
|
|
result = StructureAwareChunker().chunk(artifact(), ChunkProfile.CODE)
|
|
|
|
children = [chunk for chunk in result.chunks if chunk.parent_id is not None]
|
|
assert any("/symbol/0000" in chunk.structure_path for chunk in children)
|
|
assert any("/symbol/0001" in chunk.structure_path for chunk in children)
|
|
assert all(
|
|
chunk.parent_id
|
|
in {parent.id for parent in result.chunks if parent.parent_id is None}
|
|
for chunk in children
|
|
)
|
|
|
|
table_result = StructureAwareChunker().chunk(artifact(), ChunkProfile.TABLE)
|
|
table_children = [
|
|
chunk
|
|
for chunk in table_result.chunks
|
|
if chunk.source_structure_path == "/blocks/0004"
|
|
]
|
|
assert table_children
|
|
assert all("| A | B |" in chunk.content for chunk in table_children)
|
|
|
|
|
|
def test_token_maxima_and_profile_change_are_enforced_in_identity() -> None:
|
|
source = artifact()
|
|
long_block = CanonicalBlock(
|
|
CanonicalBlockKind.PARAGRAPH,
|
|
"/blocks/long",
|
|
" ".join(f"word{index}" for index in range(2405)),
|
|
{"line_start": 1, "line_end": 2405},
|
|
)
|
|
source = CanonicalArtifact(
|
|
source.document_version_id,
|
|
source.source_fingerprint,
|
|
(long_block,),
|
|
source.extraction_report,
|
|
)
|
|
prose = StructureAwareChunker().chunk(source, ChunkProfile.PROSE)
|
|
runbook = StructureAwareChunker().chunk(source, ChunkProfile.RUNBOOK)
|
|
|
|
assert all(chunk.token_count <= 1200 for chunk in prose.chunks)
|
|
assert all(chunk.token_count <= 1200 for chunk in runbook.chunks)
|
|
assert {chunk.id for chunk in prose.chunks}.isdisjoint(
|
|
{chunk.id for chunk in runbook.chunks}
|
|
)
|
|
assert all(
|
|
chunk.parent_id is not None
|
|
for chunk in prose.chunks
|
|
if chunk.metadata["role"] == "child"
|
|
)
|