fix(kb/chunker): derive page_number per piece in _split_oversized + chunk_circular
Commit 426b0d6 fixed page propagation through the Chunk reconstruction in
chunk(), but two downstream paths still collapsed every sub-chunk to the
parent's starting page:
1. _split_oversized operated on markerless content (since _as_dicts inside
chunk_statute/chunk_circular stripped markers before returning), so when
a single giant chunk was sliced into pieces — the common path for
ספר הליקויים / ספר המבחנים, which have no statute-style section markers —
every piece inherited c.page_number = parent's start page. 209/209 chunks
landed on page 1.
2. chunk_circular packed paragraphs into sub-chunks all tagged with
seg_page (the segment's opening page), with no attempt to track where
each sub-chunk actually begins within the original text. חוזר תקנה 37
(8 chunks) and חוזר נכות 1941 (4 chunks) — neither had header matches —
landed entirely on page 1.
Fixes:
- chunk_statute and chunk_circular now return list[Chunk] with markers
still embedded in content; _as_dicts is called once at the end of
chunk() so markers are stripped AFTER splitting.
- _split_oversized builds a local page map from markers in each parent
chunk's content, with a virtual (0, parent.page_number) entry so pieces
before the first real marker still inherit the correct page.
- chunk_circular tracks each paragraph's absolute offset in the original
text and computes each emitted sub-chunk's page_number from its first
paragraph's offset against the full-text page_map.
Verified on synthetic 3-page input: regulation path now emits pages 1,1,2,3
across 4 pieces (was 1,1,1,1); circular path emits 1,1,2,2,3,3 (was
1,1,1,1,1,1).
Refs Task Master #2
This commit is contained in:
+70
-45
@@ -147,8 +147,10 @@ def _split_oversized(chunks: list[Chunk]) -> list[Chunk]:
|
||||
"""Split any chunk whose token_count exceeds the ceiling into fixed-size pieces.
|
||||
|
||||
Preserves heading_path/section_ref on each piece and re-indexes.
|
||||
Page numbers carry across: each piece inherits the parent chunk's
|
||||
page_number (the earliest page the parent spans).
|
||||
Page numbers are re-derived per piece by scanning the parent chunk's
|
||||
content for markers: the parent's own page_number is treated as a
|
||||
virtual marker at offset 0, and any real \\x00KB_PAGE:N\\x00 markers
|
||||
preserved in content advance the active page from that point onward.
|
||||
"""
|
||||
max_chars = int(_MAX_TOKENS_PER_CHUNK / _TOKENS_PER_CHAR)
|
||||
result: list[Chunk] = []
|
||||
@@ -158,6 +160,11 @@ def _split_oversized(chunks: list[Chunk]) -> list[Chunk]:
|
||||
result.append(c)
|
||||
continue
|
||||
text = c.content
|
||||
local_map = _build_page_map(text)
|
||||
# Virtual marker at offset 0 so pieces before the first real marker
|
||||
# inherit the parent's known page.
|
||||
if c.page_number is not None and (not local_map or local_map[0][0] > 0):
|
||||
local_map = [(0, c.page_number)] + local_map
|
||||
for start in range(0, len(text), max_chars):
|
||||
piece = text[start : start + max_chars]
|
||||
result.append(
|
||||
@@ -167,7 +174,7 @@ def _split_oversized(chunks: list[Chunk]) -> list[Chunk]:
|
||||
section_ref=c.section_ref,
|
||||
content=piece,
|
||||
token_count=_count_tokens(piece),
|
||||
page_number=c.page_number,
|
||||
page_number=_page_for_offset(local_map, start) or c.page_number,
|
||||
)
|
||||
)
|
||||
return result
|
||||
@@ -208,15 +215,16 @@ def _collect_markers(text: str) -> list[dict]:
|
||||
return markers
|
||||
|
||||
|
||||
def chunk_statute(text: str) -> list[dict]:
|
||||
def chunk_statute(text: str) -> list[Chunk]:
|
||||
"""Law / regulation: one chunk per section, with hierarchical heading_path.
|
||||
|
||||
heading_path format: "חלק ט' — ביטוח נכות > פרק ב' > סימן א' > ס' 195(א)".
|
||||
section_ref stays the leaf ("ס' 195(א)") for precise citation.
|
||||
|
||||
Returns Chunk objects with page markers still embedded in content.
|
||||
chunk() calls _split_oversized next, which uses those markers to
|
||||
compute per-piece page numbers, and _as_dicts finally strips them.
|
||||
"""
|
||||
# Preserve the original text so page markers survive for content-based
|
||||
# chunks; _strip_page_markers in _as_dicts removes them from the final
|
||||
# content before storage.
|
||||
text = text.strip("\n ")
|
||||
if not text:
|
||||
return []
|
||||
@@ -226,10 +234,10 @@ def chunk_statute(text: str) -> list[dict]:
|
||||
markers = _collect_markers(text)
|
||||
section_markers = [m for m in markers if m["level"] == "section"]
|
||||
if not section_markers:
|
||||
return _as_dicts([Chunk(
|
||||
return [Chunk(
|
||||
0, "", "", text, _count_tokens(text),
|
||||
page_number=_page_for_offset(page_map, 0),
|
||||
)])
|
||||
)]
|
||||
|
||||
# Walk markers in order, tracking the active part/chapter/subchapter.
|
||||
active_part: dict | None = None
|
||||
@@ -275,15 +283,21 @@ def chunk_statute(text: str) -> list[dict]:
|
||||
page_number=_page_for_offset(page_map, start),
|
||||
)
|
||||
)
|
||||
return _as_dicts(chunks)
|
||||
return chunks
|
||||
|
||||
|
||||
def chunk_circular(
|
||||
text: str,
|
||||
target_tokens: int = 600,
|
||||
overlap_tokens: int = 80,
|
||||
) -> list[dict]:
|
||||
"""Circular: header-aware split, then pack paragraphs up to target size with overlap."""
|
||||
) -> list[Chunk]:
|
||||
"""Circular: header-aware split, then pack paragraphs up to target size with overlap.
|
||||
|
||||
Returns Chunk objects with page markers still embedded in content.
|
||||
Each sub-chunk's page_number is derived from the absolute offset of
|
||||
the first paragraph it contains — so chunks that start mid-document
|
||||
reflect the correct page, not the segment's opening page.
|
||||
"""
|
||||
text = text.strip("\n ")
|
||||
if not text:
|
||||
return []
|
||||
@@ -293,22 +307,27 @@ def chunk_circular(
|
||||
overlap_chars = int(overlap_tokens / _TOKENS_PER_CHAR)
|
||||
|
||||
headers = list(_HEADER_RE.finditer(text))
|
||||
segments: list[tuple[str, str, int]] = [] # (heading, body, offset)
|
||||
# (heading, body_raw_with_markers, body_start_offset_in_text)
|
||||
segments: list[tuple[str, str, int]] = []
|
||||
if not headers:
|
||||
segments.append(("", text, 0))
|
||||
else:
|
||||
if headers[0].start() > 0:
|
||||
segments.append(("", text[: headers[0].start()].strip(), 0))
|
||||
segments.append(("", text[: headers[0].start()], 0))
|
||||
for i, h in enumerate(headers):
|
||||
end = headers[i + 1].start() if i + 1 < len(headers) else len(text)
|
||||
heading = h.group(0).strip().rstrip(":").lstrip("#").strip()
|
||||
body = text[h.end() : end].strip()
|
||||
if body:
|
||||
segments.append((heading, body, h.end()))
|
||||
body_raw = text[h.end() : end]
|
||||
if body_raw.strip():
|
||||
segments.append((heading, body_raw, h.end()))
|
||||
|
||||
# Paragraph regex: one or more non-empty lines with no blank line between them.
|
||||
_PARA_RE = re.compile(r"[^\n]+(?:\n(?!\n)[^\n]+)*")
|
||||
|
||||
chunks: list[Chunk] = []
|
||||
for heading, body, offset in segments:
|
||||
seg_page = _page_for_offset(page_map, offset)
|
||||
for heading, body_raw, body_start in segments:
|
||||
body = body_raw.strip()
|
||||
seg_page = _page_for_offset(page_map, body_start)
|
||||
if len(body) <= target_chars:
|
||||
chunks.append(
|
||||
Chunk(
|
||||
@@ -322,11 +341,17 @@ def chunk_circular(
|
||||
)
|
||||
continue
|
||||
|
||||
# Pack paragraphs
|
||||
paragraphs = [p.strip() for p in body.split("\n\n") if p.strip()]
|
||||
# Collect paragraphs with their absolute offsets in the original text.
|
||||
paragraphs: list[tuple[str, int]] = []
|
||||
for pm in _PARA_RE.finditer(body_raw):
|
||||
p_text = pm.group(0).strip()
|
||||
if p_text:
|
||||
paragraphs.append((p_text, body_start + pm.start()))
|
||||
|
||||
buf = ""
|
||||
for p in paragraphs:
|
||||
if buf and len(buf) + len(p) + 2 > target_chars:
|
||||
buf_abs_start: int | None = None
|
||||
for p_text, p_abs in paragraphs:
|
||||
if buf and len(buf) + len(p_text) + 2 > target_chars:
|
||||
chunks.append(
|
||||
Chunk(
|
||||
chunk_index=len(chunks),
|
||||
@@ -334,13 +359,21 @@ def chunk_circular(
|
||||
section_ref=heading,
|
||||
content=buf.strip(),
|
||||
token_count=_count_tokens(buf),
|
||||
page_number=seg_page,
|
||||
page_number=(
|
||||
_page_for_offset(page_map, buf_abs_start)
|
||||
if buf_abs_start is not None else seg_page
|
||||
),
|
||||
)
|
||||
)
|
||||
# Overlap: keep last N chars
|
||||
buf = (buf[-overlap_chars:] + "\n\n" + p) if overlap_chars > 0 else p
|
||||
# Overlap: keep last N chars of prior buf, then start fresh at p.
|
||||
# The new buf's "start page" is the new paragraph's page — the
|
||||
# overlap tail is treated as prefix context, not page anchor.
|
||||
buf = (buf[-overlap_chars:] + "\n\n" + p_text) if overlap_chars > 0 else p_text
|
||||
buf_abs_start = p_abs
|
||||
else:
|
||||
buf = f"{buf}\n\n{p}" if buf else p
|
||||
if not buf:
|
||||
buf_abs_start = p_abs
|
||||
buf = f"{buf}\n\n{p_text}" if buf else p_text
|
||||
if buf.strip():
|
||||
chunks.append(
|
||||
Chunk(
|
||||
@@ -349,30 +382,22 @@ def chunk_circular(
|
||||
section_ref=heading,
|
||||
content=buf.strip(),
|
||||
token_count=_count_tokens(buf),
|
||||
page_number=seg_page,
|
||||
page_number=(
|
||||
_page_for_offset(page_map, buf_abs_start)
|
||||
if buf_abs_start is not None else seg_page
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return _as_dicts(chunks)
|
||||
return chunks
|
||||
|
||||
|
||||
def chunk(text: str, kind: Literal["law", "regulation", "circular"]) -> list[dict]:
|
||||
if kind in ("law", "regulation"):
|
||||
raw = chunk_statute(text)
|
||||
chunks = chunk_statute(text)
|
||||
else:
|
||||
raw = chunk_circular(text)
|
||||
# Reconstruct Chunk objects to apply the safety ceiling. page_number
|
||||
# must be copied back here — otherwise the reconstruction wipes it and
|
||||
# every chunk ends up with NULL in the DB.
|
||||
chunks = [
|
||||
Chunk(
|
||||
chunk_index=r["chunk_index"],
|
||||
heading_path=r["heading_path"] or "",
|
||||
section_ref=r["section_ref"] or "",
|
||||
content=r["content"],
|
||||
token_count=r["token_count"],
|
||||
page_number=r.get("page_number"),
|
||||
)
|
||||
for r in raw
|
||||
]
|
||||
chunks = chunk_circular(text)
|
||||
# _split_oversized uses page markers still embedded in content to derive
|
||||
# per-piece page numbers; _as_dicts then strips markers and emits the
|
||||
# final serializable shape.
|
||||
return _as_dicts(_split_oversized(chunks))
|
||||
|
||||
Reference in New Issue
Block a user