feat(llm): paragraph-anchored rule chunks + section metadata on every chunk

src/llm/chunk.py

@@ -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]:
+_HEADING_LINE = re.compile(r"^#{1,6}\s+(.*)$")
-    """Split at markdown headings, keeping the heading with its body."""
+
+
+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 = [
+    pieces = [
-        piece
+        (heading, piece)
-        for section in _sections(body)
+        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)
     ]

tests/llm/test_chunk.py

@@ -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]