diff --git a/api/services/kb/chunker.py b/api/services/kb/chunker.py index d8f3284..f12a5e0 100644 --- a/api/services/kb/chunker.py +++ b/api/services/kb/chunker.py @@ -11,11 +11,34 @@ from dataclasses import dataclass from typing import Literal # Page marker sentinel embedded by the PDF parser (api/services/kb/ingest.py -# _page_marker). Format: \x00KB_PAGE:\x00 where N is 1-indexed. The chunker -# strips markers from chunk.content and records the last-seen page as the -# chunk's page_number. +# _page_marker). Format: \x00KB_PAGE:\x00 where N is 1-indexed. +# _build_page_map returns a list of (offset, page_number) for each marker so +# the chunker can attach the right page to any chunk by position — _strip_page_markers +# then removes the sentinels from chunk.content. _PAGE_MARKER_RE = re.compile(r"\x00KB_PAGE:(\d+)\x00") + +def _build_page_map(text: str) -> list[tuple[int, int]]: + """Return sorted [(offset, page_number), …] of every marker in `text`.""" + return [(m.start(), int(m.group(1))) for m in _PAGE_MARKER_RE.finditer(text)] + + +def _page_for_offset(page_map: list[tuple[int, int]], offset: int) -> int | None: + """Last-seen page number at or before `offset`. None if nothing earlier.""" + if not page_map: + return None + lo, hi = 0, len(page_map) - 1 + best: int | None = None + while lo <= hi: + mid = (lo + hi) // 2 + mark_off, mark_page = page_map[mid] + if mark_off <= offset: + best = mark_page + lo = mid + 1 + else: + hi = mid - 1 + return best + # Approx tokens per char for Hebrew/mixed text (voyage-multilingual-2 tokenizer). # Voyage tends to split Hebrew into 1.5-2 tokens per word; 0.55 tokens/char # is a conservative estimate that matches what the API actually reports. @@ -83,12 +106,13 @@ def _count_tokens(text: str) -> int: def _strip_page_markers(text: str) -> tuple[str, int | None]: - """Remove page markers from `text` and return the first page number seen. + """Remove page markers from `text` and return the FIRST page number seen. - The first-marker rule matches how users think about a citation ("this - paragraph starts on page 42") — subsequent markers on the same chunk - mean the paragraph spans pages, but the jump-to-PDF link should point - at the page where the quote begins. + Only useful as a fallback when the chunker didn't attach a page via + offset math. Markers that show up mid-content (e.g. a section that + spans pages) would resolve to the spanning page rather than the + chunk's starting page — so prefer c.page_number (the caller's + offset-derived value) and fall back to this only when that is None. """ first_page: int | None = None m = _PAGE_MARKER_RE.search(text) @@ -103,14 +127,18 @@ def _strip_page_markers(text: str) -> tuple[str, int | None]: def _as_dicts(chunks: list[Chunk]) -> list[dict]: result: list[dict] = [] for c in chunks: - clean, page = _strip_page_markers(c.content) + clean, page_in_body = _strip_page_markers(c.content) + # Prefer the page the caller computed from the chunk's START offset + # (the "where does this quote begin" answer). Only fall back to any + # marker inside the body if the caller didn't set one. + page = c.page_number if c.page_number is not None else page_in_body result.append({ "chunk_index": c.chunk_index, "heading_path": c.heading_path or None, "section_ref": c.section_ref or None, "content": clean.strip(), "token_count": c.token_count, - "page_number": page if page is not None else c.page_number, + "page_number": page, }) return result @@ -119,6 +147,8 @@ 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). """ max_chars = int(_MAX_TOKENS_PER_CHUNK / _TOKENS_PER_CHAR) result: list[Chunk] = [] @@ -137,6 +167,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, ) ) return result @@ -183,14 +214,22 @@ def chunk_statute(text: str) -> list[dict]: heading_path format: "חלק ט' — ביטוח נכות > פרק ב' > סימן א' > ס' 195(א)". section_ref stays the leaf ("ס' 195(א)") for precise citation. """ - text = text.strip() + # 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 [] + page_map = _build_page_map(text) + markers = _collect_markers(text) section_markers = [m for m in markers if m["level"] == "section"] if not section_markers: - return _as_dicts([Chunk(0, "", "", text, _count_tokens(text))]) + return _as_dicts([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 @@ -233,6 +272,7 @@ def chunk_statute(text: str) -> list[dict]: section_ref=m["ref"], content=content, token_count=_count_tokens(content), + page_number=_page_for_offset(page_map, start), ) ) return _as_dicts(chunks) @@ -244,29 +284,31 @@ def chunk_circular( overlap_tokens: int = 80, ) -> list[dict]: """Circular: header-aware split, then pack paragraphs up to target size with overlap.""" - text = text.strip() + text = text.strip("\n ") if not text: return [] + page_map = _build_page_map(text) target_chars = int(target_tokens / _TOKENS_PER_CHAR) overlap_chars = int(overlap_tokens / _TOKENS_PER_CHAR) headers = list(_HEADER_RE.finditer(text)) - segments: list[tuple[str, str]] = [] # (heading, body) + segments: list[tuple[str, str, int]] = [] # (heading, body, offset) if not headers: - segments.append(("", text)) + segments.append(("", text, 0)) else: if headers[0].start() > 0: - segments.append(("", text[: headers[0].start()].strip())) + segments.append(("", text[: headers[0].start()].strip(), 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)) + segments.append((heading, body, h.end())) chunks: list[Chunk] = [] - for heading, body in segments: + for heading, body, offset in segments: + seg_page = _page_for_offset(page_map, offset) if len(body) <= target_chars: chunks.append( Chunk( @@ -275,6 +317,7 @@ def chunk_circular( section_ref=heading, content=body, token_count=_count_tokens(body), + page_number=seg_page, ) ) continue @@ -291,6 +334,7 @@ def chunk_circular( section_ref=heading, content=buf.strip(), token_count=_count_tokens(buf), + page_number=seg_page, ) ) # Overlap: keep last N chars @@ -305,6 +349,7 @@ def chunk_circular( section_ref=heading, content=buf.strip(), token_count=_count_tokens(buf), + page_number=seg_page, ) ) @@ -316,7 +361,9 @@ def chunk(text: str, kind: Literal["law", "regulation", "circular"]) -> list[dic raw = chunk_statute(text) else: raw = chunk_circular(text) - # Reconstruct Chunk objects to apply the safety ceiling. + # 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"], @@ -324,6 +371,7 @@ def chunk(text: str, kind: Literal["law", "regulation", "circular"]) -> list[dic section_ref=r["section_ref"] or "", content=r["content"], token_count=r["token_count"], + page_number=r.get("page_number"), ) for r in raw ]