Coverage for src / quber / core / figures / scan.py: 87%

1"""Scan one page: submit it whole, read the figures and the page text back. 

2 

3A nominated page is submitted on its own, as a one-page PDF. Billing is per 

4submitted page, so scanning six pages one at a time costs exactly what 

5submitting the six together would, and submitting per page is what makes the 

6response per page: one submission covering several pages returns a single 

7payload for all of them, and the fields that identify a scan — the job, the 

8model version, the credits — are document-level and would not divide. 

9 

10The whole page goes in, not a crop of the chart. A chart's title is printed 

11above it and its note at the foot of the page, and both are only related to the 

12chart by where they sit, so cropping to the chart discards them and tiling 

13several charts onto one composite page destroys the arrangement that relates 

14them. 

15 

16`page_scan` turns one response into the record for that page. Every figure 

17becomes a chart record, every table becomes a grid with a box on each of its 

18cells, and every text and marginalia chunk is kept alongside them — that is where 

19the titles and the notes are. 

20 

21Nothing filters on what a chunk contains, so a figure that is not a chart still 

22becomes a record. A page's logo comes back as a figure of its own, and on a 

23branded deck that means a record per page on top of the charts. Reading the text 

24to decide which figures are charts would mean writing a rule against 

25descriptions that have no fixed shape, and the cost of a wrong rule is a 

26silently discarded chart. 

27""" 

28 

29from __future__ import annotations 

30 

31import asyncio 

32import tempfile 

33from pathlib import Path 

34from typing import Any, List, Optional 

35 

36from loguru import logger 

37 

38from quber.core.figures.grid import scanned_tables 

39from quber.core.figures.models import ChartContext, FigureRecord, PageScan, PageStatus 

40from quber.files.pdf import slice_page 

41from quber.providers.landing.client import parse_document 

42 

43#: The figure scan's own default, independent of the parse default the ADE 

44#: client carries, and pinned to a dated version: the vendor's floating 

45#: aliases repoint silently, and reuse is keyed by generation rather than 

46#: version, so a drifting default would mix versions within one document 

47#: with nothing flagging it. Bump the pin deliberately, after evaluating the 

48#: new version. The dpt-2 line stays fully selectable by passing it as 

49#: `model`; each generation only ever consumes its own stored responses. 

50SCAN_DEFAULT_MODEL = "dpt-3-pro-20260710" 

51 

52#: A one-page scan has come back in about five seconds, so it is polled more 

53#: often than a whole document, whose scan runs for minutes. 

54PAGE_POLL_SECONDS = 2.0 

55 

56#: Content-level instruction for the figures on the page. It cannot change the 

57#: shape of what comes back — two attempts to force an HTML table returned 

58#: identical text — but it does tell the model what to expect on the page. A 

59#: superscript on a chart title routinely points at a note printed on a 

60#: different page, and saying so keeps the model from inventing a definition it 

61#: cannot see. 

62CHART_FIGURE_PROMPT = ( 

63 "This page comes from a financial document. If the figure is a chart, read its " 

64 "plotted values, its axis labels and units, and its legend. A number or symbol in " 

65 "superscript refers to a note that is often printed on another page and so is not " 

66 "visible here; report the superscript as printed and do not guess what it refers to." 

67) 

68 

69#: Chunk types kept beside the charts. `text` carries the chart titles, which the 

70#: scan returns as their own chunk with any superscript intact; `marginalia` 

71#: carries the notes printed at the foot of the page. 

72CONTEXT_TYPES = ("text", "marginalia") 

73 

74 

75async def scan_page( 

76 source: Path, 

77 page: int, 

78 model: str = SCAN_DEFAULT_MODEL, 

79 poll_seconds: float = PAGE_POLL_SECONDS, 

80) -> dict[str, Any]: 

81 """Submit page `page` of `source` as a one-page PDF and return the raw response.""" 

82 with tempfile.TemporaryDirectory(prefix="quber-chart-") as tmp: 

83 submitted = await asyncio.to_thread(slice_page, source, page, Path(tmp) / f"p{page}.pdf") 

84 logger.info("Figure scan: submitting {} page {}", source.name, page) 

85 return await asyncio.to_thread( 

86 parse_document, 

87 submitted, 

88 model, 

89 {"figure": CHART_FIGURE_PROMPT}, 

90 poll_seconds, 

91 ) 

92 

93 

94def page_scan( 

95 response: dict[str, Any], 

96 page: int, 

97 model: str, 

98 picture_classes: List[str], 

99 response_artifact: Optional[str] = None, 

100 reused: bool = False, 

101) -> PageScan: 

102 """Build the record for one scanned page from its raw response. 

103 

104 `page` is the source page number the submission came from. The response 

105 numbers its own single page from zero, so the source page is carried in 

106 rather than read back out of it. 

107 """ 

108 meta = response.get("metadata") or {} 

109 job_id = meta.get("job_id") 

110 figures: List[FigureRecord] = [] 

111 context: List[ChartContext] = [] 

112 for chunk in response.get("chunks") or []: 

113 kind = chunk.get("type") 

114 grounding = chunk.get("grounding") or {} 

115 box = grounding.get("box") 

116 text = chunk.get("markdown") or "" 

117 if kind == "figure": 

118 figures.append( 

119 FigureRecord(page=page, text=text, box=box, chunk_id=chunk.get("id"), job_id=job_id) 

120 ) 

121 elif kind in CONTEXT_TYPES: 

122 context.append(ChartContext(page=page, kind=kind, text=text, box=box, chunk_id=chunk.get("id"))) 

123 

124 tables = scanned_tables(response, page, job_id) 

125 

126 credits = meta.get("credit_usage") 

127 if credits is None: 

128 credits = (meta.get("billing") or {}).get("total_credits") 

129 status: PageStatus = "figures" if figures else "tables" if tables else "empty" 

130 return PageScan( 

131 page=page, 

132 status=status, 

133 picture_classes=picture_classes, 

134 job_id=job_id, 

135 model=model, 

136 version=meta.get("version") or meta.get("model_version"), 

137 credits=credits, 

138 reused=reused, 

139 response_artifact=response_artifact, 

140 figures=figures, 

141 tables=tables, 

142 context=context, 

143 )