Coverage for src / quber / core / figures / orchestrator.py: 50%
211 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Run the figure workflow end to end.
3The workflow is independent of the rest of the pipeline. It takes a parse and the
4source document, and gives back that same parse enriched, plus the page scans.
5Nothing about it depends on being called from inside a larger run, so a parse
6produced earlier can be enriched without reproducing it, and sequencing is a
7matter of calling the steps in order rather than a flag inside a monolithic run.
9The enriched parse is written back under the filename it was read from. The parse
10is an input every later stage reads by name, so enriching it in place means no
11stage has to be told which of two files to pick up. Run this after fusion, which
12is what makes the chart values reach the document the rest of the pipeline reads.
14The order: read the parse for the pages holding a chart, scan each of those pages
15whole and store the response as it came back, then carry each chart's text onto
16the picture already in the parse. Nomination is read purely off the parse and
17every nominated page is scanned; no judgement sits between the two. The `pages`
18parameter widens the first step: `all` sends every page of the document through
19the same flow, nomination info intact where it exists, for documents whose
20content the parse cannot be trusted to point at.
22A page is scanned once. Its stored response is what a later run reads instead of
23paying to scan the page again — provided the stored response's format matches
24the requested model's generation; a mismatched one is moved aside, never
25destroyed.
27The model selects the reading path. A dpt-3 model's responses go through the
28`dpt3` package — its digestion, its graft, its footnote resolution — while any
29other model runs the original path unchanged.
31Every nominated page reaches the output, including the pages dropped before
32scanning and the pages that scanned clean. A chart is often where a filing states
33a number that appears nowhere else, so a page that was nominated and produced
34nothing is exactly what someone needs to see.
35"""
37from __future__ import annotations
39import asyncio
40import json
41import re
42import tempfile
43from dataclasses import dataclass, field
44from pathlib import Path
45from typing import Any, Dict, List, Optional, Tuple
47from docling_core.types.doc.document import DoclingDocument
48from loguru import logger
50from quber.agents.cell_reader import CellReader, get_cell_reader
51from quber.agents.figure_correction import FigureCorrector, get_figure_corrector
52from quber.agents.llm_client import LLMClient, get_llm_client
53from quber.agents.status_inspector import StatusInspector, get_status_inspector
54from quber.core.extractors.base import ExtractedTable
55from quber.core.extractors.set_of_mark.assemble import TableAssembly
56from quber.core.figures import dpt3
57from quber.core.figures.capture import CAPTURE_DPI, capture_tables
58from quber.core.figures.correct import correct_figures
59from quber.core.figures.graft import graft_figures, graft_tables, unread_pictures
60from quber.core.figures.models import FigureRun, FigureValueRun, PageScan
61from quber.core.figures.nominate import NominatedPage, nominate_pages, table_verdicts
62from quber.core.figures.scan import SCAN_DEFAULT_MODEL, page_scan, scan_page
63from quber.core.figures.values import LocalValueReader, ScanValueParser, read_figure_values
64from quber.files.output import output_sink, write_completion_marker
65from quber.files.pdf import render_page
66from quber.providers.landing.client import is_v2_model
68#: Pages scanned at once. The scans are independent, and each spends nearly all
69#: its time waiting on the platform.
70SCAN_CONCURRENCY = 6
72#: Table corrections in flight at once, matching the table pipeline's own bound.
73CORRECT_CONCURRENCY = 16
75#: The two page breadths a run can take. `nominated` scans the pages the parse
76#: points at; `all` scans every page of the document, because a page whose
77#: content the parse never classified carries no nomination signal at all.
78PAGE_SCOPES = ("nominated", "all")
81@dataclass
82class ChartOutput:
83 """Where one figure run's artifacts landed, and what the run produced."""
85 run: FigureRun
86 document: DoclingDocument
87 tables: List[ExtractedTable] = field(default_factory=list)
88 artifacts: List[str] = field(default_factory=list)
89 values: Optional[FigureValueRun] = None
90 digests: Dict[int, dpt3.PageDigest] = field(default_factory=dict)
93async def run_figures(
94 source: Path,
95 document: DoclingDocument,
96 base: str,
97 parse_name: str,
98 output_dir: str = "output",
99 model: str = SCAN_DEFAULT_MODEL,
100 concurrency: int = SCAN_CONCURRENCY,
101 pages: str = "nominated",
102 llm: Optional[LLMClient] = None,
103 inspector: Optional[StatusInspector] = None,
104 cell_reader: Optional[CellReader] = None,
105 corrector: Optional[FigureCorrector] = None,
106 cells_pages: Optional[Dict[int, dict]] = None,
107 value_parser: Optional[ScanValueParser] = None,
108 value_reader: Optional[LocalValueReader] = None,
109) -> ChartOutput:
110 """Read the charts in `source` into `document`; write the scans and the refined parse.
112 `parse_name` is the filename the parse was read from, and the refined parse
113 is written back under that same name. The parse is an input the rest of the
114 pipeline reads by name, and this step enriches it rather than producing a
115 variant, so the name never changes and no later step has to be told which
116 file to pick up. Pointing `output_dir` at the directory the parse came from
117 therefore replaces it; pointing it elsewhere leaves the original alone.
119 A run that fails raises before anything is written, so a parse is never
120 replaced by one derived from a partial run.
121 """
122 if pages not in PAGE_SCOPES:
123 raise ValueError(f"pages must be one of {PAGE_SCOPES}, not {pages!r}")
124 nominated = nominate_pages(document, table_verdicts(document, source))
125 logger.info(
126 "Figure run: {} nominated {} page(s) of {}: {}",
127 base,
128 len(nominated),
129 len(document.pages),
130 [n.page for n in nominated],
131 )
132 if pages == "all":
133 nominated = widen_to_all_pages(document, nominated)
134 track = is_v2_model(model)
135 if track:
136 # Any page a dpt-3 scan already paid for is read too, whatever
137 # nominated it — a stat-panel page holds no picture and is never
138 # nominated, but its stored response costs nothing to digest and its
139 # block partition is what groups the page's text.
140 stored = _stored_track_pages(base, output_dir, {n.page for n in nominated})
141 if stored:
142 logger.info(
143 "Figure run: {} reading {} page(s) with a stored dpt-3 response: {}",
144 base,
145 len(stored),
146 [n.page for n in stored],
147 )
148 nominated = sorted([*nominated, *stored], key=lambda n: n.page)
149 llm = llm if llm is not None else get_llm_client(None)
150 inspector = inspector if inspector is not None else get_status_inspector()
151 cell_reader = cell_reader if cell_reader is not None else get_cell_reader()
152 corrector = corrector if corrector is not None else get_figure_corrector()
154 scans, digests, errors = await _scan_pages(source, base, nominated, model, output_dir, concurrency)
155 scans.sort(key=lambda s: s.page)
157 run = FigureRun(document=base, scans=scans, errors=errors)
158 dims = _page_dims(document)
159 # Tables first: reading one replaces a body inside an element the parse
160 # already holds, while the figures reshape the pictures around it.
161 refined = document.model_copy(deep=True)
162 assembly = TableAssembly(source, llm, asyncio.Semaphore(CORRECT_CONCURRENCY), CAPTURE_DPI, inspector)
163 tables, capture_errors = await capture_tables(
164 refined, scans, source, dims, assembly, cell_reader, orphans=track
165 )
166 run.errors.extend(capture_errors)
167 run.errors.extend(graft_tables(refined, scans, tables, dims))
168 if track:
169 # A captured grid over a region the document holds nothing for — a
170 # map's legend — is inserted at its reading-order position.
171 run.errors.extend(dpt3.insert_orphan_tables(refined, scans, tables, dims))
173 if track:
174 refined, graft_errors = dpt3.graft_figures(refined, scans, dims, source=source)
175 else:
176 refined, graft_errors = graft_figures(refined, scans, dims, source=source)
177 run.errors.extend(graft_errors)
178 # After the graft, so it sees only what the picture's own children did not
179 # already account for.
180 run.removed = await correct_figures(refined, scans, dims, source, corrector)
181 if track:
182 # After the correction, which records each figure's markers and notes on
183 # its picture — the resolution's inputs.
184 run.errors.extend(dpt3.resolve_figure_footnotes(refined, scans, digests))
185 # After the furniture sweep, so a removed fragment is never grouped.
186 dpt3.group_scanned_text(refined, scans, digests, dims)
187 run.errors.extend(dpt3.unhomed_tables(scans, refined, dims))
188 # Last, so it judges the document as it will be written: a picture the scan
189 # read carries its description by now, and a picture that turned out to be a
190 # table has its table beside it.
191 run.errors.extend(unread_pictures(refined, dims))
193 # After the graft, so each figure record carries its picture_ref and the
194 # values it produces can name the picture they belong to.
195 values_run: Optional[FigureValueRun] = None
196 if cells_pages and value_parser is not None and value_reader is not None:
197 figure_pages = sorted({s.page for s in scans if s.figures})
198 with tempfile.TemporaryDirectory() as tmp:
199 page_images: Dict[int, Path] = {}
200 for page_no in figure_pages:
201 image_path = Path(tmp) / f"values-p{page_no:04d}.png"
202 render_page(source, page_no, CAPTURE_DPI, image_path)
203 page_images[page_no] = image_path
204 values_run = await read_figure_values(
205 scans, cells_pages, page_images, value_parser, value_reader, base
206 )
207 run.errors.extend(values_run.errors)
209 artifacts = [s.response_artifact for s in scans if s.response_artifact]
210 with output_sink(output_dir) as out:
211 charts_path = out / f"{base}.figures.json"
212 charts_path.write_text(run.model_dump_json(indent=2), encoding="utf-8")
213 parse_path = out / parse_name
214 parse_path.write_text(json.dumps(refined.export_to_dict(), indent=2), encoding="utf-8")
215 artifacts.extend([charts_path.name, parse_path.name])
216 if tables:
217 tables_path = out / f"{base}.scanned-tables.json"
218 tables_path.write_text(
219 json.dumps([t.model_dump(mode="json") for t in tables], indent=2), encoding="utf-8"
220 )
221 artifacts.append(tables_path.name)
222 if values_run is not None:
223 values_path = out / f"{base}.figure-values.json"
224 values_path.write_text(values_run.model_dump_json(indent=2), encoding="utf-8")
225 artifacts.append(values_path.name)
226 if digests:
227 # The track's own records: what the projection left out of the graft
228 # contract, with each page's resolved and placed footnotes.
229 digest_path = out / f"{base}.figure-digest.json"
230 digest_path.write_text(
231 json.dumps([digests[page].model_dump(mode="json") for page in sorted(digests)], indent=2),
232 encoding="utf-8",
233 )
234 artifacts.append(digest_path.name)
236 logger.info(
237 "Figure run complete: {} nominated={} scanned={} submitted={} with_figures={} "
238 "records={} tables={} credits={}",
239 base,
240 run.nominated,
241 run.scanned,
242 run.submitted,
243 run.with_figures,
244 len(run.figures),
245 len(tables),
246 run.credits,
247 )
248 for err in run.errors:
249 logger.error(" - {}", err)
250 return ChartOutput(
251 run=run,
252 document=refined,
253 tables=tables,
254 artifacts=artifacts,
255 values=values_run,
256 digests=digests,
257 )
260def run_figures_sync(
261 source: Path,
262 document: DoclingDocument,
263 base: str,
264 parse_name: str,
265 output_dir: str = "output",
266 model: str = SCAN_DEFAULT_MODEL,
267 concurrency: int = SCAN_CONCURRENCY,
268 pages: str = "nominated",
269 llm: Optional[LLMClient] = None,
270 inspector: Optional[StatusInspector] = None,
271 cell_reader: Optional[CellReader] = None,
272 corrector: Optional[FigureCorrector] = None,
273 cells_pages: Optional[Dict[int, dict]] = None,
274 value_parser: Optional[ScanValueParser] = None,
275 value_reader: Optional[LocalValueReader] = None,
276) -> ChartOutput:
277 """`run_figures` for a synchronous caller."""
278 return asyncio.run(
279 run_figures(
280 source,
281 document,
282 base,
283 parse_name,
284 output_dir,
285 model,
286 concurrency,
287 pages,
288 llm,
289 inspector,
290 cell_reader,
291 corrector,
292 cells_pages,
293 value_parser,
294 value_reader,
295 )
296 )
299async def _scan_pages(
300 source: Path,
301 base: str,
302 kept: List[NominatedPage],
303 model: str,
304 output_dir: str,
305 concurrency: int,
306) -> Tuple[List[PageScan], Dict[int, dpt3.PageDigest], List[str]]:
307 """Scan every kept page, writing each raw response before reading it.
309 This is where the model selects the path: a dpt-3 model's responses are
310 digested by the track and the digests are returned beside the scans; any
311 other model's are read by `page_scan` as today.
313 A page whose scan or digestion raises is surfaced as an error and no record
314 is written for it, so the failure is visible rather than looking like a
315 page with no chart.
316 """
317 if not kept:
318 return [], {}, []
320 semaphore = asyncio.Semaphore(concurrency)
321 track = is_v2_model(model)
323 async def one(
324 nomination: NominatedPage,
325 ) -> Tuple[Optional[PageScan], Optional[dpt3.PageDigest], Optional[str]]:
326 async with semaphore:
327 try:
328 response, artifact, reused = await _page_response(
329 source, base, nomination.page, model, output_dir
330 )
331 except Exception as exc:
332 return (
333 None,
334 None,
335 (
336 f"page {nomination.page}: figure scan failed ({type(exc).__name__}: {exc}); "
337 "the page holds a chart the parse did not read"
338 ),
339 )
340 if track:
341 try:
342 digest = dpt3.digest_page(response, nomination.page, model, artifact)
343 except Exception as exc:
344 return (
345 None,
346 None,
347 (
348 f"page {nomination.page}: the scan was billed and its response could not be "
349 f"read ({type(exc).__name__}: {exc}); the raw response is stored as {artifact}"
350 ),
351 )
352 scan = dpt3.project_scan(digest, nomination.picture_classes, artifact, reused)
353 scan.table_refs = list(nomination.table_refs)
354 return scan, digest, None
355 scan = page_scan(response, nomination.page, model, nomination.picture_classes, artifact, reused)
356 scan.table_refs = list(nomination.table_refs)
357 return scan, None, None
359 results = await asyncio.gather(*(one(n) for n in kept))
360 scans = [s for s, _d, _e in results if s is not None]
361 digests = {d.page: d for _s, d, _e in results if d is not None}
362 errors = [e for _s, _d, e in results if e is not None]
363 return scans, digests, errors
366async def _page_response(
367 source: Path,
368 base: str,
369 page: int,
370 model: str,
371 output_dir: str,
372) -> Tuple[Dict[str, Any], str, bool]:
373 """The raw response for one page, and whether it came from a stored scan.
375 The stored response is named the way a single-page ADE run names it, so a
376 scan of that page from either path is reused instead of paid for again.
378 Reuse is keyed on the model: a stored response is reused only when its
379 format matches the requested model's generation, which the file itself
380 shows — a chunks list is dpt-2, a structure tree is dpt-3. A mismatched
381 file is moved aside to `{base}.p{N}.ade.{generation}.json` and the page is
382 rescanned, so a paid response is never destroyed and a later run under the
383 other generation finds its response by the same check. A stored file in
384 neither format raises, so a billed page can never pass as blank.
385 """
386 artifact = f"{base}.p{page}.ade.json"
387 requested = "dpt3" if is_v2_model(model) else "dpt2"
388 with output_sink(output_dir) as out:
389 stored = out / artifact
390 if stored.exists():
391 response = json.loads(stored.read_text(encoding="utf-8"))
392 generation = dpt3.response_generation(response)
393 if generation == requested:
394 logger.info("Figure scan: reusing the stored response for page {} ({})", page, artifact)
395 return response, artifact, True
396 if generation is None:
397 raise ValueError(
398 f"the stored response {artifact} (model {model}) holds neither a dpt-2 chunks "
399 "list nor a dpt-3 structure tree; move it aside to rescan the page"
400 )
401 aside = out / f"{base}.p{page}.ade.{generation}.json"
402 wanted = out / f"{base}.p{page}.ade.{requested}.json"
403 if wanted.exists():
404 candidate = json.loads(wanted.read_text(encoding="utf-8"))
405 if dpt3.response_generation(candidate) == requested:
406 stored.replace(aside)
407 wanted.replace(stored)
408 logger.info(
409 "Figure scan: page {} stored response is {}; moved aside to {} and "
410 "reusing the {} response already paid for",
411 page,
412 generation,
413 aside.name,
414 requested,
415 )
416 return candidate, artifact, True
417 stored.replace(aside)
418 logger.info(
419 "Figure scan: page {} stored response is {}; moved aside to {} and rescanning " "with {}",
420 page,
421 generation,
422 aside.name,
423 model,
424 )
426 response = await scan_page(source, page, model=model)
427 with output_sink(output_dir) as out:
428 (out / artifact).write_text(json.dumps(response, indent=1), encoding="utf-8")
429 return response, artifact, False
432def widen_to_all_pages(document: DoclingDocument, nominated: List[NominatedPage]) -> List[NominatedPage]:
433 """Every page of the document, keeping what nominated the nominated ones.
435 A page the parse pointed at keeps its picture and table refs; every other
436 page joins bare — no refs, just the page — and is scanned whole through
437 the same flow.
438 """
439 covered = {n.page for n in nominated}
440 extra = [NominatedPage(page=p) for p in sorted(document.pages) if p not in covered]
441 if extra:
442 logger.info(
443 "Figure run: scanning {} un-nominated page(s) under pages=all: {}",
444 len(extra),
445 [n.page for n in extra],
446 )
447 return sorted([*nominated, *extra], key=lambda n: n.page)
450def _stored_track_pages(base: str, output_dir: str, nominated: set[int]) -> List[NominatedPage]:
451 """Pages beyond nomination whose stored dpt-3 response already sits on disk.
453 Reading one costs nothing — the response was paid for by whichever run made
454 it, `quber ade --page` included, and it is named the way this workflow
455 names its own. Only a response already in the dpt-3 format qualifies: a
456 dpt-2 response here would be rescanned at a price, and a page nobody
457 nominated must never be billed by the back door.
458 """
459 pattern = re.compile(rf"{re.escape(base)}\.p(\d+)\.ade\.json")
460 found: List[NominatedPage] = []
461 with output_sink(output_dir) as out:
462 for path in sorted(out.glob(f"{base}.p*.ade.json")):
463 match = pattern.fullmatch(path.name)
464 if match is None:
465 continue
466 page = int(match.group(1))
467 if page in nominated:
468 continue
469 try:
470 response = json.loads(path.read_text(encoding="utf-8"))
471 except ValueError:
472 continue
473 if dpt3.response_generation(response) == "dpt3":
474 found.append(NominatedPage(page=page))
475 return found
478def _page_dims(document: DoclingDocument) -> Dict[int, Tuple[float, float]]:
479 """Each page's width and height in points, as the parse recorded them."""
480 return {
481 page_no: (page.size.width, page.size.height)
482 for page_no, page in document.pages.items()
483 if page.size is not None
484 }
487#: The names the pipeline writes a parse under, newest stage first. A directory
488#: given as the parse is read for the latest of these it holds, so a document
489#: that has been through fusion is refined rather than the raw parse beside it.
490PARSE_NAMES = ("unified.json", "docling.json")
493def resolve_parse(parse: str, base: str) -> Path:
494 """The parse file to refine, from a JSON file or a directory holding one.
496 A file is taken as given. A directory is searched for `<base>.unified.json`
497 and then `<base>.docling.json`.
498 """
499 path = Path(parse)
500 if not path.is_dir():
501 if not path.exists():
502 raise FileNotFoundError(f"no parse to refine at {path}")
503 return path
504 for name in PARSE_NAMES:
505 candidate = path / f"{base}.{name}"
506 if candidate.exists():
507 return candidate
508 wanted = " or ".join(f"{base}.{n}" for n in PARSE_NAMES)
509 raise FileNotFoundError(f"no parse to refine in {path}: expected {wanted}")
512def load_document(parse: str, base: str) -> Tuple[DoclingDocument, str]:
513 """The parse to refine and the filename it was read from."""
514 path = resolve_parse(parse, base)
515 return DoclingDocument.load_from_json(path), path.name
518def write_marker(output_dir: str, base: str, payload: Dict[str, Any]) -> None:
519 """Write the figure run's completion marker, strictly after its artifacts."""
520 write_completion_marker(output_dir, f"{base}.figure.complete.json", payload)