feat(llm): markdown-aware chunker with deterministic ids (refs #566)

src/llm/chunk.py

@@ -0,0 +1,104 @@
+"""Markdown-aware chunking with deterministic ids.
+
+Ids are ``{item_key}:{content_hash[:12]}:{seq:04d}`` — same input, same
+ids — so pgvector upserts are idempotent and ``index_state`` can skip
+unchanged items by comparing hashes.
+
+Sizes are in characters (~4 chars/token; 2000 chars ≈ 500 tokens, inside
+every candidate embed model's window).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Doc:
+    key: str
+    text: str
+    metadata: dict[str, str]
+
+
+@dataclass(frozen=True)
+class Chunk:
+    id: str
+    text: str
+    metadata: dict[str, str]
+
+
+def content_hash(text: str) -> str:
+    return hashlib.sha256(text.encode()).hexdigest()
+
+
+_FRONTMATTER = re.compile(r"\A---\n.*?\n---\n", re.DOTALL)
+_HEADING = re.compile(r"^#{1,6} ", re.MULTILINE)
+
+
+def _sections(text: str) -> list[str]:
+    """Split at markdown headings, keeping the heading with its body."""
+    starts = [m.start() for m in _HEADING.finditer(text)]
+    if not starts:
+        return [text]
+    bounds = ([0] if starts[0] != 0 else []) + starts + [len(text)]
+    return [text[a:b] for a, b in zip(bounds, bounds[1:])]
+
+
+def _hard_wrap(para: str, target: int, overlap: int) -> list[str]:
+    """Slice an oversized paragraph into <= target windows, each window's
+    leading `overlap` chars equal to the previous window's trailing
+    `overlap` chars."""
+    n = len(para)
+    pieces: list[str] = []
+    start = 0
+    while True:
+        end = min(start + target, n)
+        pieces.append(para[start:end])
+        if end == n:
+            break
+        start = end - overlap
+    return pieces
+
+
+def _pack(section: str, target: int, overlap: int) -> list[str]:
+    """Greedily pack paragraphs; hard-wrap oversized ones with overlap."""
+    paras = [p for p in re.split(r"\n\n+", section) if p.strip()]
+    pieces: list[str] = []
+    buf = ""
+    for para in paras:
+        if buf and len(buf) + len(para) + 2 > target:
+            pieces.append(buf)
+            buf = ""
+        if len(para) > target:
+            pieces.extend(_hard_wrap(para, target, overlap))
+            continue
+        buf = f"{buf}\n\n{para}" if buf else para
+    if buf:
+        pieces.append(buf)
+    return pieces
+
+
+def chunk_doc(
+    doc: Doc, *, target_chars: int = 2000, overlap_chars: int = 200
+) -> list[Chunk]:
+    if overlap_chars >= target_chars:
+        raise ValueError("overlap_chars must be smaller than target_chars")
+    body = _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)
+        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)},
+        )
+        for seq, piece in enumerate(texts)
+    ]

tests/llm/test_chunk.py

@@ -0,0 +1,57 @@
+"""llm.chunk — markdown-aware chunking with deterministic ids."""
+
+from llm.chunk import Doc, chunk_doc, content_hash
+
+FRONTMATTER_DOC = """---
+comment_id: CMS-2019-0111-0042
+docket_id: CMS-2019-0111
+---
+
+# Re: CY 2020 PFS Proposed Rule
+
+We object to the E/M consolidation.
+
+## Telehealth
+
+Originating-site rules should be relaxed.
+"""
+
+
+def _doc(text, key="K1"):
+    return Doc(key=key, text=text, metadata={"docket": "CMS-2019-0111"})
+
+
+class TestChunkDoc:
+    def test_strips_yaml_frontmatter(self):
+        chunks = chunk_doc(_doc(FRONTMATTER_DOC))
+        assert "comment_id:" not in chunks[0].text
+        assert chunks[0].text.startswith("# Re:")
+
+    def test_deterministic_ids(self):
+        a = chunk_doc(_doc(FRONTMATTER_DOC))
+        b = chunk_doc(_doc(FRONTMATTER_DOC))
+        assert [c.id for c in a] == [c.id for c in b]
+        h = content_hash(FRONTMATTER_DOC)
+        assert a[0].id == f"K1:{h[:12]}:0000"
+
+    def test_metadata_carries_item_key_and_seq(self):
+        chunks = chunk_doc(_doc(FRONTMATTER_DOC))
+        assert chunks[0].metadata["item_key"] == "K1"
+        assert chunks[0].metadata["docket"] == "CMS-2019-0111"
+        assert chunks[0].metadata["seq"] == "0"
+
+    def test_splits_on_headings_before_size(self):
+        text = "# A\n\n" + "para. " * 100 + "\n\n# B\n\nshort."
+        chunks = chunk_doc(_doc(text), target_chars=300)
+        assert all(len(c.text) <= 300 + 200 for c in chunks)
+        assert any(c.text.lstrip().startswith("# B") for c in chunks)
+
+    def test_long_paragraph_hard_wrapped_with_overlap(self):
+        text = "x" * 5000
+        chunks = chunk_doc(_doc(text), target_chars=2000, overlap_chars=200)
+        assert len(chunks) == 3
+        assert chunks[1].text[:200] == chunks[0].text[-200:]
+
+    def test_empty_and_whitespace_yield_nothing(self):
+        assert chunk_doc(_doc("")) == []
+        assert chunk_doc(_doc("  \n\n  ")) == []