Coverage for src / quber / core / figures / dpt3 / footnotes.py: 94%
161 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Resolve figure footnotes from two sources and carry them into the document.
3The correction agent reads each scanned page's image and reports both ends of
4the association: the markers printed on the figures' labels and the note lines
5printed on the page. The graft records them on the pictures. This module ties
6marker to note and writes the result where a reader finds it.
8Two sources answer each marker. The agent's image read is one; the parse's own
9footnote-labelled text on the page — retained by the track's graft — is the
10other. Both are keyed by the same canonical marker, so `(1)`, `1` and a
11superscript one meet. Agreement corroborates the note. Disagreement, an
12unresolved marker, and a note no marker points at all become review flags on
13the run, never silent drops. A marker the agent judged a cross-reference to a
14named section resolves to a heading pointer, the same way a table's does.
16Notes are pooled across a page's figures because one printed note block serves
17several charts; the marker on each figure decides which notes are its.
19The document carries the result the way it does for tables: each resolved note
20is inserted as a footnote node right after the picture, so a marker inside a
21chart and its note at the page foot are linked in the document itself. The
22parse's own footnote text stays where the page prints it.
24Placement then goes to the value, not just the chart. The agent reports which
25printed label carries each marker, and that label is matched against the
26figure's value table; the note attaches to the matching cells' records in the
27digest. A marker whose label matches no value is flagged unplaced, never
28dropped — the note still sits on the picture.
29"""
31from __future__ import annotations
33from typing import Any, Dict, List, Optional, Sequence, Tuple
35from docling_core.types.doc.document import DoclingDocument, PictureItem
36from docling_core.types.doc.labels import DocItemLabel
37from loguru import logger
39from quber.core.figures.dpt3.models import DigestedFigure, PageDigest, PlacedNote, ResolvedNote
40from quber.core.figures.graft import FOOTNOTE_MARKS_FIELD, FOOTNOTES_FIELD
41from quber.core.figures.models import PageScan
42from quber.core.figures.values import label_words
43from quber.core.fusion.footnotes import canonical_marker, resolve_footnotes, split_leading_marker
46def resolve_figure_footnotes(
47 document: DoclingDocument,
48 scans: Sequence[PageScan],
49 digests: Dict[int, PageDigest],
50) -> List[str]:
51 """Resolve every scanned page's figure markers and write the results.
53 Mutates `document` — each picture's resolved notes are inserted after it —
54 and the digests, whose figures receive their resolutions and whose value
55 cells receive the notes placed on them. Returns the review flags.
56 """
57 flags: List[str] = []
58 headings = _headings(document)
59 by_ref = {picture.self_ref: picture for picture in document.pictures}
61 for scan in scans:
62 if not scan.figures:
63 continue
64 page = scan.page
65 digest_figures = {figure.node_id: figure for figure in _page_figures(digests, page) if figure.node_id}
66 # Both ends are read before anything is inserted: inserting footnote
67 # nodes adds footnote-labelled items to the page being read.
68 page_notes = _document_notes(document, page)
69 agent_notes, general_by_ref = _agent_notes(scan, by_ref)
70 referenced: set[str] = set()
71 # One resolution per marker per page. The same marker routinely sits on
72 # several labels and several figures — the note behind it is still one
73 # printed line, and an unresolvable marker is flagged once, not once
74 # per place it is printed.
75 resolutions: Dict[str, Optional[ResolvedNote]] = {}
77 for record in scan.figures:
78 picture = by_ref.get(record.picture_ref or "")
79 if picture is None:
80 continue
81 resolved: List[ResolvedNote] = list(general_by_ref.get(picture.self_ref, []))
82 figure = digest_figures.get(record.chunk_id or "")
83 mine: Dict[str, ResolvedNote] = {}
84 labels: Dict[str, List[str]] = {}
85 for mark in getattr(picture.meta, FOOTNOTE_MARKS_FIELD, None) or []:
86 marker = str(mark.get("marker") or "")
87 key = canonical_marker(marker)
88 if not key:
89 continue
90 referenced.add(key)
91 if key not in resolutions:
92 note, flag = _resolve(
93 marker,
94 key,
95 str(mark.get("kind") or "footnote"),
96 agent_notes,
97 page_notes,
98 headings,
99 page,
100 )
101 if flag:
102 flags.append(flag)
103 resolutions[key] = note
104 template = resolutions[key]
105 if template is None:
106 continue
107 # The figure's own copy: the same marker on another figure
108 # places against that figure's values independently.
109 if key not in mine:
110 mine[key] = template.model_copy(deep=True)
111 resolved.append(mine[key])
112 labels.setdefault(key, []).append(str(mark.get("label") or ""))
113 for key, note in mine.items():
114 flags.extend(_place_on_values(note, labels.get(key, []), figure, page))
115 if resolved:
116 _insert_notes(document, picture, resolved)
117 if figure is not None:
118 figure.footnotes = resolved
120 for key, (marker, _text) in agent_notes.items():
121 if key and key not in referenced:
122 flags.append(
123 f"page {page}: the page prints note {marker!r} and no marker on its "
124 "figures points at it"
125 )
127 return flags
130def _resolve(
131 marker: str,
132 key: str,
133 kind: str,
134 agent_notes: Dict[str, Tuple[str, str]],
135 page_notes: Dict[str, str],
136 headings: List[Tuple[str, int]],
137 page: int,
138) -> Tuple[Optional[ResolvedNote], Optional[str]]:
139 """One marker's resolution and, when something needs review, its flag."""
140 if kind == "section":
141 pointer = resolve_footnotes([marker], [], page, headings=headings, section_keys={key})
142 if pointer.resolved:
143 return ResolvedNote(marker=marker, text=pointer.resolved[0].text, sources=["headings"]), None
144 return None, (
145 f"page {page}: figure marker {marker!r} names a section of the document and no "
146 "heading opens with it"
147 )
149 agent = agent_notes.get(key)
150 printed = page_notes.get(key)
151 if agent is None and printed is None:
152 return None, (
153 f"page {page}: figure marker {marker!r} is defined by no note — neither the "
154 "agent's read of the page nor the page's own footnote text has it"
155 )
156 if agent is None:
157 return ResolvedNote(marker=marker, text=printed or "", sources=["document"]), None
158 if printed is None:
159 return ResolvedNote(marker=marker, text=agent[1], sources=["agent"]), None
160 if _agree(agent[1], printed):
161 return (
162 ResolvedNote(marker=marker, text=agent[1], sources=["agent", "document"], corroborated=True),
163 None,
164 )
165 return (
166 ResolvedNote(marker=marker, text=agent[1], sources=["agent", "document"]),
167 f"page {page}: note {marker!r} reads differently in the two sources — the agent read "
168 f"{agent[1]!r}, the page's own text says {printed!r}",
169 )
172def _agree(one: str, other: str) -> bool:
173 """Do two reads of one note say the same thing?
175 Whitespace and case are rendering; one read extending the other is the
176 same note cut at a different point, not a different note.
177 """
178 a = " ".join(one.split()).casefold().rstrip(".")
179 b = " ".join(other.split()).casefold().rstrip(".")
180 return bool(a) and bool(b) and (a == b or a in b or b in a)
183def _place_on_values(
184 note: ResolvedNote,
185 labels: Sequence[str],
186 figure: Optional[DigestedFigure],
187 page: int,
188) -> List[str]:
189 """Attach the note to the value cells whose text carries one of its labels.
191 A label the marker sits on is the join: every word of it (minus the marker
192 itself) must appear in the cell's text. The marker follows its label
193 whether the model drew one figure or two around it. A note none of its
194 labels place is flagged, and it stays on the picture; a marker printed on
195 a figure's title has no value to sit on, which is what the flag records.
196 """
197 if figure is None or not figure.values:
198 return []
199 tried = False
200 for label in labels:
201 words = label_words(label.replace(note.marker, ""))
202 if not words:
203 continue
204 tried = True
205 for table in figure.values:
206 for row in table.rows:
207 for cell in row:
208 if words <= label_words(cell.text):
209 cell.footnotes.append(PlacedNote(marker=note.marker, text=note.text))
210 note.placed = True
211 if note.placed or not tried:
212 return []
213 printed = ", ".join(repr(label) for label in labels if label)
214 return [
215 f"page {page}: note {note.marker!r} sits on {printed} and the figure's value table "
216 "has no value under that label; the note is on the picture only"
217 ]
220def _insert_notes(document: DoclingDocument, picture: PictureItem, notes: Sequence[ResolvedNote]) -> None:
221 """Insert each resolved note as a footnote node right after the picture.
223 The same carriage a table's footnotes get: a body node after the element
224 renders as a paragraph below it and keeps the reading order. The node
225 carries no provenance box, so nothing sweeping printed regions matches it.
227 A note already sitting after the picture is not inserted again, so a rerun
228 over an enriched parse — which is what reusing stored scans produces —
229 leaves the document as it was rather than doubling every note.
230 """
231 anchor, existing = _trailing_notes(document, picture)
232 for note in notes:
233 text = f"{note.marker} {note.text}".strip() if note.marker else note.text
234 if not text.strip() or text in existing:
235 continue
236 anchor = document.insert_text(label=DocItemLabel.FOOTNOTE, text=text, sibling=anchor, after=True)
237 logger.info(
238 "Figure footnotes: picture {} carries {} resolved note(s) ({})",
239 picture.self_ref,
240 len(notes),
241 ", ".join(note.marker or "unmarked" for note in notes),
242 )
245def _trailing_notes(document: DoclingDocument, picture: PictureItem) -> Tuple[Any, set[str]]:
246 """The run of footnote nodes already following the picture: the last one,
247 to anchor new insertions behind it, and their texts, to skip re-inserting.
249 Only a prov-less footnote belongs to the run — one this module inserted on
250 an earlier pass. A printed footnote the parse positioned on the page never
251 matches, so it is never mistaken for an insertion.
252 """
253 anchor: Any = picture
254 existing: set[str] = set()
255 parent = picture.parent.resolve(document) if picture.parent else None
256 if parent is None:
257 return anchor, existing
258 refs = [child.cref for child in parent.children]
259 if picture.self_ref not in refs:
260 return anchor, existing
261 for ref in refs[refs.index(picture.self_ref) + 1 :]:
262 item = document.texts[int(ref.rsplit("/", 1)[1])] if ref.startswith("#/texts/") else None
263 if item is None or item.label != DocItemLabel.FOOTNOTE or item.prov:
264 break
265 existing.add(item.text)
266 anchor = item
267 return anchor, existing
270def _agent_notes(
271 scan: PageScan, by_ref: Dict[str, PictureItem]
272) -> Tuple[Dict[str, Tuple[str, str]], Dict[str, List[ResolvedNote]]]:
273 """The notes the agent read off this page, pooled, plus the unmarked ones.
275 Marked notes pool across the page's figures — one printed note block
276 serves several charts — keyed by canonical marker, each keeping its
277 printed marker and text. An unmarked general note has no marker to pool
278 under and stays with the figure the agent read it for.
279 """
280 pooled: Dict[str, Tuple[str, str]] = {}
281 general: Dict[str, List[ResolvedNote]] = {}
282 for record in scan.figures:
283 picture = by_ref.get(record.picture_ref or "")
284 if picture is None:
285 continue
286 for entry in getattr(picture.meta, FOOTNOTES_FIELD, None) or []:
287 marker = str(entry.get("marker") or "")
288 text = str(entry.get("text") or "").strip()
289 if not text:
290 continue
291 key = canonical_marker(marker)
292 if key:
293 pooled.setdefault(key, (marker, text))
294 else:
295 general.setdefault(picture.self_ref, []).append(
296 ResolvedNote(marker="", text=text, sources=["agent"])
297 )
298 return pooled, general
301def _document_notes(document: DoclingDocument, page: int) -> Dict[str, str]:
302 """The page's own footnote text, keyed by the marker each line opens with.
304 These are the footnote-labelled items the parse produced and the track's
305 graft retained. A line opening with no marker defines nothing to key on.
306 """
307 notes: Dict[str, str] = {}
308 for item in document.texts:
309 if item.label != DocItemLabel.FOOTNOTE:
310 continue
311 if not item.prov or item.prov[0].page_no != page:
312 continue
313 parsed = split_leading_marker(item.text or "")
314 if parsed is not None:
315 notes.setdefault(parsed[0], parsed[1])
316 return notes
319def _headings(document: DoclingDocument) -> List[Tuple[str, int]]:
320 """The document's section headings with their pages, for section pointers."""
321 found: List[Tuple[str, int]] = []
322 for item in document.texts:
323 if item.label != DocItemLabel.SECTION_HEADER or not (item.text or "").strip():
324 continue
325 found.append((item.text, item.prov[0].page_no if item.prov else 0))
326 return found
329def _page_figures(digests: Dict[int, PageDigest], page: int) -> List[DigestedFigure]:
330 digest = digests.get(page)
331 return list(digest.figures) if digest is not None else []