Coverage for src / quber / core / figures / dpt3 / digest.py: 94%
146 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Turn one page's dpt-3 response into records, and project them for the graft.
3A dpt-3 response has two parts: one long markdown string holding all the text
4the model read off the page, and a tree describing what sits where. Each tree
5node names its kind, gives its rectangle on the page, and points at the stretch
6of the long string that belongs to it — the node itself stores no text, so
7getting a node's text means cutting its stretch out of the string.
9`digest_page` walks the tree and produces the track's records: a figure node
10becomes a `DigestedFigure` with its kind, description and value tables kept
11separate; a table node and its cell nodes become a `DigestedTable` with a box
12on every cell; everything else on the page is kept as context. A response
13without the tree raises `UnrecognizedResponse` naming the stored file and the
14model, so a billed page can never pass as blank because the code did not
15recognize the format. A page node reporting anything but success raises too,
16with the response's own reason.
18`project_scan` turns a digest into the `PageScan` every downstream consumer
19already reads — the graft, the correction sweep, the value reconciliation. A
20figure is projected as its description followed by its value tables rendered
21as markdown; a table is projected grid-for-grid. The rest of the digest stays
22in the track models.
24The figure markup is the response's own labelling — the model wraps each
25figure's content in tags it emits itself — so reading it back is reading the
26format, not inferring anything about the page.
27"""
29from __future__ import annotations
31import re
32from html.parser import HTMLParser
33from typing import Any, Dict, List, Optional, Tuple
35from quber.core.figures.dpt3.models import (
36 DigestedFigure,
37 DigestedTable,
38 DigestedText,
39 FigureValueCell,
40 FigureValueTable,
41 PageDigest,
42)
43from quber.core.figures.models import Box, ChartContext, FigureRecord, PageScan, PageStatus, ScannedTable
46class UnrecognizedResponse(ValueError):
47 """A response in neither the dpt-2 nor the dpt-3 format."""
50#: The opening tag the response wraps a figure's content in, carrying its kind.
51_FIGURE_OPEN_RE = re.compile(r'<figure\b[^>]*\btype="([^"]*)"[^>]*>')
53#: One description block inside a figure's content.
54_DESCRIPTION_RE = re.compile(r"<description>(.*?)</description>", re.DOTALL)
56#: One value table inside a figure's content.
57_VALUE_TABLE_RE = re.compile(r"<table>(.*?)</table>", re.DOTALL)
60def response_generation(response: Dict[str, Any]) -> Optional[str]:
61 """Which format a raw response holds: `dpt2`, `dpt3`, or None for neither.
63 The format is visible in the response itself: a dpt-2 response carries a
64 flat `chunks` list, a dpt-3 response carries the `structure` tree.
65 """
66 if isinstance(response.get("chunks"), list):
67 return "dpt2"
68 if isinstance(response.get("structure"), dict):
69 return "dpt3"
70 return None
73def digest_page(
74 response: Dict[str, Any],
75 page: int,
76 model: str,
77 artifact: Optional[str] = None,
78) -> PageDigest:
79 """Digest one page's raw dpt-3 response into the track's records.
81 `page` is the 1-based source page the submission came from; the response
82 numbers its own single page from one, so the source page is carried in
83 rather than read back out of it. `artifact` names the stored response file
84 for error messages.
85 """
86 structure = response.get("structure")
87 if not isinstance(structure, dict):
88 raise UnrecognizedResponse(
89 f"{artifact or 'the response'} (model {model}) holds neither a dpt-2 chunks list "
90 "nor a dpt-3 structure tree"
91 )
92 markdown = response.get("markdown") or ""
93 meta = response.get("metadata") or {}
95 digest = PageDigest(
96 page=page,
97 job_id=meta.get("job_id"),
98 model=model,
99 version=meta.get("model_version") or meta.get("version"),
100 credits=(meta.get("billing") or {}).get("total_credits", meta.get("credit_usage")),
101 )
102 for page_node in structure.get("children") or []:
103 if page_node.get("type") != "page":
104 continue
105 status = page_node.get("status")
106 if status not in (None, "ok"):
107 raise ValueError(
108 f"{artifact or 'the response'} (model {model}) reports page status {status!r}"
109 f" ({page_node.get('reason')}); the page was billed and cannot pass as blank"
110 )
111 for node in page_node.get("children") or []:
112 kind = node.get("type")
113 if kind == "figure":
114 digest.figures.append(_figure(node, markdown, page))
115 elif kind == "table":
116 digest.tables.append(_table(node, markdown, page))
117 else:
118 digest.context.append(
119 DigestedText(
120 page=page,
121 kind=kind or "",
122 text=_slice(node, markdown),
123 box=_box(node),
124 node_id=node.get("id"),
125 )
126 )
127 return digest
130def project_scan(
131 digest: PageDigest,
132 picture_classes: List[str],
133 response_artifact: Optional[str] = None,
134 reused: bool = False,
135) -> PageScan:
136 """The digest as the `PageScan` every downstream consumer already reads.
138 The figure records carry text, id and rectangle — the graft's whole
139 contract. The digest keeps everything else.
140 """
141 figures = [
142 FigureRecord(
143 page=digest.page,
144 text=figure_text(figure),
145 box=figure.box,
146 chunk_id=figure.node_id,
147 job_id=digest.job_id,
148 )
149 for figure in digest.figures
150 ]
151 tables = [
152 ScannedTable(
153 page=digest.page,
154 cells=table.cells,
155 cell_boxes=table.cell_boxes,
156 box=table.box,
157 chunk_id=table.node_id,
158 job_id=digest.job_id,
159 )
160 for table in digest.tables
161 ]
162 context = [
163 ChartContext(page=digest.page, kind=item.kind, text=item.text, box=item.box, chunk_id=item.node_id)
164 for item in digest.context
165 if item.text.strip()
166 ]
167 status: PageStatus = "figures" if figures else "tables" if tables else "empty"
168 return PageScan(
169 page=digest.page,
170 status=status,
171 picture_classes=picture_classes,
172 job_id=digest.job_id,
173 model=digest.model,
174 version=digest.version,
175 credits=digest.credits,
176 reused=reused,
177 response_artifact=response_artifact,
178 figures=figures,
179 tables=tables,
180 context=context,
181 )
184def figure_text(figure: DigestedFigure) -> str:
185 """A figure's readable content: its description, then its value tables.
187 This is what the graft puts on the picture and what the correction agent
188 and the value reconciliation read, so the values are rendered as markdown
189 rows rather than left in the response's markup.
190 """
191 parts: List[str] = []
192 if figure.description.strip():
193 parts.append(figure.description.strip())
194 for table in figure.values:
195 rendered = _markdown_rows(table)
196 if rendered:
197 parts.append(rendered)
198 return "\n\n".join(parts)
201def _markdown_rows(table: FigureValueTable) -> str:
202 """A value table as pipe-delimited rows, one line per row."""
203 lines = []
204 for row in table.rows:
205 cells = [" ".join(cell.text.split()).replace("|", "/") for cell in row]
206 lines.append("| " + " | ".join(cells) + " |")
207 if len(lines) == 1 and len(table.rows) > 1:
208 lines.append("|" + " --- |" * len(row))
209 return "\n".join(lines)
212def _figure(node: Dict[str, Any], markdown: str, page: int) -> DigestedFigure:
213 """One figure node: kind, descriptions and value tables out of its stretch."""
214 content = _slice(node, markdown)
215 kind_match = _FIGURE_OPEN_RE.search(content)
216 descriptions = [d.strip() for d in _DESCRIPTION_RE.findall(content) if d.strip()]
217 tables = [_value_table(html) for html in _VALUE_TABLE_RE.findall(content)]
218 return DigestedFigure(
219 page=page,
220 kind=(kind_match.group(1).lower() if kind_match else ""),
221 description="\n\n".join(descriptions),
222 box=_box(node),
223 values=[t for t in tables if t.rows],
224 node_id=node.get("id"),
225 )
228def _table(node: Dict[str, Any], markdown: str, page: int) -> DigestedTable:
229 """One table node as a dense grid, cells and boxes read off its cell nodes.
231 A cell that spans is written at its own row and column; the positions it
232 covers stay blank and carry no box, so the grid never states a value twice.
233 """
234 placed: List[Tuple[int, int, str, Optional[Box]]] = []
235 rows = 0
236 columns = 0
237 for cell in node.get("children") or []:
238 if cell.get("type") != "table_cell":
239 continue
240 row = int(cell.get("row") or 0)
241 col = int(cell.get("col") or 0)
242 rows = max(rows, row + int(cell.get("rowspan") or 1))
243 columns = max(columns, col + int(cell.get("colspan") or 1))
244 placed.append((row, col, " ".join(_slice(cell, markdown).split()), _box(cell)))
246 grid = [["" for _ in range(columns)] for _ in range(rows)]
247 boxes: List[List[Optional[Box]]] = [[None for _ in range(columns)] for _ in range(rows)]
248 for row, col, text, box in placed:
249 if 0 <= row < rows and 0 <= col < columns:
250 grid[row][col] = text
251 boxes[row][col] = box
253 if not placed:
254 # A table node without cell nodes still states its grid in its stretch.
255 value_table = _value_table(_slice(node, markdown))
256 grid = [[cell.text for cell in row] for row in value_table.rows]
257 boxes = [[None for _ in row] for row in grid]
259 return DigestedTable(page=page, cells=grid, cell_boxes=boxes, box=_box(node), node_id=node.get("id"))
262def _value_table(html: str) -> FigureValueTable:
263 """One value table's rows, read out of the response's own table markup."""
264 parser = _GridText()
265 parser.feed(html)
266 parser.close()
267 return FigureValueTable(rows=[[FigureValueCell(text=cell) for cell in row] for row in parser.rows if row])
270class _GridText(HTMLParser):
271 """Collects `td`/`th` text row by row. Markup inside a cell is dropped and
272 its text kept, so a value typeset with a line break reads as one string."""
274 def __init__(self) -> None:
275 super().__init__(convert_charrefs=True)
276 self.rows: List[List[str]] = []
277 self._row: Optional[List[str]] = None
278 self._parts: Optional[List[str]] = None
280 def handle_starttag(self, tag: str, attrs: List[Tuple[str, Optional[str]]]) -> None:
281 if tag == "tr":
282 self._flush_row()
283 self._row = []
284 elif tag in ("td", "th"):
285 self._flush_cell()
286 self._parts = []
287 elif tag == "br" and self._parts is not None:
288 self._parts.append(" ")
290 def handle_endtag(self, tag: str) -> None:
291 if tag in ("td", "th"):
292 self._flush_cell()
293 elif tag == "tr":
294 self._flush_row()
296 def handle_data(self, data: str) -> None:
297 if self._parts is not None:
298 self._parts.append(data)
300 def close(self) -> None:
301 super().close()
302 self._flush_row()
304 def _flush_cell(self) -> None:
305 if self._parts is None:
306 return
307 if self._row is None:
308 self._row = []
309 self._row.append(" ".join("".join(self._parts).split()))
310 self._parts = None
312 def _flush_row(self) -> None:
313 self._flush_cell()
314 if self._row is not None:
315 self.rows.append(self._row)
316 self._row = None
319def _slice(node: Dict[str, Any], markdown: str) -> str:
320 """A node's text: the stretch of the response's one long string it points at."""
321 grounding = node.get("grounding") or {}
322 rng = grounding.get("range") or {}
323 start, end = rng.get("start"), rng.get("end")
324 if start is None or end is None:
325 return ""
326 return markdown[start:end]
329def _box(node: Dict[str, Any]) -> Optional[Box]:
330 """A node's rectangle in the workflow's own box keys, or None without one."""
331 grounding = node.get("grounding") or {}
332 box = grounding.get("box") or {}
333 try:
334 return {
335 "left": box["xmin"],
336 "top": box["ymin"],
337 "right": box["xmax"],
338 "bottom": box["ymax"],
339 }
340 except KeyError:
341 return None