feat(llm): paragraph-anchored rule chunks + section metadata on every chunk
This commit is contained in:
kert
@@ -15,11 +15,27 @@ import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Paragraph:
|
||||
"""One FR paragraph anchor (mirrors ``bib.frlink.Anchor``)."""
|
||||
|
||||
p_id: int
|
||||
page: int
|
||||
ordinal: int
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Doc:
|
||||
key: str
|
||||
text: str
|
||||
metadata: dict[str, str]
|
||||
# Rules: the anchor paragraphs that make up ``text``; when set,
|
||||
# chunk_doc packs whole paragraphs and stamps p_id/page per chunk.
|
||||
paragraphs: tuple[Paragraph, ...] = ()
|
||||
# (section heading, local file path) pairs for PDF page lookup —
|
||||
# ``llm.pages.enrich_pdf_pages`` matches chunks' ``section`` to these.
|
||||
files: tuple[tuple[str, str], ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -41,13 +57,21 @@ _HEADING = re.compile(r"^#{1,6} ", re.MULTILINE)
|
||||
_CONTROL = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]")
|
||||
|
||||
|
||||
def _sections(text: str) -> list[str]:
|
||||
"""Split at markdown headings, keeping the heading with its body."""
|
||||
_HEADING_LINE = re.compile(r"^#{1,6}\s+(.*)$")
|
||||
|
||||
|
||||
def _sections(text: str) -> list[tuple[str, str]]:
|
||||
"""Split at markdown headings → (heading text, section incl. heading)."""
|
||||
starts = [m.start() for m in _HEADING.finditer(text)]
|
||||
if not starts:
|
||||
return [text]
|
||||
return [("", text)]
|
||||
bounds = ([0] if starts[0] != 0 else []) + starts + [len(text)]
|
||||
return [text[a:b] for a, b in zip(bounds, bounds[1:])]
|
||||
out = []
|
||||
for a, b in zip(bounds, bounds[1:]):
|
||||
section = text[a:b]
|
||||
m = _HEADING_LINE.match(section.split("\n", 1)[0])
|
||||
out.append((m.group(1).strip() if m else "", section))
|
||||
return out
|
||||
|
||||
|
||||
def _hard_wrap(para: str, target: int, overlap: int) -> list[str]:
|
||||
@@ -89,25 +113,82 @@ def chunk_doc(
|
||||
) -> list[Chunk]:
|
||||
"""Chunk *doc* into <= target_chars windows with overlap between them.
|
||||
|
||||
Rule docs (``doc.paragraphs`` set) pack whole FR paragraphs instead
|
||||
and carry ``p_id``/``p_id_last``/``page``/``ordinal`` metadata.
|
||||
Every chunk carries ``section`` — the markdown heading it sits under
|
||||
(``""`` when none), which is how comment chunks know their attachment.
|
||||
|
||||
Raises ``ValueError`` when ``overlap_chars >= target_chars`` (an
|
||||
overlap that large or larger would never let the window advance).
|
||||
"""
|
||||
if overlap_chars >= target_chars:
|
||||
raise ValueError("overlap_chars must be smaller than target_chars")
|
||||
if doc.paragraphs:
|
||||
return _chunk_paragraphs(doc, target_chars, overlap_chars)
|
||||
body = _CONTROL.sub("", _FRONTMATTER.sub("", doc.text)).strip()
|
||||
if not body:
|
||||
return []
|
||||
prefix = f"{doc.key}:{content_hash(doc.text)[:12]}"
|
||||
texts = [
|
||||
piece
|
||||
for section in _sections(body)
|
||||
pieces = [
|
||||
(heading, piece)
|
||||
for heading, section in _sections(body)
|
||||
for piece in _pack(section, target_chars, overlap_chars)
|
||||
]
|
||||
return [
|
||||
Chunk(
|
||||
id=f"{prefix}:{seq:04d}",
|
||||
text=piece,
|
||||
metadata={**doc.metadata, "item_key": doc.key, "seq": str(seq)},
|
||||
metadata={
|
||||
**doc.metadata,
|
||||
"item_key": doc.key,
|
||||
"seq": str(seq),
|
||||
"section": heading,
|
||||
},
|
||||
)
|
||||
for seq, piece in enumerate(texts)
|
||||
for seq, (heading, piece) in enumerate(pieces)
|
||||
]
|
||||
|
||||
|
||||
def _chunk_paragraphs(doc: Doc, target: int, overlap: int) -> list[Chunk]:
|
||||
"""Greedy pack of whole paragraphs; an oversized paragraph is
|
||||
hard-wrapped with overlap, every piece keeping its own anchor."""
|
||||
packed: list[tuple[str, Paragraph, Paragraph]] = [] # text, first, last
|
||||
buf: list[Paragraph] = []
|
||||
|
||||
def flush() -> None:
|
||||
if buf:
|
||||
packed.append(("\n\n".join(p.text for p in buf), buf[0], buf[-1]))
|
||||
buf.clear()
|
||||
|
||||
for para in doc.paragraphs:
|
||||
text = _CONTROL.sub("", para.text).strip()
|
||||
if not text:
|
||||
continue
|
||||
para = Paragraph(para.p_id, para.page, para.ordinal, text)
|
||||
if len(text) > target:
|
||||
flush()
|
||||
for piece in _hard_wrap(text, target, overlap):
|
||||
packed.append((piece, para, para))
|
||||
continue
|
||||
if buf and sum(len(p.text) + 2 for p in buf) + len(text) > target:
|
||||
flush()
|
||||
buf.append(para)
|
||||
flush()
|
||||
prefix = f"{doc.key}:{content_hash(doc.text)[:12]}"
|
||||
return [
|
||||
Chunk(
|
||||
id=f"{prefix}:{seq:04d}",
|
||||
text=text,
|
||||
metadata={
|
||||
**doc.metadata,
|
||||
"item_key": doc.key,
|
||||
"seq": str(seq),
|
||||
"section": "",
|
||||
"p_id": str(first.p_id),
|
||||
"p_id_last": str(last.p_id),
|
||||
"page": str(first.page),
|
||||
"ordinal": str(first.ordinal),
|
||||
},
|
||||
)
|
||||
for seq, (text, first, last) in enumerate(packed)
|
||||
]
|
||||
|
||||
@@ -77,3 +77,77 @@ class TestControlCharStripping:
|
||||
doc = _doc("line1\n\nline2\twith tab")
|
||||
joined = "".join(c.text for c in chunk_doc(doc))
|
||||
assert "\t" in joined and "line2" in joined
|
||||
|
||||
|
||||
# ── section metadata + paragraph-anchored rule chunks ────────────────
|
||||
|
||||
from llm.chunk import Paragraph # noqa: E402
|
||||
|
||||
|
||||
def _para(p_id, page, ordinal, text):
|
||||
return Paragraph(p_id=p_id, page=page, ordinal=ordinal, text=text)
|
||||
|
||||
|
||||
class TestSectionMetadata:
|
||||
def test_heading_text_recorded_per_chunk(self):
|
||||
doc = Doc(
|
||||
key="K",
|
||||
text="intro para\n\n## attachment_1.pdf\n\nbody one\n\n## attachment_2.docx\n\nbody two",
|
||||
metadata={},
|
||||
)
|
||||
chunks = chunk_doc(doc, target_chars=60, overlap_chars=5)
|
||||
assert [c.metadata["section"] for c in chunks] == [
|
||||
"",
|
||||
"attachment_1.pdf",
|
||||
"attachment_2.docx",
|
||||
]
|
||||
|
||||
def test_no_heading_is_empty_section(self):
|
||||
(c,) = chunk_doc(Doc(key="K", text="plain", metadata={}))
|
||||
assert c.metadata["section"] == ""
|
||||
|
||||
|
||||
class TestParagraphChunks:
|
||||
def test_packs_whole_paragraphs_and_stamps_anchor_metadata(self):
|
||||
paras = (
|
||||
_para(10, 100, 1, "A" * 30),
|
||||
_para(11, 100, 2, "B" * 30),
|
||||
_para(12, 101, 1, "C" * 30),
|
||||
)
|
||||
doc = Doc(
|
||||
key="R",
|
||||
text="\n\n".join(p.text for p in paras),
|
||||
metadata={"kind": "rule"},
|
||||
paragraphs=paras,
|
||||
)
|
||||
chunks = chunk_doc(doc, target_chars=70, overlap_chars=5)
|
||||
assert [c.text for c in chunks] == ["A" * 30 + "\n\n" + "B" * 30, "C" * 30]
|
||||
assert chunks[0].metadata["p_id"] == "10"
|
||||
assert chunks[0].metadata["p_id_last"] == "11"
|
||||
assert chunks[0].metadata["page"] == "100"
|
||||
assert chunks[0].metadata["ordinal"] == "1"
|
||||
assert chunks[1].metadata["p_id"] == "12"
|
||||
assert chunks[1].metadata["page"] == "101"
|
||||
assert chunks[0].metadata["kind"] == "rule"
|
||||
assert chunks[0].metadata["item_key"] == "R"
|
||||
assert chunks[0].id.endswith(":0000") and chunks[1].id.endswith(":0001")
|
||||
|
||||
def test_oversized_paragraph_is_wrapped_but_keeps_its_anchor(self):
|
||||
paras = (_para(5, 200, 3, "X" * 100),)
|
||||
doc = Doc(key="R", text=paras[0].text, metadata={}, paragraphs=paras)
|
||||
chunks = chunk_doc(doc, target_chars=40, overlap_chars=10)
|
||||
assert len(chunks) == 3 # windows 0-40, 30-70, 60-100
|
||||
assert {c.metadata["p_id"] for c in chunks} == {"5"}
|
||||
assert chunks[1].text[:10] == chunks[0].text[-10:]
|
||||
|
||||
def test_empty_paragraph_text_skipped(self):
|
||||
paras = (_para(1, 1, 1, " "), _para(2, 1, 2, "real"))
|
||||
doc = Doc(key="R", text="real", metadata={}, paragraphs=paras)
|
||||
(c,) = chunk_doc(doc)
|
||||
assert c.metadata["p_id"] == "2"
|
||||
|
||||
def test_ids_are_deterministic(self):
|
||||
paras = (_para(1, 1, 1, "same"),)
|
||||
a = chunk_doc(Doc(key="R", text="same", metadata={}, paragraphs=paras))
|
||||
b = chunk_doc(Doc(key="R", text="same", metadata={}, paragraphs=paras))
|
||||
assert [c.id for c in a] == [c.id for c in b]
|
||||
|
||||
Reference in New Issue
Block a user