wehub-resource-sync
1352 lines
50 KiB
Python
1352 lines
50 KiB
Python
"""Regression tests for paragraph-semantic TableRowSplit oversized-table handling."""
|
|
|
|
import json
|
|
import logging
|
|
|
|
import pytest
|
|
|
|
from lightrag.chunker.paragraph_semantic import (
|
|
_detect_table_format,
|
|
_expand_block_with_table_splits,
|
|
_merge_small_blocks,
|
|
_new_block,
|
|
_split_html_rows,
|
|
_split_long_block,
|
|
_split_rows_by_tokens,
|
|
_split_table_text,
|
|
chunking_by_paragraph_semantic,
|
|
)
|
|
from lightrag.utils import Tokenizer, TokenizerInterface
|
|
|
|
|
|
class _CharTokenizer(TokenizerInterface):
|
|
"""1:1 character-to-token mapping — keeps math obvious in assertions."""
|
|
|
|
def encode(self, content: str):
|
|
return [ord(ch) for ch in content]
|
|
|
|
def decode(self, tokens):
|
|
return "".join(chr(t) for t in tokens)
|
|
|
|
|
|
def _make_tokenizer() -> Tokenizer:
|
|
return Tokenizer(model_name="char", tokenizer=_CharTokenizer())
|
|
|
|
|
|
@pytest.fixture
|
|
def _propagate_lightrag_logger(monkeypatch):
|
|
# The ``lightrag`` logger sets ``propagate=False`` (see lightrag.utils), so
|
|
# caplog's root handler captures nothing by default. Flip it on for the
|
|
# duration of a test that asserts on the degrade warning.
|
|
monkeypatch.setattr(logging.getLogger("lightrag"), "propagate", True)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_rows_by_tokens_few_rows_huge_total_no_empty_slice():
|
|
# Reproduces the bug where target_chunks > len(rows) made target_rows
|
|
# < 1, so int((i+1)*target_rows) collapsed to start and the loop
|
|
# appended empty slices (which would later serialise as <table>[]…).
|
|
tokenizer = _make_tokenizer()
|
|
# 3 rows that each individually exceed target_max — forces
|
|
# math.ceil(total/target_ideal) and math.ceil(total/target_max) to
|
|
# both be much greater than len(rows).
|
|
rows = [
|
|
[{"col": "x" * 800}],
|
|
[{"col": "y" * 800}],
|
|
[{"col": "z" * 800}],
|
|
]
|
|
|
|
chunks = _split_rows_by_tokens(
|
|
rows,
|
|
tokenizer,
|
|
target_max=200,
|
|
target_ideal=150,
|
|
last_min=64,
|
|
)
|
|
|
|
assert chunks, "expected at least one chunk"
|
|
for chunk in chunks:
|
|
assert chunk, "TableRowSplit must never emit an empty row slice"
|
|
# Concatenation preserves all rows in order.
|
|
flat: list = []
|
|
for chunk in chunks:
|
|
flat.extend(chunk)
|
|
assert flat == rows
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_rows_by_tokens_balanced_split_yields_one_row_per_chunk():
|
|
# When target_chunks gets capped at len(rows), each chunk holds one
|
|
# row — verifies the cap kicks in and forward progress is preserved.
|
|
tokenizer = _make_tokenizer()
|
|
rows = [[{"col": "a" * 300}] for _ in range(4)]
|
|
|
|
chunks = _split_rows_by_tokens(
|
|
rows,
|
|
tokenizer,
|
|
target_max=200,
|
|
target_ideal=150,
|
|
last_min=10, # low enough that the tail-merge step doesn't fire
|
|
)
|
|
|
|
assert all(chunk for chunk in chunks)
|
|
# Each row appears exactly once across the chunks.
|
|
flat: list = []
|
|
for chunk in chunks:
|
|
flat.extend(chunk)
|
|
assert flat == rows
|
|
|
|
|
|
def _build_oversized_table_text(num_rows: int, row_payload_size: int) -> str:
|
|
rows = [[f"r{idx}-" + "x" * row_payload_size] for idx in range(num_rows)]
|
|
return f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
|
|
def _write_blocks_jsonl(tmp_path, content: str) -> str:
|
|
path = tmp_path / "doc.blocks.jsonl"
|
|
row = {
|
|
"type": "content",
|
|
"heading": "Section",
|
|
"parent_headings": [],
|
|
"level": 2,
|
|
"content": content,
|
|
}
|
|
path.write_text(json.dumps(row, ensure_ascii=False), encoding="utf-8")
|
|
return str(path)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_expand_block_assigns_first_and_last_roles_to_glued_blocks():
|
|
# An oversized table sandwiched between leading and trailing paragraphs
|
|
# produces three slices: "first" (glued with leading paras),
|
|
# "middle" (standalone), "last" (glued with trailing paras). Before
|
|
# the fix, the first/last blocks defaulted to "none" and lost their
|
|
# directional merge-protection.
|
|
tokenizer = _make_tokenizer()
|
|
table_text = _build_oversized_table_text(num_rows=6, row_payload_size=200)
|
|
block = {
|
|
"heading": "Section",
|
|
"parent_headings": ["Doc"],
|
|
"level": 2,
|
|
"paragraphs": [
|
|
{"text": "lead paragraph", "is_table": False},
|
|
{"text": table_text, "is_table": True},
|
|
{"text": "trailing paragraph", "is_table": False},
|
|
],
|
|
}
|
|
|
|
out = _expand_block_with_table_splits(
|
|
block,
|
|
tokenizer=tokenizer,
|
|
table_max=400,
|
|
table_ideal=300,
|
|
table_min_last=128,
|
|
)
|
|
|
|
roles = [b["table_chunk_role"] for b in out]
|
|
assert roles[0] == "first", f"expected leading block role=first, got {roles}"
|
|
assert roles[-1] == "last", f"expected trailing block role=last, got {roles}"
|
|
assert all(r == "middle" for r in roles[1:-1]), (
|
|
f"expected middle slices between first/last, got {roles}"
|
|
)
|
|
|
|
# Boundary glue still works: leading text sits inside the first block,
|
|
# trailing text sits inside the last block.
|
|
assert any(p["text"] == "lead paragraph" for p in out[0]["paragraphs"]), (
|
|
"leading paragraph must glue with the first table slice"
|
|
)
|
|
assert any(p["text"] == "trailing paragraph" for p in out[-1]["paragraphs"]), (
|
|
"trailing paragraph must glue with the last table slice"
|
|
)
|
|
assert all("表格片段" not in b["heading"] for b in out), (
|
|
"TableRowSplit should not expose legacy table-fragment heading suffixes"
|
|
)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_expand_block_two_oversized_tables_separates_last_and_first_roles():
|
|
# Two oversized tables in the same heading block: the tail of the first
|
|
# split must carry role="last" and not be silently merged into the
|
|
# head of the second split (which must carry role="first").
|
|
tokenizer = _make_tokenizer()
|
|
block = {
|
|
"heading": "Section",
|
|
"parent_headings": [],
|
|
"level": 2,
|
|
"paragraphs": [
|
|
{
|
|
"text": _build_oversized_table_text(num_rows=4, row_payload_size=200),
|
|
"is_table": True,
|
|
},
|
|
{"text": "between tables", "is_table": False},
|
|
{
|
|
"text": _build_oversized_table_text(num_rows=4, row_payload_size=200),
|
|
"is_table": True,
|
|
},
|
|
],
|
|
}
|
|
|
|
out = _expand_block_with_table_splits(
|
|
block,
|
|
tokenizer=tokenizer,
|
|
table_max=400,
|
|
table_ideal=300,
|
|
table_min_last=128,
|
|
)
|
|
|
|
roles = [b["table_chunk_role"] for b in out]
|
|
# We expect the role sequence to start with "first", end with "last",
|
|
# and contain at least one "last" -> "first" transition (the boundary
|
|
# between the two oversized tables) without any boundary block losing
|
|
# its role.
|
|
assert roles[0] == "first"
|
|
assert roles[-1] == "last"
|
|
assert "last" in roles
|
|
# The transition: there must be a "last" immediately followed by a
|
|
# "first" somewhere in the middle of the role sequence.
|
|
transitions = list(zip(roles, roles[1:]))
|
|
assert ("last", "first") in transitions, (
|
|
f"expected a last->first boundary between the two split tables, got {roles}"
|
|
)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_expand_block_duplicates_short_text_between_oversized_tables():
|
|
tokenizer = _make_tokenizer()
|
|
bridge = "between tables"
|
|
block = {
|
|
"heading": "Section",
|
|
"parent_headings": [],
|
|
"level": 2,
|
|
"paragraphs": [
|
|
{
|
|
"text": _build_oversized_table_text(num_rows=4, row_payload_size=200),
|
|
"is_table": True,
|
|
},
|
|
{"text": bridge, "is_table": False},
|
|
{
|
|
"text": _build_oversized_table_text(num_rows=4, row_payload_size=200),
|
|
"is_table": True,
|
|
},
|
|
],
|
|
}
|
|
|
|
out = _expand_block_with_table_splits(
|
|
block,
|
|
tokenizer=tokenizer,
|
|
table_max=400,
|
|
table_ideal=300,
|
|
table_min_last=128,
|
|
target_max=800,
|
|
chunk_overlap_token_size=100,
|
|
)
|
|
|
|
roles = [b["table_chunk_role"] for b in out]
|
|
boundary_idx = next(
|
|
i
|
|
for i, (left, right) in enumerate(zip(roles, roles[1:]))
|
|
if (left, right) == ("last", "first")
|
|
)
|
|
assert bridge in out[boundary_idx]["content"]
|
|
assert bridge in out[boundary_idx + 1]["content"]
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_expand_block_emits_middle_text_when_table_bridge_is_long():
|
|
tokenizer = _make_tokenizer()
|
|
bridge = ("A" * 45) + ("B" * 50) + ("C" * 45)
|
|
block = {
|
|
"heading": "Section",
|
|
"parent_headings": [],
|
|
"level": 2,
|
|
"paragraphs": [
|
|
{
|
|
"text": _build_oversized_table_text(num_rows=6, row_payload_size=120),
|
|
"is_table": True,
|
|
},
|
|
{"text": bridge, "is_table": False},
|
|
{
|
|
"text": _build_oversized_table_text(num_rows=6, row_payload_size=120),
|
|
"is_table": True,
|
|
},
|
|
],
|
|
}
|
|
|
|
out = _expand_block_with_table_splits(
|
|
block,
|
|
tokenizer=tokenizer,
|
|
table_max=260,
|
|
table_ideal=180,
|
|
table_min_last=32,
|
|
target_max=400,
|
|
chunk_overlap_token_size=45,
|
|
)
|
|
|
|
# The standalone middle block carries R-style overlap with the text that
|
|
# went left (prefix) and right (suffix): because each side's slice is
|
|
# itself ≤ the overlap budget, the middle re-covers the whole bridge — the
|
|
# bridge is never fragmented, only its head/tail are *also* duplicated into
|
|
# the neighbouring table blocks.
|
|
middle_idx = next(
|
|
i
|
|
for i, blk in enumerate(out)
|
|
if blk["table_chunk_role"] == "none" and blk["content"] == bridge
|
|
)
|
|
assert out[middle_idx - 1]["table_chunk_role"] == "last"
|
|
assert "A" * 45 in out[middle_idx - 1]["content"]
|
|
assert "B" * 50 not in out[middle_idx - 1]["content"]
|
|
assert out[middle_idx + 1]["table_chunk_role"] == "first"
|
|
assert out[middle_idx + 1]["content"].startswith("C" * 45)
|
|
assert "B" * 50 not in out[middle_idx + 1]["content"]
|
|
# The overlap never drags table markup into the middle text block.
|
|
assert "<table" not in out[middle_idx]["content"]
|
|
assert all(b["tokens"] <= 400 for b in out), [b["tokens"] for b in out]
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_bridge_single_side_budget_capped_at_half_target_max():
|
|
# §7 guarantee: even with a huge configured overlap, the bridge text
|
|
# duplicated into each table boundary block is capped at target_max // 2
|
|
# so it can never dominate the block. The bridge is a long run of a
|
|
# character that never appears in the table payload, so counting it in
|
|
# each boundary block measures exactly how much bridge text was copied in.
|
|
tokenizer = _make_tokenizer()
|
|
bridge = "Z" * 1000
|
|
block = {
|
|
"heading": "Section",
|
|
"parent_headings": [],
|
|
"level": 2,
|
|
"paragraphs": [
|
|
{
|
|
"text": _build_oversized_table_text(num_rows=8, row_payload_size=30),
|
|
"is_table": True,
|
|
},
|
|
{"text": bridge, "is_table": False},
|
|
{
|
|
"text": _build_oversized_table_text(num_rows=8, row_payload_size=30),
|
|
"is_table": True,
|
|
},
|
|
],
|
|
}
|
|
target_max = 400
|
|
half = target_max // 2
|
|
|
|
out = _expand_block_with_table_splits(
|
|
block,
|
|
tokenizer=tokenizer,
|
|
table_max=200,
|
|
table_ideal=140,
|
|
table_min_last=48,
|
|
target_max=target_max,
|
|
chunk_overlap_token_size=10_000, # huge → only the half-cap can bind
|
|
)
|
|
|
|
# The long bridge yields [... "last"+prefix, "none" full-bridge middle,
|
|
# suffix+"first" ...]; locate the middle and read its neighbours.
|
|
middle_idx = next(
|
|
i
|
|
for i, blk in enumerate(out)
|
|
if blk["table_chunk_role"] == "none" and blk["content"] == bridge
|
|
)
|
|
left = out[middle_idx - 1]
|
|
right = out[middle_idx + 1]
|
|
assert left["table_chunk_role"] == "last"
|
|
assert right["table_chunk_role"] == "first"
|
|
left_z = left["content"].count("Z")
|
|
right_z = right["content"].count("Z")
|
|
# The huge overlap would have copied far more than `half` without the cap;
|
|
# the bridge (1000) and the small table slices leave ample headroom, so the
|
|
# half-cap is the binding constraint and each side gets exactly `half`.
|
|
assert left_z == half, left_z
|
|
assert right_z == half, right_z
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_public_chunking_adds_part_suffixes_to_all_table_split_fragments(tmp_path):
|
|
tokenizer = _make_tokenizer()
|
|
body = "\n".join(
|
|
[
|
|
"lead paragraph",
|
|
_build_oversized_table_text(num_rows=6, row_payload_size=200),
|
|
"trailing paragraph",
|
|
]
|
|
)
|
|
blocks_path = _write_blocks_jsonl(tmp_path, body)
|
|
|
|
chunks = chunking_by_paragraph_semantic(
|
|
tokenizer,
|
|
body,
|
|
chunk_token_size=800,
|
|
blocks_path=blocks_path,
|
|
chunk_overlap_token_size=0,
|
|
)
|
|
|
|
assert len(chunks) > 1
|
|
assert [chunk["heading"]["heading"] for chunk in chunks] == [
|
|
f"Section [part {idx}]" for idx in range(1, len(chunks) + 1)
|
|
]
|
|
assert all("表格片段" not in chunk["heading"]["heading"] for chunk in chunks)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Table-aware fallback tests (row-boundary first, character last).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_detect_table_format_explicit_attr():
|
|
assert _detect_table_format('id="t1" format="json"', "[]") == "json"
|
|
assert _detect_table_format("format='html'", "<tr></tr>") == "html"
|
|
# Unknown formats fall through (force the caller to use char fallback).
|
|
assert _detect_table_format('format="markdown"', "...") is None
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_detect_table_format_sniff_when_attrs_silent():
|
|
assert _detect_table_format("", '[{"a":1}]') == "json"
|
|
assert _detect_table_format("", "<tr><td>x</td></tr>") == "html"
|
|
# Body that doesn't look like JSON or HTML → unknown.
|
|
assert _detect_table_format("", "plain text rows") is None
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_html_rows_extracts_tr_elements():
|
|
body = (
|
|
"<thead><tr><th>h</th></tr></thead>"
|
|
"<tbody><tr><td>a</td></tr><tr><td>b</td></tr></tbody>"
|
|
)
|
|
rows = _split_html_rows(body)
|
|
assert rows is not None
|
|
assert len(rows) == 3
|
|
# Each row carries its parent wrapper so the chunk serialiser can
|
|
# rebuild <thead>/<tbody> instead of dropping them silently.
|
|
assert [w for w, _ in rows] == ["thead", "tbody", "tbody"]
|
|
assert all(tr.startswith("<tr") and tr.endswith("</tr>") for _, tr in rows)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_html_rows_no_tr_returns_none():
|
|
assert _split_html_rows("just text, no rows") is None
|
|
assert _split_html_rows("") is None
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_single_row_oversized_falls_to_character_split():
|
|
# A 1-row table whose single cell is huge cannot be reduced via row
|
|
# boundary — the function must fall to character splitting and respect
|
|
# target_max on every output piece.
|
|
tokenizer = _make_tokenizer()
|
|
rows = [[{"col": "x" * 2000}]]
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=500,
|
|
target_ideal=350,
|
|
last_min=128,
|
|
)
|
|
|
|
assert len(pieces) >= 2, "single-row oversized table must produce multiple pieces"
|
|
# Every piece honors the cap (this is the contract violation the user
|
|
# reported when the previous code emitted a single 2000-token table).
|
|
assert all(_count_tokens(tokenizer, p) <= 500 for p in pieces)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_multirow_one_huge_row_degrades_whole_table_to_recursive():
|
|
# A multi-row table where most rows fit but one row is itself huge. A single
|
|
# row that exceeds the cap can no longer be expressed at row boundaries, so
|
|
# the WHOLE table degrades to a recursive character split (rather than
|
|
# emitting a mix of <table> slices and orphaned character fragments).
|
|
tokenizer = _make_tokenizer()
|
|
small_row = [{"col": "ok"}]
|
|
huge_row = [{"col": "z" * 2000}]
|
|
rows = [small_row, huge_row, small_row]
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=500,
|
|
target_ideal=350,
|
|
last_min=64,
|
|
)
|
|
|
|
assert len(pieces) >= 2, "oversized row must force a multi-piece recursive split"
|
|
assert all(_count_tokens(tokenizer, p) <= 500 for p in pieces)
|
|
# The whole table degraded: it was shredded as one recursive split, so no
|
|
# piece survives as a complete re-wrapped <table>...</table> slice (the old
|
|
# mixed output kept legal per-row table slices alongside text fragments).
|
|
assert not any(p.startswith("<table ") and p.endswith("</table>") for p in pieces)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_html_table_split_by_tr():
|
|
# HTML-format table: rows are <tr>...</tr>; each output fragment must
|
|
# remain a legal <table {attrs}>{rows}</table> string.
|
|
tokenizer = _make_tokenizer()
|
|
body = "".join(f"<tr><td>{'r' * 200}</td></tr>" for _ in range(5))
|
|
table_text = f'<table id="tb-h1" format="html">{body}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=500,
|
|
target_ideal=350,
|
|
last_min=64,
|
|
)
|
|
|
|
assert len(pieces) >= 2
|
|
# All pieces should be legal <table>...</table> fragments (none of the
|
|
# rows individually exceeds target_max, so no character fallback).
|
|
assert all(p.startswith("<table ") and p.endswith("</table>") for p in pieces)
|
|
assert all(_count_tokens(tokenizer, p) <= 500 for p in pieces)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_html_preserves_thead_tbody_wrappers():
|
|
# When an HTML table mixes <thead> and <tbody>, the row splitter
|
|
# used to drop the wrappers entirely — the chunked output came back
|
|
# as bare <tr> sequences. The fix re-emits each wrapper around its
|
|
# rows in every chunk so the table structure survives splitting.
|
|
tokenizer = _make_tokenizer()
|
|
head_row = "<tr><th>" + ("h" * 80) + "</th></tr>"
|
|
body_rows = "".join(f"<tr><td>{'b' * 80}{i}</td></tr>" for i in range(4))
|
|
body = f"<thead>{head_row}</thead><tbody>{body_rows}</tbody>"
|
|
table_text = f'<table id="tb-mixed" format="html">{body}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=400,
|
|
target_ideal=280,
|
|
last_min=64,
|
|
)
|
|
|
|
# Multiple chunks expected and every chunk must remain a legal
|
|
# <table>-wrapped fragment.
|
|
assert len(pieces) >= 2
|
|
assert all(p.startswith("<table ") and p.endswith("</table>") for p in pieces)
|
|
# Every chunk that contains the header row must still wrap it in
|
|
# <thead>...</thead>; every chunk with body rows must wrap them in
|
|
# <tbody>...</tbody>. Before the fix, both wrappers vanished.
|
|
for piece in pieces:
|
|
if "<th>" in piece:
|
|
assert "<thead>" in piece and "</thead>" in piece, piece
|
|
if "<td>" in piece:
|
|
assert "<tbody>" in piece and "</tbody>" in piece, piece
|
|
# Round-trip: concatenating just the row payloads from every chunk
|
|
# recovers the original row sequence in order.
|
|
extracted_rows: list[str] = []
|
|
import re
|
|
|
|
for piece in pieces:
|
|
extracted_rows.extend(
|
|
re.findall(r"<tr\b[^>]*>.*?</tr>", piece, re.DOTALL | re.IGNORECASE)
|
|
)
|
|
expected_rows = re.findall(r"<tr\b[^>]*>.*?</tr>", body, re.DOTALL | re.IGNORECASE)
|
|
assert extracted_rows == expected_rows
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_html_injects_thead_into_nonfirst_slices():
|
|
"""A long HTML table with an HTML <thead> header_body: the first slice keeps
|
|
its own in-body <thead> (not double-injected); every middle/last slice has
|
|
the raw <thead> spliced in exactly once, with rowspan/colspan preserved;
|
|
every slice stays within the cap."""
|
|
tokenizer = _make_tokenizer()
|
|
header = (
|
|
'<thead><tr><th rowspan="2">Metric</th><th colspan="2">Group</th></tr></thead>'
|
|
)
|
|
head_row = '<tr><th rowspan="2">Metric</th><th colspan="2">Group</th></tr>'
|
|
body_rows = "".join(f"<tr><td>{'d' * 120}{i}</td></tr>" for i in range(6))
|
|
# Body carries its own <thead> (the source header) plus the data rows.
|
|
body = f"<thead>{head_row}</thead><tbody>{body_rows}</tbody>"
|
|
table_text = f'<table id="tb-h" format="html">{body}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=500,
|
|
target_ideal=350,
|
|
last_min=64,
|
|
header_body=header,
|
|
)
|
|
|
|
assert len(pieces) >= 3 # first / middle(s) / last
|
|
assert all(p.startswith("<table ") and p.endswith("</table>") for p in pieces)
|
|
assert all(_count_tokens(tokenizer, p) <= 500 for p in pieces)
|
|
|
|
# First slice keeps its own header and is NOT injected with a second copy.
|
|
assert pieces[0].count("<thead>") == 1
|
|
# Every non-first slice gets the header spliced in exactly once, verbatim
|
|
# (rowspan/colspan preserved), immediately after the opening <table …>.
|
|
for piece in pieces[1:]:
|
|
assert piece.count("<thead>") == 1, piece
|
|
assert header in piece, piece
|
|
assert 'rowspan="2"' in piece and 'colspan="2"' in piece, piece
|
|
# The <thead> sits at the front of the body, ahead of the data rows.
|
|
assert piece.index("<thead>") < piece.index("<tbody>"), piece
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_raises_on_html_table_with_json_header():
|
|
"""Cross-format guard: an HTML table fed a JSON-array header_body must raise
|
|
(the sidecar header does not belong to this table)."""
|
|
tokenizer = _make_tokenizer()
|
|
body = "".join(f"<tr><td>{'r' * 200}</td></tr>" for _ in range(5))
|
|
table_text = f'<table id="tb-h1" format="html">{body}</table>'
|
|
with pytest.raises(ValueError):
|
|
_split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=500,
|
|
target_ideal=350,
|
|
last_min=64,
|
|
header_body='[["H1", "H2"]]',
|
|
)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_unknown_format_falls_to_character():
|
|
# No format attr, body that doesn't look like JSON/HTML → unknown.
|
|
tokenizer = _make_tokenizer()
|
|
table_text = '<table id="weird">' + ("plain row text " * 300) + "</table>"
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=500,
|
|
target_ideal=350,
|
|
last_min=64,
|
|
)
|
|
|
|
assert len(pieces) >= 2
|
|
assert all(_count_tokens(tokenizer, p) <= 500 for p in pieces)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_expand_block_single_row_table_no_longer_left_intact():
|
|
# TableRowSplit integration: previously a single-row oversized table was
|
|
# appended back to cur_paras unchanged, leading the block to reach
|
|
# AnchorSplit with the table whole and the character fallback shredding
|
|
# the <table> tag. After the fix, TableRowSplit itself produces multiple
|
|
# pieces for such a table.
|
|
tokenizer = _make_tokenizer()
|
|
rows = [[{"col": "x" * 2000}]] # single huge row
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
block = {
|
|
"heading": "Section",
|
|
"parent_headings": [],
|
|
"level": 2,
|
|
"paragraphs": [
|
|
{"text": "lead", "is_table": False},
|
|
{"text": table_text, "is_table": True},
|
|
{"text": "trail", "is_table": False},
|
|
],
|
|
}
|
|
|
|
out = _expand_block_with_table_splits(
|
|
block,
|
|
tokenizer=tokenizer,
|
|
table_max=400,
|
|
table_ideal=300,
|
|
table_min_last=128,
|
|
)
|
|
|
|
# Multiple sub-blocks must be produced; the oversized table no longer
|
|
# passes through whole.
|
|
assert len(out) >= 2
|
|
# First/last role protection still fires when the table was reduced.
|
|
roles = [b["table_chunk_role"] for b in out]
|
|
assert "first" in roles or "last" in roles, (
|
|
f"expected first/last role assignment after table split, got {roles}"
|
|
)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_long_block_table_dominant_no_anchor_keeps_some_table_markup():
|
|
# AnchorSplit integration: a block dominated by an oversized table with no
|
|
# anchor candidates used to be character-split end-to-end, destroying
|
|
# the <table> tag. After the fix, at least some output sub-blocks
|
|
# retain legal <table>...</table> markup for the rows that fit.
|
|
tokenizer = _make_tokenizer()
|
|
# Many small rows -> row-boundary split produces multiple legal
|
|
# <table> fragments, none of which individually exceed target_max.
|
|
rows = [[{"col": f"r{i}-" + "v" * 200}] for i in range(8)]
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
paragraphs = [
|
|
{"text": "Sufficiently long lead paragraph " * 30, "is_table": False},
|
|
{"text": table_text, "is_table": True},
|
|
]
|
|
|
|
sub_blocks = _split_long_block(
|
|
paragraphs,
|
|
heading="Heading",
|
|
parent_headings=[],
|
|
level=2,
|
|
table_chunk_role="none",
|
|
tokenizer=tokenizer,
|
|
target_max=600,
|
|
target_ideal=450,
|
|
)
|
|
|
|
# Every sub-block respects the cap.
|
|
assert all(b["tokens"] <= 600 for b in sub_blocks)
|
|
# At least one sub-block keeps an unbroken <table> fragment somewhere
|
|
# in its content (proof that row-boundary preservation kicked in).
|
|
contents = [b["content"] for b in sub_blocks]
|
|
assert any(("<table " in c and "</table>" in c) for c in contents), (
|
|
"expected at least one sub-block to retain a legal <table> fragment"
|
|
)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_budgets_wrapper_overhead_for_target_max():
|
|
# ``_split_rows_by_tokens`` measures only the body (json.dumps(rows));
|
|
# the surrounding ``<table {attrs}></table>`` wrapper costs tokens too.
|
|
# Without wrapper-aware budgeting, a chunk whose body just fits
|
|
# target_max would overflow once wrapped and trigger character
|
|
# fallback — shredding the row structure for no good reason.
|
|
tokenizer = _make_tokenizer()
|
|
# A long attrs string forces a non-trivial wrapper overhead so the
|
|
# body-only budget previously chosen (==target_max) overflows when
|
|
# the wrapper is added back in.
|
|
attrs_padding = "x" * 80
|
|
rows = [[{"col": "y" * 80}] for _ in range(4)]
|
|
table_text = f'<table id="{attrs_padding}" format="json">{json.dumps(rows)}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=250,
|
|
target_ideal=180,
|
|
last_min=64,
|
|
)
|
|
|
|
# Every output piece honors the cap.
|
|
assert all(_count_tokens(tokenizer, p) <= 250 for p in pieces), [
|
|
_count_tokens(tokenizer, p) for p in pieces
|
|
]
|
|
# Row structure preserved — none of the pieces fell back to
|
|
# character fragments because of accidental wrapper overflow.
|
|
assert all(p.startswith("<table ") and p.endswith("</table>") for p in pieces)
|
|
|
|
|
|
def _count_tokens(tokenizer: Tokenizer, text: str) -> int:
|
|
return len(tokenizer.encode(text))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# HeaderRecovery — the source table's repeating header is budgeted out of the
|
|
# per-slice cap and re-injected into the <table> body of every non-first slice
|
|
# AT SPLIT TIME (not backfilled at output), so a slice + its header always
|
|
# stays ≤ target_max, and split slices are frozen against LevelMerge so the
|
|
# header is never duplicated by a re-merge.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_HEADER_BODY = '[["H1", "H2"]]'
|
|
# After injection (or naturally on the first slice) a JSON slice's <table> body
|
|
# begins with the header row.
|
|
_INJECTED_PREFIX = '<table id="tb-1" format="json">[["H1", "H2"]'
|
|
|
|
# Sentinel distinguishing "write tables.json without a table_header" from
|
|
# "do not write tables.json at all".
|
|
_NO_SIDECAR = object()
|
|
|
|
|
|
def _write_tables_json(tmp_path, headers: dict) -> None:
|
|
"""Write a ``doc.tables.json`` beside ``doc.blocks.jsonl``.
|
|
|
|
``headers`` maps table id -> header string; a ``None`` value emits an entry
|
|
WITHOUT the ``table_header`` field (a table that has no repeating header).
|
|
"""
|
|
tables: dict = {}
|
|
for tid, header in headers.items():
|
|
entry = {"id": tid, "format": "json", "content": "[]", "caption": ""}
|
|
if header is not None:
|
|
entry["table_header"] = header
|
|
tables[tid] = entry
|
|
path = tmp_path / "doc.tables.json"
|
|
path.write_text(
|
|
json.dumps({"version": "1.0", "tables": tables}, ensure_ascii=False),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def _build_table_with_header(num_data_rows: int, payload_size: int) -> str:
|
|
# First row is the real header; the rest are data rows. After row-splitting,
|
|
# only the "first" slice keeps the header row; middle/last slices must have
|
|
# it re-injected.
|
|
rows = [["H1", "H2"]] + [
|
|
[f"r{idx}-" + "x" * payload_size, "y"] for idx in range(num_data_rows)
|
|
]
|
|
return f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
|
|
def _chunk_with_oversized_table(
|
|
tmp_path,
|
|
*,
|
|
heading: str = "Section",
|
|
sidecar=_HEADER_BODY,
|
|
num_data_rows: int = 6,
|
|
payload_size: int = 200,
|
|
trailing: str = "trailing paragraph",
|
|
chunk_token_size: int = 800,
|
|
):
|
|
tokenizer = _make_tokenizer()
|
|
body = "\n".join(
|
|
[
|
|
"lead paragraph",
|
|
_build_table_with_header(
|
|
num_data_rows=num_data_rows, payload_size=payload_size
|
|
),
|
|
trailing,
|
|
]
|
|
)
|
|
if sidecar is _NO_SIDECAR:
|
|
pass # leave the directory without a tables.json
|
|
elif sidecar is None:
|
|
_write_tables_json(tmp_path, {"tb-1": None}) # entry without table_header
|
|
else:
|
|
_write_tables_json(tmp_path, {"tb-1": sidecar})
|
|
path = tmp_path / "doc.blocks.jsonl"
|
|
row = {
|
|
"type": "content",
|
|
"heading": heading,
|
|
"parent_headings": [],
|
|
"level": 2 if heading else 0,
|
|
"content": body,
|
|
}
|
|
path.write_text(json.dumps(row, ensure_ascii=False), encoding="utf-8")
|
|
chunks = chunking_by_paragraph_semantic(
|
|
tokenizer,
|
|
body,
|
|
chunk_token_size=chunk_token_size,
|
|
blocks_path=str(path),
|
|
chunk_overlap_token_size=0,
|
|
)
|
|
return chunks
|
|
|
|
|
|
def _table_chunks(chunks):
|
|
return [c for c in chunks if '<table id="tb-1"' in c["content"]]
|
|
|
|
|
|
def _slice_starts_with_header(chunk) -> bool:
|
|
content = chunk["content"]
|
|
tbl = content[content.index('<table id="tb-1"') :]
|
|
return tbl.startswith(_INJECTED_PREFIX)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_extract_table_id_variants():
|
|
from lightrag.table_markup import extract_table_id
|
|
|
|
assert extract_table_id('id="tb-1" format="json"') == "tb-1"
|
|
assert extract_table_id("format='html' id='tb-h'") == "tb-h"
|
|
assert extract_table_id('format="json"') is None
|
|
assert extract_table_id("") is None
|
|
# ``id`` that is only the tail of another attribute name must not match —
|
|
# neither a hyphen-joined one (``data-id``) nor a word-char prefix
|
|
# (``grid``); a standalone ``id`` alongside such an attribute still wins.
|
|
assert extract_table_id('data-id="x"') is None
|
|
assert extract_table_id('grid="y"') is None
|
|
assert extract_table_id('data-id="x" id="tb-9"') == "tb-9"
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_inject_header_into_table_slice_json_and_html():
|
|
from lightrag.chunker.paragraph_semantic import _inject_header_into_table_slice
|
|
|
|
# JSON: header rows are prepended to the row array, attrs (incl. id) kept.
|
|
json_slice = '<table id="tb-1" format="json">[["a", "b"]]</table>'
|
|
assert (
|
|
_inject_header_into_table_slice(json_slice, '[["H1", "H2"]]')
|
|
== '<table id="tb-1" format="json">[["H1", "H2"], ["a", "b"]]</table>'
|
|
)
|
|
|
|
# HTML: the raw <thead> header is spliced verbatim ahead of the body,
|
|
# preserving any rowspan/colspan markup it carries.
|
|
html_slice = (
|
|
'<table id="tb-h" format="html"><tbody><tr><td>a</td></tr></tbody></table>'
|
|
)
|
|
html_header = '<thead><tr><th colspan="2">H</th></tr></thead>'
|
|
assert _inject_header_into_table_slice(html_slice, html_header) == (
|
|
'<table id="tb-h" format="html">'
|
|
'<thead><tr><th colspan="2">H</th></tr></thead>'
|
|
"<tbody><tr><td>a</td></tr></tbody></table>"
|
|
)
|
|
|
|
# Unparseable header → no injection.
|
|
assert _inject_header_into_table_slice(json_slice, "not json") is None
|
|
# Non-<table> fragment → no injection.
|
|
assert _inject_header_into_table_slice("just text", '[["H1", "H2"]]') is None
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_inject_header_raises_on_cross_format():
|
|
"""A header whose format disagrees with the slice's content format means the
|
|
sidecar header does not belong to this table — inject must raise, not
|
|
silently emit a malformed slice."""
|
|
from lightrag.chunker.paragraph_semantic import _inject_header_into_table_slice
|
|
|
|
json_slice = '<table id="tb-1" format="json">[["a", "b"]]</table>'
|
|
html_slice = (
|
|
'<table id="tb-h" format="html"><tbody><tr><td>a</td></tr></tbody></table>'
|
|
)
|
|
html_header = "<thead><tr><th>H</th></tr></thead>"
|
|
|
|
# HTML header into a JSON slice → mismatch.
|
|
with pytest.raises(ValueError):
|
|
_inject_header_into_table_slice(json_slice, html_header)
|
|
# JSON header into an HTML slice → mismatch.
|
|
with pytest.raises(ValueError):
|
|
_inject_header_into_table_slice(html_slice, '[["H1", "H2"]]')
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_inject_header_skips_html_slice_that_already_has_thead():
|
|
from lightrag.chunker.paragraph_semantic import _inject_header_into_table_slice
|
|
|
|
# HTML: a slice that already carries a <thead> is left untouched rather
|
|
# than gaining a second one.
|
|
slice_with_thead = (
|
|
'<table id="tb-h" format="html">'
|
|
"<thead><tr><th>H2a</th><th>H2b</th></tr></thead>"
|
|
"<tbody><tr><td>d</td></tr></tbody></table>"
|
|
)
|
|
html_header = "<thead><tr><th>H1a</th><th>H1b</th></tr></thead>"
|
|
assert _inject_header_into_table_slice(slice_with_thead, html_header) is None
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_full_pipeline_injects_html_thead_into_split_html_table(tmp_path):
|
|
"""End-to-end through chunking_by_paragraph_semantic: an oversized HTML table
|
|
whose tables.json carries a raw <thead> header has that header re-injected
|
|
(verbatim, spans preserved) into every non-first table chunk."""
|
|
html_header = (
|
|
'<thead><tr><th rowspan="2">Metric</th><th colspan="2">Group</th></tr></thead>'
|
|
)
|
|
head_row = '<tr><th rowspan="2">Metric</th><th colspan="2">Group</th></tr>'
|
|
data_rows = "".join(f"<tr><td>{'d' * 160}{i}</td></tr>" for i in range(6))
|
|
table = (
|
|
'<table id="tb-h" format="html">'
|
|
f"<thead>{head_row}</thead><tbody>{data_rows}</tbody></table>"
|
|
)
|
|
body = "\n".join(["lead paragraph", table, "trailing paragraph"])
|
|
|
|
# tables.json sidecar with an HTML-format header.
|
|
(tmp_path / "doc.tables.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"version": "1.0",
|
|
"tables": {
|
|
"tb-h": {
|
|
"id": "tb-h",
|
|
"format": "html",
|
|
"content": table,
|
|
"caption": "",
|
|
"table_header": html_header,
|
|
}
|
|
},
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
blocks_path = tmp_path / "doc.blocks.jsonl"
|
|
blocks_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"type": "content",
|
|
"heading": "Section",
|
|
"parent_headings": [],
|
|
"level": 2,
|
|
"content": body,
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
tokenizer = _make_tokenizer()
|
|
chunks = chunking_by_paragraph_semantic(
|
|
tokenizer,
|
|
body,
|
|
chunk_token_size=900,
|
|
blocks_path=str(blocks_path),
|
|
chunk_overlap_token_size=0,
|
|
)
|
|
|
|
table_chunks = [c for c in chunks if '<table id="tb-h"' in c["content"]]
|
|
assert len(table_chunks) >= 2 # the table was split
|
|
# Every table chunk carries exactly one <thead> with spans intact; the
|
|
# first keeps its own, later ones got the sidecar header spliced in.
|
|
for chunk in table_chunks:
|
|
assert chunk["content"].count("<thead>") == 1, chunk["content"]
|
|
assert 'rowspan="2"' in chunk["content"]
|
|
assert 'colspan="2"' in chunk["content"]
|
|
# At least one non-first chunk literally contains the injected header.
|
|
assert any(html_header in c["content"] for c in table_chunks[1:])
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_multi_row_header_pinned_not_duplicated():
|
|
# A multi-row header is pinned out of the split body, so every slice gets
|
|
# the full header exactly once — never a cut remnant duplicated into a
|
|
# later slice (e.g. [H1, H2, H2, ...]).
|
|
from lightrag.table_markup import TABLE_TAG_RE
|
|
|
|
tokenizer = _make_tokenizer()
|
|
h1 = ["H1", "h1"]
|
|
h2 = ["H2", "h2"]
|
|
header_rows = [h1, h2]
|
|
header_body = json.dumps(header_rows)
|
|
rows = [h1, h2] + [[f"D{i}", "x" * 30] for i in range(6)]
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=160,
|
|
target_ideal=120,
|
|
last_min=10,
|
|
header_body=header_body,
|
|
)
|
|
|
|
assert len(pieces) >= 2, "expected the data rows to be split across slices"
|
|
for piece in pieces:
|
|
body_rows = json.loads(TABLE_TAG_RE.match(piece).group("body"))
|
|
# Full header present exactly once at the top of every slice.
|
|
assert body_rows[:2] == header_rows, body_rows
|
|
assert body_rows[2:].count(h1) == 0 and body_rows[2:].count(h2) == 0, body_rows
|
|
assert _count_tokens(tokenizer, piece) <= 160
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_preserves_data_row_equal_to_header():
|
|
# A genuine data row whose values equal the (single-row) header must be
|
|
# preserved, not mistaken for a header remnant and dropped. Pinning the
|
|
# header out of the body guarantees this: the body is split as pure data
|
|
# and the header is prepended verbatim.
|
|
from lightrag.table_markup import TABLE_TAG_RE
|
|
|
|
tokenizer = _make_tokenizer()
|
|
header_row = ["Label", "Value"]
|
|
header_body = json.dumps([header_row])
|
|
# The 3rd data row is identical to the header (e.g. a repeated label row).
|
|
data_rows = [
|
|
["a", "x" * 40],
|
|
["b", "y" * 40],
|
|
["Label", "Value"],
|
|
["c", "z" * 40],
|
|
]
|
|
rows = [header_row] + data_rows
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=130,
|
|
target_ideal=90,
|
|
last_min=10,
|
|
header_body=header_body,
|
|
)
|
|
|
|
assert len(pieces) >= 2, "expected the data to be split across slices"
|
|
# Reconstruct the data by stripping the one prepended header row from each
|
|
# slice; every original data row (including the header-lookalike) survives.
|
|
recovered: list = []
|
|
for piece in pieces:
|
|
body_rows = json.loads(TABLE_TAG_RE.match(piece).group("body"))
|
|
assert body_rows[0] == header_row, body_rows
|
|
recovered.extend(body_rows[1:])
|
|
assert recovered == data_rows, recovered
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_resplits_reducible_slice_to_keep_header():
|
|
# A reducible multi-row non-first slice whose wrapped table fits target_max
|
|
# but whose table+header would exceed it must be split further at row
|
|
# boundaries so each sub-slice still carries the header — rather than being
|
|
# accepted whole and then silently left header-less. Here the splitter
|
|
# emits two 2-row chunks at 119 tokens (target_max=120, effective cap=106),
|
|
# which must be re-split so every non-first slice keeps its header.
|
|
from lightrag.table_markup import TABLE_TAG_RE
|
|
|
|
tokenizer = _make_tokenizer()
|
|
header_row = ["HH", "HH"]
|
|
header_body = json.dumps([header_row])
|
|
rows = [header_row] + [[f"D{i}", "x" * 28] for i in range(4)]
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=120,
|
|
target_ideal=120,
|
|
last_min=10,
|
|
header_body=header_body,
|
|
)
|
|
|
|
assert len(pieces) >= 3
|
|
# Every piece honors the hard cap.
|
|
assert all(_count_tokens(tokenizer, p) <= 120 for p in pieces), [
|
|
_count_tokens(tokenizer, p) for p in pieces
|
|
]
|
|
# No non-first slice is left header-less: each carries the recovered header
|
|
# (every data row here is small enough that header + row fits target_max).
|
|
for piece in pieces[1:]:
|
|
body_rows = json.loads(TABLE_TAG_RE.match(piece).group("body"))
|
|
assert body_rows[0] == header_row, body_rows
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_budgets_header_before_splitting():
|
|
# With a header supplied, every emitted slice must stay ≤ target_max even
|
|
# though the header is prepended into the non-first slices — proving the
|
|
# header tokens were budgeted out BEFORE the split, not backfilled after.
|
|
tokenizer = _make_tokenizer()
|
|
header_body = '[["' + "H" * 60 + '", "' + "K" * 60 + '"]]' # non-trivial header
|
|
rows = [["H" * 60, "K" * 60]] + [[f"r{i}-" + "y" * 60, "z"] for i in range(8)]
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=300,
|
|
target_ideal=220,
|
|
last_min=64,
|
|
header_body=header_body,
|
|
)
|
|
|
|
assert len(pieces) >= 3, "expected the oversized table to split into 3+ slices"
|
|
# Every slice honors the cap including its (injected) header.
|
|
assert all(_count_tokens(tokenizer, p) <= 300 for p in pieces), [
|
|
_count_tokens(tokenizer, p) for p in pieces
|
|
]
|
|
# Non-first slices carry the recovered header row.
|
|
assert all(
|
|
p.startswith('<table id="tb-1" format="json">[["' + "H" * 60)
|
|
for p in pieces[1:]
|
|
)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_degrades_when_header_would_exceed_cap(
|
|
caplog, _propagate_lightrag_logger
|
|
):
|
|
# Cap edge: a single-row slice whose wrapped <table> sits just under
|
|
# target_max yet leaves no room for the header it would need. Rather than
|
|
# silently keeping it header-less (the old size-guard behaviour), the whole
|
|
# table degrades to a recursive split so the header is preserved as text,
|
|
# the hard cap still holds, and a warning is logged.
|
|
tokenizer = _make_tokenizer()
|
|
target_max = 200
|
|
header_body = '[["HHHHHHHHHH", "KKKKKKKKKK"]]' # ~30 tokens of header
|
|
# row0 is the real header (pinned out of the body); row1 is small; row2 is a
|
|
# single big row whose wrapped <table> size sits just under target_max, so
|
|
# it fits the bare cap yet leaves no room for the ~30-token header.
|
|
rows = [["HHHHHHHHHH", "KKKKKKKKKK"], ["s", "y"], ["x" * 140, "y"]]
|
|
table_text = f'<table id="tb-1" format="json">{json.dumps(rows)}</table>'
|
|
|
|
with caplog.at_level(logging.WARNING, logger="lightrag"):
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=target_max,
|
|
target_ideal=150,
|
|
last_min=64,
|
|
header_body=header_body,
|
|
)
|
|
|
|
assert len(pieces) >= 2
|
|
# The cap holds for every piece, and the whole table degraded: no piece
|
|
# survives as a complete re-wrapped <table>...</table> slice.
|
|
assert all(_count_tokens(tokenizer, p) <= target_max for p in pieces), [
|
|
_count_tokens(tokenizer, p) for p in pieces
|
|
]
|
|
assert not any(p.startswith("<table ") and p.endswith("</table>") for p in pieces)
|
|
# The header content survives in the recursive fallback (not dropped).
|
|
assert "HHHHHHHHHH" in "".join(pieces)
|
|
assert any("degrading the whole table" in rec.message for rec in caplog.records), [
|
|
rec.message for rec in caplog.records
|
|
]
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_text_pinned_header_oversized_row_preserves_header(
|
|
caplog, _propagate_lightrag_logger
|
|
):
|
|
# Codex P2 regression: a pinned-header JSON table whose data portion holds a
|
|
# single row larger than target_max. The old code character-split the
|
|
# header-less ``wrapped`` body, so the later injection pass (which can't
|
|
# repair a non-<table> fragment) dropped the header entirely. The whole
|
|
# table must now degrade to a recursive split that keeps the header text.
|
|
tokenizer = _make_tokenizer()
|
|
target_max = 300
|
|
header_body = '[["HEADERCOL_A", "HEADERCOL_B"]]'
|
|
# row0 repeats the header (pinned out of the body); row1 is a single huge
|
|
# data row whose content alone blows past target_max.
|
|
rows = [["HEADERCOL_A", "HEADERCOL_B"], ["x" * 2000, "y"]]
|
|
table_text = f'<table id="tb-7" format="json">{json.dumps(rows)}</table>'
|
|
|
|
with caplog.at_level(logging.WARNING, logger="lightrag"):
|
|
pieces = _split_table_text(
|
|
table_text,
|
|
tokenizer=tokenizer,
|
|
target_max=target_max,
|
|
target_ideal=220,
|
|
last_min=64,
|
|
header_body=header_body,
|
|
)
|
|
|
|
assert len(pieces) >= 2
|
|
assert all(_count_tokens(tokenizer, p) <= target_max for p in pieces), [
|
|
_count_tokens(tokenizer, p) for p in pieces
|
|
]
|
|
# The header column names survive somewhere in the emitted fragments.
|
|
assert "HEADERCOL_A" in "".join(pieces)
|
|
assert any(
|
|
"tb-7" in rec.message and "degrading the whole table" in rec.message
|
|
for rec in caplog.records
|
|
), [rec.message for rec in caplog.records]
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_header_injected_into_every_table_slice(tmp_path):
|
|
chunks = _chunk_with_oversized_table(tmp_path)
|
|
table_chunks = _table_chunks(chunks)
|
|
|
|
assert len(table_chunks) >= 2, "expected the oversized table to be split"
|
|
# Every slice's <table> now begins with the header row: the "first" slice
|
|
# naturally (it kept the real header row), middle/last via re-injection.
|
|
assert all(_slice_starts_with_header(c) for c in table_chunks), [
|
|
c["content"][:90] for c in table_chunks
|
|
]
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_every_chunk_respects_cap_after_header_injection(tmp_path):
|
|
# The Codex P2 finding: a near-cap slice + a recovered header must not push
|
|
# the chunk past the hard cap. With split-time budgeting this holds for a
|
|
# table that has substantial trailing prose glued to its last slice too.
|
|
chunk_token_size = 800
|
|
long_trailing = " ".join(f"word{i}" for i in range(120)) # ~ hundreds of tokens
|
|
chunks = _chunk_with_oversized_table(
|
|
tmp_path,
|
|
num_data_rows=10,
|
|
payload_size=120,
|
|
trailing=long_trailing,
|
|
chunk_token_size=chunk_token_size,
|
|
)
|
|
tokenizer = _make_tokenizer()
|
|
assert _table_chunks(chunks), "expected a split header-bearing table"
|
|
for c in chunks:
|
|
assert _count_tokens(tokenizer, c["content"]) <= chunk_token_size, c["content"][
|
|
:120
|
|
]
|
|
assert c["tokens"] <= chunk_token_size
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_split_table_slices_never_remerge_no_duplicate_header(tmp_path):
|
|
# Split slices are frozen against LevelMerge. Two slices of one table are
|
|
# never recombined into a single chunk (which, post-injection, would carry
|
|
# the header twice). The header appears at most once per chunk.
|
|
chunks = _chunk_with_oversized_table(
|
|
tmp_path, num_data_rows=8, payload_size=120, trailing="", chunk_token_size=800
|
|
)
|
|
table_chunks = _table_chunks(chunks)
|
|
|
|
assert len(table_chunks) >= 2, "table must remain split, not reassembled"
|
|
for c in table_chunks:
|
|
# The header signature must not be duplicated inside any single chunk.
|
|
assert c["content"].count('"H1", "H2"') <= 1, c["content"]
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_levelmerge_freezes_first_and_last_slices():
|
|
# Deterministic freeze guard at the LevelMerge layer: a small "first" block
|
|
# immediately followed by a small "last" block of the SAME 2-slice table —
|
|
# both under target_ideal and same parent path — would have been reassembled
|
|
# into one whole table before the freeze. Now neither merge direction fires,
|
|
# so the two slices stay as two separate chunks (no duplicated header).
|
|
tokenizer = _make_tokenizer()
|
|
first = _new_block(
|
|
heading="Section",
|
|
parent_headings=["Doc"],
|
|
level=2,
|
|
paragraphs=[
|
|
{
|
|
"text": '<table id="tb-1" format="json">[["H1","H2"],["a","b"]]</table>',
|
|
"is_table": True,
|
|
}
|
|
],
|
|
table_chunk_role="first",
|
|
tokenizer=tokenizer,
|
|
)
|
|
last = _new_block(
|
|
heading="Section",
|
|
parent_headings=["Doc"],
|
|
level=2,
|
|
paragraphs=[
|
|
{
|
|
"text": '<table id="tb-1" format="json">[["H1","H2"],["c","d"]]</table>',
|
|
"is_table": True,
|
|
}
|
|
],
|
|
table_chunk_role="last",
|
|
tokenizer=tokenizer,
|
|
)
|
|
|
|
merged = _merge_small_blocks(
|
|
[first, last],
|
|
tokenizer=tokenizer,
|
|
target_max=10_000, # generous cap so size never blocks a merge
|
|
target_ideal=10_000, # both blocks count as below-ideal
|
|
small_tail_threshold=10_000,
|
|
)
|
|
|
|
assert len(merged) == 2, "frozen split slices must not be reassembled"
|
|
assert all(b["content"].count('"H1","H2"') == 1 for b in merged)
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_no_injection_when_source_table_has_no_header(tmp_path):
|
|
# tables.json lists the table but WITHOUT a table_header field → nothing is
|
|
# fabricated; only the first slice (which kept the real header) starts with it.
|
|
chunks = _chunk_with_oversized_table(tmp_path, sidecar=None)
|
|
table_chunks = _table_chunks(chunks)
|
|
|
|
assert len(table_chunks) >= 2
|
|
starts = [_slice_starts_with_header(c) for c in table_chunks]
|
|
assert starts[0] is True, "first slice keeps the table's own header row"
|
|
assert any(s is False for s in starts[1:]), "later slices must not be injected"
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_no_injection_when_tables_json_missing(tmp_path):
|
|
# No tables.json at all → silent degrade: no error, no injection.
|
|
chunks = _chunk_with_oversized_table(tmp_path, sidecar=_NO_SIDECAR)
|
|
table_chunks = _table_chunks(chunks)
|
|
|
|
assert len(table_chunks) >= 2
|
|
starts = [_slice_starts_with_header(c) for c in table_chunks]
|
|
assert starts[0] is True
|
|
assert any(s is False for s in starts[1:])
|
|
|
|
|
|
@pytest.mark.offline
|
|
def test_injection_works_when_source_heading_empty(tmp_path):
|
|
# A preamble block (empty source heading) still gets the header injected.
|
|
chunks = _chunk_with_oversized_table(tmp_path, heading="")
|
|
table_chunks = _table_chunks(chunks)
|
|
|
|
assert len(table_chunks) >= 2
|
|
assert all(_slice_starts_with_header(c) for c in table_chunks)
|