File size: 3,312 Bytes
acf77ab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | from __future__ import annotations
from markdown_it import MarkdownIt
from markdown_it.token import Token
from pydantic import BaseModel, ConfigDict
MIN_CHUNK_CHARS = 80
MAX_HEADING_LEVEL = 3
_md = MarkdownIt()
class SectionChunk(BaseModel):
model_config = ConfigDict(frozen=True)
section_path: tuple[str, ...]
section_body: str
def _segment_tokens(
tokens: list[Token],
) -> list[tuple[int, str, list[str]]]:
"""Walk token stream, returning heading-delimited segments.
Each segment is (heading_level, heading_text, body_texts).
Level 0 is a sentinel for content before the first H1-H3.
"""
segments: list[tuple[int, str, list[str]]] = []
current: tuple[int, str, list[str]] | None = None
for i, tok in enumerate(tokens):
if tok.type == "heading_open":
level = int(tok.tag[1]) # "h2" -> 2
heading_text = tokens[i + 1].content.strip()
if level <= MAX_HEADING_LEVEL:
if current is not None:
segments.append(current)
current = (level, heading_text, [])
else:
if current is None:
current = (0, "", [])
current[2].append(heading_text)
elif tok.type.endswith("_open"):
continue
elif tok.type == "inline":
if current is None:
current = (0, "", [])
current[2].append(tok.content)
if current is not None:
segments.append(current)
return segments
def chunk_body(body: str) -> list[SectionChunk]:
"""Split markdown body into section chunks by H1-H3 headings."""
tokens = _md.parse(body)
segments = _segment_tokens(tokens)
if not segments:
text = body.strip()
if text:
return [SectionChunk(section_path=(), section_body=text)]
return []
stack: list[str] = []
chunks: list[SectionChunk] = []
for level, title, body_parts in segments:
if level == 0:
chunks.append(
SectionChunk(
section_path=(),
section_body="\n\n".join(body_parts).strip(),
)
)
continue
while len(stack) >= level:
stack.pop()
stack.append(title)
chunks.append(
SectionChunk(
section_path=tuple(stack),
section_body="\n\n".join(body_parts).strip(),
)
)
return _merge_small(chunks)
def _merge_small(chunks: list[SectionChunk]) -> list[SectionChunk]:
"""Merge chunks smaller than MIN_CHUNK_CHARS into their successor."""
if not chunks:
return chunks
merged: list[SectionChunk] = []
carry: str = ""
for c in chunks:
body = (carry + "\n\n" + c.section_body).strip() if carry else c.section_body
if len(body) < MIN_CHUNK_CHARS and c is not chunks[-1]:
carry = body
continue
merged.append(SectionChunk(section_path=c.section_path, section_body=body))
carry = ""
if carry and merged:
last = merged[-1]
merged[-1] = SectionChunk(
section_path=last.section_path,
section_body=(last.section_body + "\n\n" + carry).strip(),
)
return merged
|