384 lines
14 KiB
Python
384 lines
14 KiB
Python
"""Unit tests: format-aware chunker (PLAN §5 policy, A9 formats).
|
|
|
|
The markdown policy tests are the original contract (md output stays
|
|
unchanged); the per-format tests cover the phase-09 dispatcher
|
|
(yaml/yml, json, py, txt) and the 1200-char hard cap for every format.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from itertools import pairwise
|
|
|
|
import pytest
|
|
|
|
from app.rag.chunker import (
|
|
HARD_MAX_CHARS,
|
|
chunk_document,
|
|
chunk_json,
|
|
chunk_markdown,
|
|
chunk_python,
|
|
chunk_text,
|
|
chunk_yaml,
|
|
extract_title,
|
|
)
|
|
|
|
ANCHOR = "## Big"
|
|
ANCHOR_PREFIX = f"{ANCHOR}\n\n"
|
|
|
|
|
|
def _paras(n: int, char: str = "l", width: int = 300) -> list[str]:
|
|
return [f"paragraph {i} " + char * (width - 12) for i in range(n)]
|
|
|
|
|
|
def test_short_document_is_single_chunk() -> None:
|
|
doc = "# Title\n\nJust some intro, no section headings at all."
|
|
chunks = chunk_markdown(doc)
|
|
assert chunks == [doc.strip()]
|
|
|
|
|
|
def test_empty_and_whitespace_only_content() -> None:
|
|
assert chunk_markdown("") == []
|
|
assert chunk_markdown(" \n\n \n") == []
|
|
|
|
|
|
def test_invalid_params_raise() -> None:
|
|
with pytest.raises(ValueError):
|
|
chunk_markdown("# x", target_chars=0)
|
|
with pytest.raises(ValueError):
|
|
chunk_markdown("# x", overlap_chars=-1)
|
|
|
|
|
|
def test_splits_on_headings_and_keeps_nearest_heading() -> None:
|
|
doc = (
|
|
"# Title\n"
|
|
"intro line\n"
|
|
"## Alpha\n"
|
|
"alpha body\n"
|
|
"### Beta\n"
|
|
"beta body\n"
|
|
"## Gamma\n"
|
|
"gamma body\n"
|
|
)
|
|
chunks = chunk_markdown(doc)
|
|
assert chunks[0] == "# Title\nintro line"
|
|
assert chunks[1] == "## Alpha\nalpha body"
|
|
assert chunks[2] == "### Beta\nbeta body"
|
|
assert chunks[3] == "## Gamma\ngamma body"
|
|
|
|
|
|
def test_document_without_h1_starts_at_first_section() -> None:
|
|
chunks = chunk_markdown("## Only\n\nbody")
|
|
assert chunks == ["## Only\n\nbody"]
|
|
|
|
|
|
def test_long_section_splits_with_overlap_and_anchor_on_every_chunk() -> None:
|
|
body = "\n\n".join(_paras(10, width=138))
|
|
doc = f"{ANCHOR}\n\n{body}"
|
|
chunks = chunk_markdown(doc, target_chars=800, overlap_chars=100)
|
|
|
|
assert len(chunks) == 3
|
|
# Every chunk keeps its nearest preceding heading (the section anchor).
|
|
assert all(c.startswith(ANCHOR) for c in chunks)
|
|
# All chunks respect the target budget (anchor + packed body).
|
|
assert all(len(c) <= 800 for c in chunks)
|
|
# Overlap: the tail of each chunk is at the start of the next one.
|
|
for prev, nxt in pairwise(chunks):
|
|
assert nxt[len(ANCHOR_PREFIX) :].startswith(prev[-100:])
|
|
|
|
|
|
def test_overlap_zero_disables_tail_carryover() -> None:
|
|
body = "\n\n".join(_paras(8, width=200))
|
|
chunks = chunk_markdown(f"{ANCHOR}\n\n{body}", target_chars=800, overlap_chars=0)
|
|
assert len(chunks) >= 2
|
|
for prev, nxt in pairwise(chunks):
|
|
assert not nxt[len(ANCHOR_PREFIX) :].startswith(prev[-50:])
|
|
|
|
|
|
def test_code_fences_stay_intact() -> None:
|
|
doc = (
|
|
"## Section\n"
|
|
"before fence\n"
|
|
"```\n"
|
|
"## fake heading inside fence\n"
|
|
"\n"
|
|
"still in fence\n"
|
|
"```\n"
|
|
"after fence\n"
|
|
"## Other\n"
|
|
"other body\n"
|
|
)
|
|
chunks = chunk_markdown(doc)
|
|
assert any(c.startswith("## Other") for c in chunks)
|
|
# The fake heading inside the fence never opens a section…
|
|
assert not any(c.startswith("## fake heading") for c in chunks)
|
|
# …and the fence itself is whole in the chunk that contains it.
|
|
fenced = [c for c in chunks if "still in fence" in c]
|
|
assert len(fenced) == 1
|
|
assert "## fake heading inside fence" in fenced[0]
|
|
assert fenced[0].count("```") == 2
|
|
# Blank lines inside the fence did not create extra paragraph chunks.
|
|
assert not any(c.startswith("before fence\n\n") for c in chunks)
|
|
|
|
|
|
def test_fence_block_is_atomic_across_forced_split() -> None:
|
|
fence = "```\n" + "\n".join(f"code line {i}" for i in range(60)) + "\n```"
|
|
doc = (
|
|
f"{ANCHOR}\n\npara A "
|
|
+ "a" * 300
|
|
+ f"\n\n{fence}\n\npara B "
|
|
+ "b" * 300
|
|
+ "\n\npara C "
|
|
+ "c" * 300
|
|
)
|
|
chunks = chunk_markdown(doc, target_chars=1000, overlap_chars=100)
|
|
assert len(chunks) >= 2
|
|
# The whole fence (first and last code line) lives in one chunk — a
|
|
# chunk boundary never falls inside a code block.
|
|
assert any("code line 0" in c and "code line 59" in c for c in chunks)
|
|
|
|
|
|
def test_oversized_fence_block_is_split_to_stay_under_hard_cap() -> None:
|
|
"""aipi's embedding endpoint caps requests at ~1024 input tokens — a
|
|
multi-KB fenced code block must not survive chunking as one piece."""
|
|
code = "\n".join(f"int value_{i:03d} = {i}; // padding to grow the line" for i in range(160))
|
|
doc = (
|
|
"# Big Doc\n\n"
|
|
"## Usage Example\n\n"
|
|
f"```cpp\n{code}\n```\n\n"
|
|
"## After\n\nDone.\n"
|
|
)
|
|
chunks = chunk_markdown(doc)
|
|
assert len(chunks) >= 3
|
|
# No chunk exceeds the hard cap (heading anchor adds a little).
|
|
assert all(len(c) <= HARD_MAX_CHARS + 60 for c in chunks)
|
|
# Content survives the split, and later sections are untouched.
|
|
joined = "\n".join(chunks)
|
|
assert "value_000" in joined
|
|
assert "value_159" in joined
|
|
assert any(c.startswith("## After") for c in chunks)
|
|
|
|
|
|
def test_unclosed_fence_does_not_break_sections() -> None:
|
|
doc = "## A\n\n```\nunterminated fence\n\n## B\n\nbody\n"
|
|
chunks = chunk_markdown(doc)
|
|
# "## B" is inside the unterminated fence → not a real heading.
|
|
assert len(chunks) == 1
|
|
assert "## B" in chunks[0]
|
|
|
|
|
|
def test_extract_title_prefers_h1() -> None:
|
|
assert extract_title("# My Title\n\nbody") == "My Title"
|
|
assert extract_title(" # Indented H1\nbody") == "" # ATX must be at col 0
|
|
assert extract_title("## not a title\n\nbody") == ""
|
|
assert extract_title("## sub only", fallback="stem") == "stem"
|
|
assert extract_title("", fallback="fallback") == "fallback"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Format dispatcher (chunk_document) — A9 multi-format ingestion
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_dispatch_by_lowercased_suffix() -> None:
|
|
md = "# T\n\n## A\n\nbody\n"
|
|
assert chunk_document(md, "notes/Doc.MD") == chunk_markdown(md)
|
|
assert chunk_document(md, "notes/doc.MARKDOWN") == chunk_markdown(md)
|
|
assert chunk_document("p1\n\np2\n", "x.TXT") == chunk_text("p1\n\np2\n")
|
|
assert chunk_document("a: 1\n", "x.YAML") == chunk_yaml("a: 1\n")
|
|
assert chunk_document("a: 1\n", "x.Yml") == chunk_yaml("a: 1\n")
|
|
assert chunk_document('{"a": 1}', "x.Json") == chunk_json('{"a": 1}')
|
|
assert chunk_document("def f(): pass\n", "x.PY") == chunk_python("def f(): pass\n")
|
|
|
|
|
|
def test_dispatch_unknown_suffix_falls_back_to_paragraphs() -> None:
|
|
assert chunk_document("hello\n\nworld", "data.csv") == ["hello\nworld"]
|
|
|
|
|
|
def test_dispatch_ignores_directory_part_of_path() -> None:
|
|
assert chunk_document("def f(): pass\n", "a/b/c/script.py") == chunk_python("def f(): pass\n")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# yaml / yml
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_yaml_splits_on_top_level_keys_and_keeps_key_anchors() -> None:
|
|
doc = (
|
|
"# leading comment\n"
|
|
"services:\n"
|
|
" gitlab:\n"
|
|
" image: gitlab/gitlab-ce\n"
|
|
" prometheus:\n"
|
|
" image: prom/prometheus\n"
|
|
"volumes:\n"
|
|
" gitlab-data:\n"
|
|
)
|
|
chunks = chunk_yaml(doc)
|
|
joined = "\n".join(chunks)
|
|
for key in ("services:", "volumes:"):
|
|
assert key in joined
|
|
# Indented keys are NOT block starts — they stay inside their parent block.
|
|
assert not any(c.startswith(" gitlab:") for c in chunks)
|
|
# The leading comment stays with the first block (preamble).
|
|
assert chunks[0].startswith("# leading comment")
|
|
assert "gitlab/gitlab-ce" in joined and "prom/prometheus" in joined
|
|
|
|
|
|
def test_yaml_document_separators_start_new_blocks() -> None:
|
|
a = "site_a: " + "a" * 500 + "\n"
|
|
b = "site_b: " + "b" * 500 + "\n"
|
|
chunks = chunk_yaml(a + "---\n" + b, target_chars=600, overlap_chars=0)
|
|
# Each site is long enough to force its own chunk; the separator must not
|
|
# glue them into one over-budget chunk.
|
|
assert len(chunks) >= 2
|
|
assert all(len(c) <= 600 for c in chunks)
|
|
assert not any("site_a" in c and "site_b" in c for c in chunks)
|
|
|
|
|
|
def test_yaml_oversized_key_block_is_split_under_hard_cap() -> None:
|
|
doc = "big_list:\n" + (" - " + "x" * 60 + "\n") * 60 # one ~3800-char block
|
|
chunks = chunk_yaml(doc)
|
|
assert len(chunks) >= 2
|
|
assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
|
|
# Overlap re-prints (≤50 chars per split), so only a little content is
|
|
# re-stated — the bulk of the block must survive.
|
|
assert sum(len(c) for c in chunks) >= len(doc) - 300
|
|
|
|
|
|
def test_yaml_empty_content() -> None:
|
|
assert chunk_yaml("") == []
|
|
assert chunk_yaml("\n\n \n") == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# json
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_json_splits_on_top_level_keys_pretty_printed() -> None:
|
|
doc = '{"hosts": {"kafkabridge": "10.0.3.7"}, "count": 3}'
|
|
chunks = chunk_json(doc, target_chars=45, overlap_chars=0) # force 1 chunk/block
|
|
assert len(chunks) == 2
|
|
first, second = chunks
|
|
assert '"hosts"' in first and "kafkabridge" in first
|
|
assert '"count"' in second
|
|
# Pretty-printed (indent=2), not the compact input form.
|
|
assert '"kafkabridge": "10.0.3.7"' in first
|
|
assert not any('{"hosts"' in c for c in chunks)
|
|
|
|
|
|
def test_json_each_key_block_is_self_contained() -> None:
|
|
doc = '{"a": "x", "b": "y"}'
|
|
chunks = chunk_json(doc, target_chars=13, overlap_chars=0) # force 1 chunk/block
|
|
assert [c for c in chunks if '"a"' in c] and [c for c in chunks if '"b"' in c]
|
|
assert not any('"a"' in c and '"b"' in c for c in chunks)
|
|
|
|
|
|
def test_json_oversized_value_falls_under_hard_cap() -> None:
|
|
doc = '{"blob": "' + "z" * 4000 + '"}'
|
|
chunks = chunk_json(doc)
|
|
assert len(chunks) >= 2
|
|
assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
|
|
assert "".join(chunks).count("z") >= 4000
|
|
|
|
|
|
def test_json_top_level_list_is_one_pretty_block() -> None:
|
|
chunks = chunk_json("[1, 2, 3]")
|
|
assert chunks == ["[\n 1,\n 2,\n 3\n]"]
|
|
|
|
|
|
def test_json_unparseable_falls_back_to_paragraph_packing() -> None:
|
|
doc = "{broken json\n\nsecond paragraph here\n"
|
|
assert chunk_json(doc) == chunk_text(doc)
|
|
assert chunk_json("not json at all") == chunk_text("not json at all")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# python
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_python_splits_on_top_level_defs_and_classes() -> None:
|
|
doc = (
|
|
'"""Module doc."""\n'
|
|
"import asyncio\n"
|
|
"\n"
|
|
"CONST = 1\n"
|
|
"\n"
|
|
"def alpha():\n"
|
|
" return 1\n"
|
|
"\n"
|
|
"class Beta:\n"
|
|
" def run(self):\n"
|
|
" return 2\n"
|
|
)
|
|
chunks = chunk_python(doc, target_chars=60, overlap_chars=0) # force 1 chunk/block
|
|
assert len(chunks) == 3
|
|
assert chunks[0].startswith('"""Module doc."""')
|
|
assert "CONST = 1" in chunks[0] # preamble ends at the first def/class
|
|
assert chunks[1].startswith("def alpha")
|
|
assert chunks[2].startswith("class Beta")
|
|
assert "def run" in chunks[2] # nested def stays inside the class block
|
|
|
|
|
|
def test_python_decorators_stay_with_their_definition() -> None:
|
|
doc = "@app.get('/x')\ndef handler():\n return 'x'\n"
|
|
chunks = chunk_python(doc)
|
|
assert chunks[0].startswith("@app.get")
|
|
|
|
|
|
def test_python_oversized_function_falls_back_to_line_packing() -> None:
|
|
doc = "def big():\n" + "\n".join(f" val_{i:03d} = {i} # padding" for i in range(80))
|
|
chunks = chunk_python(doc)
|
|
assert len(chunks) >= 2
|
|
assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
|
|
assert "val_000" in chunks[0]
|
|
assert "val_079" in chunks[-1]
|
|
assert sum(len(c) for c in chunks) >= len(doc) - 100
|
|
|
|
|
|
def test_python_unparseable_source_falls_back_to_paragraphs() -> None:
|
|
src = "def broken(:\n\nstill text\n"
|
|
assert chunk_python(src) == chunk_text(src)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# txt
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_txt_paragraph_packing() -> None:
|
|
doc = "alpha\n\nbeta\n\ngamma\n"
|
|
chunks = chunk_text(doc)
|
|
assert chunks == ["alpha\nbeta\ngamma"] # all three fit the target
|
|
|
|
|
|
def test_txt_long_doc_packs_with_overlap() -> None:
|
|
doc = "\n\n".join(f"para {i} " + "l" * 300 for i in range(6))
|
|
chunks = chunk_text(doc, target_chars=800, overlap_chars=100)
|
|
assert len(chunks) >= 2
|
|
assert all(len(c) <= 800 for c in chunks)
|
|
assert all(f"para {i}" in "\n".join(chunks) for i in range(6))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Hard cap across every format (aipi ~1024-token request cap)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("content", "path"),
|
|
[
|
|
("# T\n\n" + "word " * 1200, "big.md"),
|
|
("key: " + "v" * 5000 + "\n", "big.yaml"),
|
|
('{"blob": "' + "z" * 5000 + '"}', "big.json"),
|
|
("def f():\n" + " x = 1\n" * 1000, "big.py"),
|
|
("line of text\n\n" * 800, "big.txt"),
|
|
],
|
|
)
|
|
def test_hard_cap_holds_for_every_format(content: str, path: str) -> None:
|
|
chunks = chunk_document(content, path)
|
|
assert chunks, "expected at least one chunk"
|
|
for c in chunks:
|
|
assert len(c) <= HARD_MAX_CHARS, f"{path}: {len(c)} chars"
|