Coverage for src / quber / core / extractors / set_of_mark / pipeline.py: 56%

1"""Set-of-Mark vision-guided extraction as a pydantic-graph. 

2 

3Topology (mirrors the GraphBuilder shape of the LLM pipeline; every map edge 

4pairs with its own closest-fork join): 

5 

6 start -> locate_tables -> [map] capture_table -> (join) -> finalize -> end 

7 

8- locate_tables: render-fed page images in, the grid locator names every table 

9 and its region; emits one ref per table. 

10- capture_table (mapped per table): region-constrained Camelot extracts the 

11 cell grid inside the table's box, then the grounded structure-correction step 

12 cleans it (spans, merged-symbol columns, multi-level headers) without 

13 touching values. 

14- finalize: sort to reading order. 

15 

16Vision owns identity and region; Camelot owns the values inside that region; 

17the LLM only cleans structure under the no-number guard. The legacy llm/ and 

18correspondence/ pipelines remain importable but are deprecated and dormant. 

19""" 

20 

21from __future__ import annotations 

22 

23import asyncio 

24from dataclasses import dataclass 

25from functools import cache 

26from pathlib import Path 

27from typing import Dict, List, Optional, Tuple 

28 

29from pydantic import BaseModel, ConfigDict 

30from pydantic_graph import GraphBuilder, StepContext, reduce_list_append 

31from pydantic_graph.graph_builder import Graph 

32from pydantic_graph.id_types import JoinID, NodeID 

33 

34from quber.agents.capture_advisor import CaptureAdvisor 

35from quber.agents.grid_locator import GridLocator, LocatedTable 

36from quber.agents.llm_client import LLMClient 

37from quber.agents.status_inspector import StatusInspector 

38from quber.core.extractors.base import ExtractedTable, table_address 

39from quber.core.extractors.camelot.correspondence.correction import printed_title 

40from quber.core.extractors.camelot.correspondence.geometry import norm_bbox_to_table_area 

41from quber.core.extractors.camelot.correspondence.recovery import camelot_targeted 

42from quber.core.extractors.camelot.recapture import repair_capture 

43from quber.core.extractors.set_of_mark.assemble import CamelotOrigin, TableAssembly, assemble_table 

44from quber.core.extractors.set_of_mark.extent import apply_content_regions 

45 

46 

47@dataclass 

48class SetOfMarkDeps: 

49 """Run-scoped dependencies and the pre-rendered page images.""" 

50 

51 locator: GridLocator 

52 llm: LLMClient 

53 source: Path 

54 page_images: Dict[int, Path] 

55 page_dims: Dict[int, Tuple[float, float]] 

56 dpi: int 

57 correct_sem: asyncio.Semaphore 

58 page_sem: asyncio.Semaphore 

59 # Recommends a Camelot retry when a capture drop is detected; None 

60 # disables the repair loop (backend `off`). 

61 advisor: Optional[CaptureAdvisor] = None 

62 # Verifies each unboxed cell's proposed status against the table image; 

63 # None disables inspection (backend `off`) and proposed statuses stand. 

64 inspector: Optional[StatusInspector] = None 

65 

66 

67@dataclass 

68class SetOfMarkState: 

69 """No cross-step mutable state; everything run-scoped rides on deps.""" 

70 

71 

72class LocatedTableRef(BaseModel): 

73 """One located table threaded into the per-table map stage.""" 

74 

75 model_config = ConfigDict(arbitrary_types_allowed=True) 

76 page: int 

77 located: LocatedTable 

78 

79 

80class CapturedTable(BaseModel): 

81 """A captured table plus its page/ordinal, so finalize can reading-order it.""" 

82 

83 page: int 

84 ordinal: int 

85 table: ExtractedTable 

86 

87 

88async def locate_tables(ctx: StepContext[SetOfMarkState, SetOfMarkDeps, Path]) -> List[LocatedTableRef]: 

89 deps = ctx.deps 

90 

91 async def per_page(page: int) -> List[LocatedTableRef]: 

92 async with deps.page_sem: 

93 located = await deps.locator.locate(deps.page_images[page], deps.source, page) 

94 return [LocatedTableRef(page=page, located=t) for t in sorted(located, key=lambda t: t.ordinal)] 

95 

96 per_page_refs = await asyncio.gather(*(per_page(p) for p in sorted(deps.page_images))) 

97 return [ref for page_refs in per_page_refs for ref in page_refs] 

98 

99 

100async def capture_table(ctx: StepContext[SetOfMarkState, SetOfMarkDeps, LocatedTableRef]) -> CapturedTable: 

101 deps = ctx.deps 

102 ref = ctx.inputs 

103 page, located = ref.page, ref.located 

104 page_w, page_h = deps.page_dims[page] 

105 area = norm_bbox_to_table_area(located.region, page_w, page_h) 

106 try: 

107 cand = await asyncio.to_thread(camelot_targeted, str(deps.source), page, area, located.ordinal) 

108 except Exception: 

109 cand = None 

110 

111 if cand is None: 

112 # Region produced no grid: keep the visual identity, empty body. 

113 table = ExtractedTable( 

114 table_id=table_address(deps.source, page, located.ordinal), 

115 title=await printed_title(located.title, str(deps.source), page), 

116 markdown="", 

117 page=page, 

118 source=str(deps.source), 

119 flavor="stream", 

120 som_region=located.region, 

121 ) 

122 return CapturedTable(page=page, ordinal=located.ordinal, table=table) 

123 

124 # Capture-drop repair: values Camelot's cell assignment dropped (present in 

125 # the region's text layer, absent from the grid) trigger an advised retry, 

126 # accepted only when it provably captures more. See camelot/recapture.py. 

127 cand = await repair_capture( 

128 cand, str(deps.source), page, deps.page_images[page], deps.dpi, located.ordinal, deps.advisor 

129 ) 

130 

131 table = await assemble_table( 

132 TableAssembly(deps.source, deps.llm, deps.correct_sem, deps.dpi, deps.inspector), 

133 page=page, 

134 ordinal=located.ordinal, 

135 page_image=deps.page_images[page], 

136 page_dims=(page_w, page_h), 

137 cells=cand.cells, 

138 cell_boxes=cand.cell_boxes, 

139 bbox=cand.bbox, 

140 title=located.title, 

141 som_region=located.region, 

142 camelot=CamelotOrigin(cand.bbox, cand.flavor, cand.accuracy), 

143 ) 

144 return CapturedTable(page=page, ordinal=located.ordinal, table=table) 

145 

146 

147async def finalize( 

148 ctx: StepContext[SetOfMarkState, SetOfMarkDeps, List[CapturedTable]], 

149) -> List[ExtractedTable]: 

150 ordered = sorted(ctx.inputs, key=lambda c: (c.page, c.ordinal)) 

151 tables = [c.table for c in ordered] 

152 # The locator and Camelot both bound a table generously, so its box can run 

153 # past the last data row (e.g. enclosing footnote lines below the grid). 

154 # Record the true end of the table content from its corrected last row. 

155 await asyncio.to_thread(apply_content_regions, tables, str(ctx.deps.source)) 

156 return tables 

157 

158 

159@cache 

160def build_set_of_mark_graph() -> Graph[SetOfMarkState, SetOfMarkDeps, Path, List[ExtractedTable]]: 

161 """Build the vision-guided extraction graph (cached; immutable).""" 

162 g = GraphBuilder( 

163 name="set_of_mark.pipeline", 

164 state_type=SetOfMarkState, 

165 deps_type=SetOfMarkDeps, 

166 input_type=Path, 

167 output_type=List[ExtractedTable], 

168 ) 

169 

170 locate_step = g.step(locate_tables) 

171 capture_step = g.step(capture_table) 

172 finalize_step = g.step(finalize) 

173 

174 j_captured = g.join( 

175 reduce_list_append, 

176 initial_factory=list, 

177 node_id="j_captured", 

178 preferred_parent_fork="closest", 

179 ) 

180 

181 g.add_edge(g.start_node, locate_step) 

182 g.add_mapping_edge(locate_step, capture_step, downstream_join_id=JoinID(NodeID("j_captured"))) 

183 g.add_edge(capture_step, j_captured) 

184 g.add_edge(j_captured, finalize_step) 

185 g.add_edge(finalize_step, g.end_node) 

186 

187 return g.build()