diff --git a/api/services/kb/chunker.py b/api/services/kb/chunker.py index f12a5e0..d416568 100644 --- a/api/services/kb/chunker.py +++ b/api/services/kb/chunker.py @@ -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))