fix(import): keep short pdfs, stitch long tables and stop doubled or dropped charts

wikify/engine/images.py

@@ -4,6 +4,7 @@
 import unicodedata
 from collections import Counter
 from dataclasses import dataclass
+from itertools import count
 
 import fitz
 
@@ -36,6 +37,7 @@
 )
 _SKIPPED_MARKUP_RE = re.compile(r"!\[[^\]]*\]\([^)]*\)|<[^<>]*>|&#?\w+;")
 _FENCE_RE = re.compile(r"^[ \t]*```", re.MULTILINE)
+_UNHOSTED_IMAGE_RE = re.compile(r"!\[([^\]]*)\]\((?!/files/|/private/files/)[^)]*\)")
 
 
 @dataclass
@@ -279,16 +281,20 @@ def drop_echoes_after_figures(markdown: str, page_text: str) -> str:
 	return re.sub(r"\n{3,}", "\n\n", "\n".join(line for index, line in enumerate(lines) if index not in drop))
 
 
-def drop_figure_echoes(pages: list[tuple[int, str]], pdf_path: str) -> list[tuple[int, str]]:
+def text_layers(pdf_path: str) -> dict[int, str]:
+	with fitz.open(pdf_path) as pdf_document:
+		return {page.number + 1: page.get_text() for page in pdf_document}
+
+
+def drop_figure_echoes(pages: list[tuple[int, str]], page_texts: dict[int, str]) -> list[tuple[int, str]]:
 	"""Drop the labels and captions a parser copied out of a picture into the text after it. Only on pages
 	with a text layer to check against, since on a scanned page every line would look drawn."""
-	with fitz.open(pdf_path) as pdf_document:
-		cleaned = []
-		for page_no, markdown in pages:
-			page_text = pdf_document[page_no - 1].get_text() if 0 < page_no <= pdf_document.page_count else ""
-			if "![" in markdown and len(normalized(layer_text(page_text))) >= MIN_TEXT_LAYER_CHARS:
-				markdown = drop_echoes_after_figures(markdown, page_text)
-			cleaned.append((page_no, markdown))
+	cleaned = []
+	for page_no, markdown in pages:
+		page_text = page_texts.get(page_no, "")
+		if "![" in markdown and len(normalized(layer_text(page_text))) >= MIN_TEXT_LAYER_CHARS:
+			markdown = drop_echoes_after_figures(markdown, page_text)
+		cleaned.append((page_no, markdown))
 	return cleaned
 
 
@@ -341,10 +347,23 @@ def insert_at(markdown: str, position: int, tag: str) -> str:
 	return f"{markdown[:position].rstrip()}\n\n{tag}\n\n{markdown[position:].lstrip()}"
 
 
+def unhosted_images_to_tokens(markdown: str) -> str:
+	if FIGURE_TOKEN_RE.search(markdown):
+		return _UNHOSTED_IMAGE_RE.sub("", markdown)
+	numbers = count(1)
+
+	def to_token(match: re.Match) -> str:
+		title = " ".join(match.group(1).replace("]", ")").split())
+		return f"[[FIGURE {next(numbers)}{': ' + title if title else ''}]]"
+
+	return _UNHOSTED_IMAGE_RE.sub(to_token, markdown)
+
+
 def place_figures(markdown: str, figures: list[Figure], page_no: int) -> str:
 	if not figures:
 		return FIGURE_TOKEN_RE.sub("", markdown or "")
 	placed: set[int] = set()
+	markdown = unhosted_images_to_tokens(markdown or "")
 
 	def swap_token(match: re.Match) -> str:
 		number = int(match.group(1))

wikify/engine/loader/cleanup.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import re
+import unicodedata
 from collections import Counter
 from itertools import groupby
 
@@ -67,6 +68,8 @@
 MAX_TABLE_TITLE_LINES = 2
 MAX_TABLE_TITLE_CHARS = 120
 MIN_TRANSCRIBED_ROWS = 4
+TYPED_TABLE_WORD_SHARE = 0.6
+_TABLE_WORD = re.compile(r"[^\W_]{3,}")
 
 
 def _plain(line: str) -> str:
@@ -141,12 +144,22 @@ def _edge_lines(md: str) -> list[str]:
 def find_boilerplate(pages: list[tuple[int, str]]) -> set[str]:
 	"""Normalized lines that recur at the top or bottom edge of a large fraction of pages."""
 	counts: Counter[str] = Counter()
+	outer_counts: Counter[str] = Counter()
 	for _, md in pages:
-		for nl in {_norm(line) for line in _edge_lines(md)}:
+		edge = _edge_lines(md)
+		for nl in {_norm(line) for line in edge}:
 			counts[nl] += 1
+		outer = [line for line in edge if not _is_page_of(line) and not _PAGE_NUMBER.match(_plain(line))]
+		for nl in {_norm(line) for line in outer[:1] + outer[-1:]}:
+			outer_counts[nl] += 1
+	threshold = max(3, int(0.30 * len(pages)))
 	# Capped so a header that alternates between styles across a long document still counts.
-	threshold = max(3, min(int(0.30 * len(pages)), BOILERPLATE_MAX_PAGES))
-	return {line for line, c in counts.items() if c >= threshold and len(line) <= 90}
+	outer_threshold = min(threshold, BOILERPLATE_MAX_PAGES)
+	return {
+		line
+		for line, c in counts.items()
+		if len(line) <= 90 and (c >= threshold or outer_counts[line] >= outer_threshold)
+	}
 
 
 def _is_page_of(line: str) -> bool:
@@ -402,7 +415,17 @@ def _transcription_after_figure(lines: list[str], figure: int) -> tuple[str, int
 	return None
 
 
-def drop_transcribed_figures(md: str) -> str:
+def _typed_on_page(lines: list[str], start: int, rows: int, page_text: str) -> bool:
+	if not page_text:
+		return False
+	words = _TABLE_WORD.findall(unicodedata.normalize("NFKC", " ".join(lines[start : start + rows])).lower())
+	page_words = set(_TABLE_WORD.findall(unicodedata.normalize("NFKC", page_text).lower()))
+	if not words:
+		return False
+	return sum(word in page_words for word in words) >= TYPED_TABLE_WORD_SHARE * len(words)
+
+
+def drop_transcribed_figures(md: str, page_text: str = "") -> str:
 	"""Keep one copy of a picture the parser also transcribed: the crop for a diagram it redrew as mermaid,
 	the typed table for a table printed as an image, since that copy is exact and searchable."""
 	lines = md.splitlines()
@@ -414,17 +437,20 @@ def drop_transcribed_figures(md: str) -> str:
 		kind, start, end = transcription
 		if kind == "mermaid":
 			drop.update(range(start, end + 1))
-		elif end >= MIN_TRANSCRIBED_ROWS:
+		elif end >= MIN_TRANSCRIBED_ROWS and not _typed_on_page(lines, start, end, page_text):
 			drop.add(index)
 	if not drop:
 		return md
 	return "\n".join(line for index, line in enumerate(lines) if index not in drop)
 
 
-def clean_pages(pages: list[tuple[int, str]]) -> list[tuple[int, str]]:
+def clean_pages(
+	pages: list[tuple[int, str]], page_texts: dict[int, str] | None = None
+) -> list[tuple[int, str]]:
 	pages = [(page_no, html_tables_to_markdown(md)) for page_no, md in pages]
+	page_texts = page_texts or {}
 	stripped = [
-		(page_no, contents_lines_to_table(drop_transcribed_figures(md)))
+		(page_no, contents_lines_to_table(drop_transcribed_figures(md, page_texts.get(page_no, ""))))
 		for page_no, md in strip_boilerplate(pages, find_boilerplate(pages))
 	]
 	stitched = stitch_cross_page_tables(stripped)

wikify/engine/loader/page_plan.py

@@ -140,7 +140,7 @@ def enforce_invariants(
 		or index in child_counts
 		or (sections[index].markdown.strip() and child_counts[parents[index]] > 1)
 	}
-	return pages, parents
+	return pages or {0}, parents
 
 
 def add_numbered_siblings(
@@ -343,6 +343,7 @@ def fold_sections(sections: list[Section], parents: list[int | None], pages: set
 	depth_below_page: dict[int, int] = {}
 	page_parents = {parents[index] for index in pages}
 	leading: list[str] = []
+	leading_sections: set[int] = set()
 	for index, section in enumerate(sections):
 		if index in pages:
 			owner[index], depth_below_page[index] = index, 0
@@ -352,9 +353,10 @@ def fold_sections(sections: list[Section], parents: list[int | None], pages: set
 				leading = []
 			continue
 		parent = parents[index]
-		if parent is None and index - 1 not in owner:
+		if (parent is None and index - 1 not in owner) or parent in leading_sections:
 			title = "" if section.title == PREAMBLE_TITLE else f"**{section.title}**"
 			leading.extend(part for part in (title, section.markdown.strip()) if part)
+			leading_sections.add(index)
 			continue
 		if parent is None:
 			owner[index], depth_below_page[index] = owner[index - 1], 1

wikify/engine/loader/sectionizer.py

@@ -30,6 +30,7 @@
 	r"^(?:[-*+]\s+)?\|?\s*\d+(?:\.\d+)*\.?\s*\|?\s*[A-Za-z(].*?(?:\.{3,}|\s[-\u2013\u2014]\s|\||<br>|\s)\s*(\d{1,4})\s*\|?$"
 )
 _TOC_LEADER = re.compile(r"[A-Za-z].*\.{4,}\s*(\d{1,4})$")
+_TOC_HEADER_WORD = re.compile(r"(?i)\b(?:page|pg|pp|contents|index)\b")
 TOC_MIN_ENTRIES = 5
 TOC_MIN_ENTRY_SHARE = 0.6
 TOC_MIN_ASCENDING_SHARE = 0.75
@@ -313,13 +314,18 @@ def toc_end_line(lines: list[str]) -> int:
 	real headings, so reading them as headings opens phantom sections and advances chapter numbering."""
 	entries: list[tuple[int, int]] = []
 	content_lines = 0
+	in_data_table = False
 	for index, line in enumerate(lines):
 		text = line.strip()
 		if not text or _SEP_ONLY.match(text):
 			continue
 		content_lines += 1
 		entry = _TOC_ENTRY.match(text) or _TOC_LEADER.search(text)
-		if entry:
+		if not text.startswith("|"):
+			in_data_table = False
+		elif index + 1 < len(lines) and _SEP_ONLY.match(lines[index + 1].strip()):
+			in_data_table = not entry and not _TOC_HEADER_WORD.search(text)
+		if entry and not in_data_table:
 			entries.append((index, int(entry.group(1))))
 	if len(entries) < TOC_MIN_ENTRIES or len(entries) < TOC_MIN_ENTRY_SHARE * content_lines:
 		return -1
@@ -386,6 +392,8 @@ def flush():
 		# raw HTML, or a lazy continuation of a list item that ended the previous page.
 		if buf and buf[-1].strip():
 			buf.append("")
+		if current is not None and not md.strip():
+			current.page_end = page_no
 		lines = md.splitlines()
 		toc_end = toc_end_line(lines)
 		first_line = next((index for index, line in enumerate(lines) if line.strip()), -1)

wikify/engine/loader/table_stitch.py

@@ -24,6 +24,7 @@
 _WRAPPER_TAGS = ("thead", "tbody", "tfoot", "span", "p", "u")
 _LIST_TAGS = ("ul", "ol")
 MIN_NUMBERED_ROWS = 2
+_SUB_STEP = re.compile(r"^\(?(?:[a-z]|[ivx]{1,4})[.)]?$")
 
 
 class _TableReader(HTMLParser):
@@ -194,7 +195,7 @@ def _continues(row: str, after_page_break: bool) -> bool:
 	"""A row that carries on the previous step instead of starting one: no step number and, in the
 	first non-empty cell, a lowercase start (or, right after a page break, any non-capital start)."""
 	cells = [_plain_cell(cell) for cell in _cells(row)]
-	if _STEP.match(cells[0]):
+	if _STEP.match(cells[0]) or _SUB_STEP.match(cells[0]):
 		return False
 	if after_page_break:
 		text = next((cell for cell in cells if cell), "")
@@ -284,27 +285,36 @@ def _with_blank_cell(row: str, column: int) -> str:
 
 def stitch_cross_page_tables(pages: list[tuple[int, str]]) -> list[tuple[int, str]]:
 	out = [[pno, md] for pno, md in pages]
-	for k in range(len(out) - 1):
-		a = _trailing_table(out[k][1])
-		b = _leading_table(out[k + 1][1])
-		if not a or not b:
+	target = 0
+	for k in range(1, len(out)):
+		if not out[k][1].strip():
 			continue
-		a_before, a_tbl = a
-		b_tbl, b_after = b
-		if not a_tbl or not b_tbl:
-			continue
-		dropped_column = _empty_column(a_tbl) if _ncols(b_tbl[0]) == _ncols(a_tbl[0]) - 1 else None
-		if _ncols(a_tbl[0]) != _ncols(b_tbl[0]) and dropped_column is None:
-			continue
-		cont = b_tbl
-		if len(b_tbl) >= 2 and _SEP.match(b_tbl[1]):
-			repeated = _header_key(b_tbl[0]) in ("", _header_key(a_tbl[0]))
-			cont = b_tbl[2:] if repeated else [b_tbl[0], *b_tbl[2:]]
-		if dropped_column is not None:
-			cont = [_with_blank_cell(row, dropped_column) for row in cont]
-		if cont and not _SEP.match(a_tbl[-1]) and _is_numbered(a_tbl) and _continues(cont[0], True):
-			a_tbl = [*a_tbl[:-1], _merge_rows(a_tbl[-1], cont[0])]
-			cont = cont[1:]
-		out[k][1] = "\n".join(a_before + a_tbl + cont).strip()
-		out[k + 1][1] = "\n".join(b_after).strip()
+		_stitch_pair(out, target, k)
+		if out[k][1].strip():
+			target = k
 	return [(pno, md) for pno, md in out]
+
+
+def _stitch_pair(out: list[list], previous: int, current: int) -> None:
+	a = _trailing_table(out[previous][1])
+	b = _leading_table(out[current][1])
+	if not a or not b:
+		return
+	a_before, a_tbl = a
+	b_tbl, b_after = b
+	if not a_tbl or not b_tbl:
+		return
+	dropped_column = _empty_column(a_tbl) if _ncols(b_tbl[0]) == _ncols(a_tbl[0]) - 1 else None
+	if _ncols(a_tbl[0]) != _ncols(b_tbl[0]) and dropped_column is None:
+		return
+	cont = b_tbl
+	if len(b_tbl) >= 2 and _SEP.match(b_tbl[1]):
+		repeated = _header_key(b_tbl[0]) in ("", _header_key(a_tbl[0]))
+		cont = b_tbl[2:] if repeated else [b_tbl[0], *b_tbl[2:]]
+	if dropped_column is not None:
+		cont = [_with_blank_cell(row, dropped_column) for row in cont]
+	if cont and not _SEP.match(a_tbl[-1]) and _is_numbered(a_tbl) and _continues(cont[0], True):
+		a_tbl = [*a_tbl[:-1], _merge_rows(a_tbl[-1], cont[0])]
+		cont = cont[1:]
+	out[previous][1] = "\n".join(a_before + a_tbl + cont).strip()
+	out[current][1] = "\n".join(b_after).strip()

wikify/engine/sectionize.py

@@ -32,7 +32,10 @@
 def sectionize_document(source_document: str, pdf_path: str, project_context: str = "") -> int:
 	"""Rebuild the Source Section tree from the doc's canonical pages. Returns the count."""
 	level_map = toc_level_map(str(pdf_path))
-	pages = clean_pages(images.drop_figure_echoes(store.get_canonical_pages(source_document), pdf_path))
+	page_texts = images.text_layers(pdf_path)
+	pages = clean_pages(
+		images.drop_figure_echoes(store.get_canonical_pages(source_document), page_texts), page_texts
+	)
 	sections = plan_pages(
 		sectionize(pages, level_map),
 		project_context,

wikify/tests/test_images.py

@@ -131,6 +131,27 @@ def test_inserts_a_figure_the_vlm_left_out_after_its_preceding_text(self):
 			f"## Admissions\n\n{TEXT_BEFORE}\n\n![Figure 2.1](/private/files/chart.png)\n\n{TEXT_AFTER}",
 		)
 
+	def test_an_image_tag_the_vlm_made_up_becomes_the_crop(self):
+		markdown = f"{TEXT_BEFORE}\n\n![Figure 2. Admissions per year](image2.png)\n\n{TEXT_AFTER}"
+
+		placed = images.place_figures(markdown, [chart_figure()], 2)
+
+		self.assertEqual(
+			placed,
+			f"{TEXT_BEFORE}\n\n![Figure 2. Admissions per year](/private/files/chart.png)\n\n{TEXT_AFTER}",
+		)
+
+	def test_a_made_up_image_tag_beside_a_token_is_dropped(self):
+		markdown = (
+			f"{TEXT_BEFORE}\n\n[[FIGURE 1]]\n\n![chart](attachment:chart.png)\n\n![Logo](/files/logo.png)"
+		)
+
+		placed = images.place_figures(markdown, [chart_figure()], 2)
+
+		self.assertNotIn("attachment:chart.png", placed)
+		self.assertIn("![Logo](/files/logo.png)", placed)
+		self.assertEqual(placed.count("/private/files/chart.png"), 1)
+
 	def test_a_figure_already_in_the_markdown_is_not_added_twice(self):
 		markdown = f"{TEXT_BEFORE}\n\n![Admissions per year](/private/files/chart.png)\n\n{TEXT_AFTER}"
 
@@ -171,6 +192,16 @@ def test_keeps_a_chart_followed_by_a_paragraph_and_a_table(self):
 
 		self.assertEqual(drop_transcribed_figures(markdown), markdown)
 
+	def test_keeps_a_chart_when_the_table_after_it_is_typed_on_the_page(self):
+		markdown = (
+			"![Figure 1](/private/files/chart.png)\n\nFigure 2. eGFR trajectory\n\n"
+			"| Outcome | HR |\n|---|---|\n| eGFR decline | 0.61 |\n| Death | 0.69 |"
+		)
+		page_text = "Results\nOutcome HR\neGFR decline 0.61\nDeath 0.69\nTable 1. Hazard ratios."
+
+		self.assertEqual(drop_transcribed_figures(markdown, page_text), markdown)
+		self.assertNotIn("chart.png", drop_transcribed_figures(markdown, "Results"))
+
 
 class TestFigureEchoes(FrappeTestCase):
 	def test_drops_labels_read_off_the_picture_and_keeps_the_printed_text(self):
@@ -196,7 +227,7 @@ def test_drops_labels_read_off_the_picture_and_keeps_the_printed_text(self):
 			"Assemble needle and manometer. Attach the 3-way stopcock to manometer"
 		)
 
-		[(_, cleaned)] = images.drop_figure_echoes([(1, markdown)], str(path))
+		[(_, cleaned)] = images.drop_figure_echoes([(1, markdown)], images.text_layers(str(path)))
 
 		self.assertEqual(
 			cleaned,

wikify/tests/test_page_plan.py

@@ -115,6 +115,17 @@ def test_heading_depth_is_capped_at_six(self):
 		self.assertEqual(child.title, "Level 2")
 		self.assertIn("###### Level 7\n\nbody of Level 7\n\n###### Level 8", child.markdown)
 
+	def test_a_short_memo_of_unnumbered_headings_keeps_its_text(self):
+		sections = [
+			_sec(["Preamble"], 1, markdown="Memo to all renal staff."),
+			_sec(["Clinic closure"], 1, markdown="The renal clinic is closed on Friday."),
+			_sec(["Contact"], 1, markdown="Call extension 4455."),
+		]
+		pages = plan_pages(sections)
+		self.assertEqual(_titles(pages), ["Preamble"])
+		for text in ("Memo to all renal staff.", "## Clinic closure", "Call extension 4455."):
+			self.assertIn(text, pages[0].markdown)
+
 	def test_hierarchy_of_the_remaining_pages_is_unchanged(self):
 		pages = plan_pages(_outline())
 		self.assertEqual(pages[-1].hierarchy_path, ["6.1 Protocols", "6.1.2 Lupus nephritis"])
@@ -422,6 +433,18 @@ def test_sections_sharing_a_pdf_page_share_a_wiki_page(self):
 		for title in ("2.5 Research staff", "2.6 Office staff", "2.7 Support staff"):
 			self.assertIn(f"## {title}", pages[-1].markdown)
 
+	def test_children_of_an_unpicked_leading_section_fold_with_it(self):
+		sections = [
+			_sec(["Preamble"], 1, markdown=_words(10)),
+			_sec(["Introduction"], 1, markdown=_words(5)),
+			_sec(["Introduction", "Background"], 1, markdown="background " + _words(200)),
+			_sec(["Introduction", "Scope"], 2, markdown="scope " + _words(200)),
+		]
+		pages, _ = self._plan(_llm_reply({"pages": [2, 3]}), sections)
+		self.assertEqual(len(pages), 1)
+		self.assertIn("background", pages[0].markdown)
+		self.assertIn("scope", pages[0].markdown)
+
 	def test_a_long_numbered_section_keeps_its_page_when_it_starts_mid_page(self):
 		sections = [
 			_sec(["3. Job descriptions"], 34, markdown=_words(80)),