Coverage for src / quber / core / fusion / heading_review.py: 92%

1"""Document heading review — demote page decoration the parser called a heading. 

2 

3The parser labels certain printed lines as section headings, and downstream 

4consumers treat a heading as governing every chunk beneath it. When the label 

5is wrong — a banner the page layout repeats at the top of its pages, or a 

6table column label that leaked into the heading stream — that false context 

7spreads across a large share of the document's text. 

8 

9Demotion is deliberately hard, because wrongly demoting a real heading strips 

10section context from everything under it. Two witnesses must agree: 

11 

121. Geometry nominates: the same text appears as a heading on 

13 ``NOMINATION_PAGES`` or more pages, computed from the parse. 

142. The heading-review agent confirms, judging only the nominated texts. 

15 

16A heading that appears on fewer pages cannot be demoted by any verdict, and 

17an unavailable review demotes nothing. Every demotion is returned as a flag 

18record naming the text, its pages, the verdict and reason, and how many text 

19items it governed — the run's record that the relabel happened and why. 

20""" 

21 

22from __future__ import annotations 

23 

24from collections import defaultdict 

25from typing import TYPE_CHECKING, Dict, List, Set 

26 

27from docling_core.types.doc.labels import DocItemLabel 

28from loguru import logger 

29 

30from quber.agents.llm_client import LLMClient 

31from quber.core.extractors.base import CellFlag 

32 

33if TYPE_CHECKING: 

34 from docling_core.types.doc.document import DoclingDocument 

35 

36# A real heading is printed once, where its section starts. Two occurrences 

37# can still be a legitimate continuation heading; from three pages up the 

38# repetition pattern is layout, not authorship, and the text goes to review. 

39NOMINATION_PAGES = 3 

40 

41# How far, as a share of page height, a heading may sit from a demoted running 

42# header's position band and still be that running header's first half. 

43SPLIT_HEADER_TOLERANCE = 0.02 

44 

45 

46async def review_headings(document: "DoclingDocument", llm: LLMClient, source: str) -> List[CellFlag]: 

47 """Demote confirmed page decoration from the document's heading stream. 

48 

49 Relabels demoted headings to ``page_header`` in place on ``document`` and 

50 returns one flag record per demotion. Returns an empty list when nothing 

51 is nominated, the review backend does not support it, or the agent keeps 

52 everything. 

53 """ 

54 heading_pages: Dict[str, Set[int]] = defaultdict(set) 

55 heading_tops: Dict[str, List[float]] = defaultdict(list) 

56 heading_count: Dict[str, int] = defaultdict(int) 

57 page_heights = {no: p.size.height or 1.0 for no, p in (document.pages or {}).items()} 

58 for item in document.texts: 

59 if item.label == DocItemLabel.SECTION_HEADER and item.text.strip(): 

60 page = item.prov[0].page_no if item.prov else 0 

61 heading_pages[item.text.strip()].add(page) 

62 heading_count[item.text.strip()] += 1 

63 if item.prov and item.prov[0].bbox is not None: 

64 # Distance from the page top, 0..1. Docling boxes measure up 

65 # from the page bottom, so the top edge is height minus t. 

66 h = page_heights.get(page, 1.0) 

67 heading_tops[item.text.strip()].append(max(0.0, (h - item.prov[0].bbox.t) / h)) 

68 

69 nominated = {text: pages for text, pages in heading_pages.items() if len(pages) >= NOMINATION_PAGES} 

70 if not nominated: 

71 return [] 

72 

73 def position_evidence(text: str) -> str: 

74 tops = heading_tops.get(text) 

75 if not tops: 

76 return "position unknown" 

77 lo, hi = min(tops), max(tops) 

78 band = f"{lo:.0%}-{hi:.0%} down the page" if hi - lo > 0.05 else f"{lo:.0%} down the page" 

79 fixed = "a fixed position" if hi - lo <= 0.05 else "varying positions" 

80 return f"printed at {fixed}, {band}" 

81 

82 lines = [ 

83 f'- "{text}" — {heading_count[text]} occurrence(s) as a heading across {len(pages)} pages ' 

84 f"({', '.join(str(p) for p in sorted(pages))}); {position_evidence(text)}" 

85 for text, pages in sorted(nominated.items(), key=lambda kv: -len(kv[1])) 

86 ] 

87 review = await llm.review_headings("Nominated headings of this document:\n\n" + "\n".join(lines)) 

88 if review is None: 

89 logger.warning( 

90 "heading review unavailable on this backend; keeping all {} nominated heading(s)", 

91 len(nominated), 

92 ) 

93 return [] 

94 

95 demoted = { 

96 v.text.strip(): v 

97 for v in review.verdicts 

98 if v.verdict != "section_heading" and v.text.strip() in nominated 

99 } 

100 if not demoted: 

101 return [] 

102 

103 # Blast radius: how many text items each demoted heading governs, counted 

104 # before the relabel with the same walk consumers use — every non-heading 

105 # item inherits the most recent heading above it. 

106 governed: Dict[str, int] = defaultdict(int) 

107 current = "" 

108 for item in document.texts: 

109 if item.label == DocItemLabel.SECTION_HEADER: 

110 current = item.text.strip() 

111 elif current in demoted: 

112 governed[current] += 1 

113 

114 flags: List[CellFlag] = [] 

115 for text, verdict in demoted.items(): 

116 relabeled = 0 

117 for item in document.texts: 

118 if item.label == DocItemLabel.SECTION_HEADER and item.text.strip() == text: 

119 # SectionHeaderItem pins its label type, so the relabel goes 

120 # around assignment validation; the serialized artifact then 

121 # carries page_header like any other furniture text. 

122 object.__setattr__(item, "label", DocItemLabel.PAGE_HEADER) 

123 relabeled += 1 

124 pages = sorted(nominated[text]) 

125 logger.info( 

126 "heading demoted ({}): {!r} — a heading on {} pages, governed {} text item(s)", 

127 verdict.verdict, 

128 text[:60], 

129 len(pages), 

130 governed[text], 

131 ) 

132 flags.append( 

133 CellFlag( 

134 source=source, 

135 page=pages[0], 

136 title=text, 

137 text=text, 

138 status="heading_demoted", 

139 note=( 

140 f"{verdict.verdict} on {len(pages)} pages " 

141 f"({', '.join(str(p) for p in pages)}); relabeled {relabeled} heading " 

142 f"item(s) governing {governed[text]} text item(s); {verdict.reason}" 

143 ), 

144 ) 

145 ) 

146 

147 # A running header the layout split in two. The parser labels its first 

148 # half a heading on one page, and that half is also a real heading elsewhere 

149 # (a company name on the cover and above each statement), so the per-text 

150 # verdict cannot demote it without demoting the real ones. The occurrence is 

151 # judged by where it is printed instead: a heading whose text opens a 

152 # demoted running header, printed where that running header sits, is the 

153 # running header. 

154 bands = { 

155 text: (min(heading_tops[text]), max(heading_tops[text])) 

156 for text, verdict in demoted.items() 

157 if verdict.verdict == "running_header" and heading_tops.get(text) 

158 } 

159 for item in document.texts: 

160 if item.label != DocItemLabel.SECTION_HEADER or not item.prov or item.prov[0].bbox is None: 

161 continue 

162 text = item.text.strip() 

163 page = item.prov[0].page_no 

164 h = page_heights.get(page, 1.0) 

165 top = max(0.0, (h - item.prov[0].bbox.t) / h) 

166 for header, (lo, hi) in bands.items(): 

167 if not (header.startswith(text) and len(text) < len(header)): 

168 continue 

169 if not (lo - SPLIT_HEADER_TOLERANCE <= top <= hi + SPLIT_HEADER_TOLERANCE): 

170 continue 

171 object.__setattr__(item, "label", DocItemLabel.PAGE_HEADER) 

172 logger.info( 

173 "heading demoted (split running header): {!r} on page {} at {:.0%} down the page opens {!r}", 

174 text[:60], 

175 page, 

176 top, 

177 header[:60], 

178 ) 

179 flags.append( 

180 CellFlag( 

181 source=source, 

182 page=page, 

183 title=text, 

184 text=text, 

185 status="heading_demoted", 

186 note=( 

187 f"first half of the running header {header!r}, printed on page {page} " 

188 f"at {top:.0%} down the page where that header sits ({lo:.0%}-{hi:.0%})" 

189 ), 

190 ) 

191 ) 

192 break 

193 return flags