Coverage for src / quber / core / extractors / camelot / correspondence / orchestrator.py: 33%
217 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""
2CamelotCorrespondenceExtractor: orchestrates the per-page correspondence
3flow — concurrent Camelot + detection, chunk assignment, guarded
4recovery, the completeness/extend/fill loop, and output assembly. The
5deterministic logic each step rests on lives in the sibling modules
6(`geometry`, `matching`, `recovery`, `correction`).
8DEPRECATED and dormant: superseded by `SetOfMarkExtractor` (vision-guided,
9in-region Camelot, grounded correction), which drops this flow's
10detector/match/recovery escalation. Kept importable for comparison only; do
11not build new work on it. Note its `correction`/`geometry` helpers are still
12reused live by the Set-of-Mark pipeline.
13"""
15from __future__ import annotations
17import asyncio
18import tempfile
19from pathlib import Path
20from typing import Dict, List, Literal, Optional, Tuple
22from loguru import logger
24from quber.agents.completeness import (
25 CompletenessAuditor,
26 extracted_span_pts,
27 get_completeness,
28 page_words,
29 repair_box,
30 round_off_grid,
31 tabular_bands,
32)
33from quber.agents.detector import DetectedTable, TableDetector, get_detector
34from quber.agents.grid_locator import GridLocator, LocatedTable, get_grid_locator
35from quber.agents.llm_client import FootnoteDef, LLMClient
36from quber.core.extractors.base import ExtractedTable, ExtractionRecord, FilledCell
37from quber.core.extractors.camelot.acquire import (
38 CAMELOT_FLAVOR_TIMEOUT_S,
39 CamelotCandidate,
40 Flavor,
41 grid_to_markdown,
42 is_content_empty,
43 render_pages,
44 run_camelot_flavors_parallel,
45)
46from quber.core.extractors.camelot.correspondence.correction import (
47 StructureCorrection,
48 correct_structure,
49)
50from quber.core.extractors.camelot.correspondence.geometry import (
51 neighbor_bounded_bbox,
52 norm_bbox_to_table_area,
53 page_size_pts,
54)
55from quber.core.extractors.camelot.correspondence.matching import (
56 assemble_cells,
57 assign_chunks,
58 chunk_top,
59)
60from quber.core.extractors.camelot.correspondence.recovery import (
61 camelot_targeted,
62 nearest_grid_region,
63)
66class CamelotCorrespondenceExtractor:
67 detector: TableDetector
68 completeness: CompletenessAuditor
69 dpi: int
70 run_completeness: bool
71 run_recovery: bool
72 run_fill: bool
73 run_box_repair: bool
74 run_grid_recovery: bool
75 grid_locator: Optional[GridLocator]
76 max_concurrent: int
77 flavor_timeout_s: float
78 llm: Optional[LLMClient]
79 run_llm_correction: bool
80 correct_sem: Optional[asyncio.Semaphore]
82 def __init__(
83 self,
84 detector: Optional[TableDetector] = None,
85 completeness: Optional[CompletenessAuditor] = None,
86 dpi: int = 200,
87 run_completeness: bool = True,
88 run_recovery: bool = True,
89 run_fill: bool = True,
90 run_box_repair: bool = True,
91 run_grid_recovery: bool = True,
92 grid_locator: Optional[GridLocator] = None,
93 max_concurrent: int = 5,
94 flavor_timeout_s: float = CAMELOT_FLAVOR_TIMEOUT_S,
95 llm: Optional[LLMClient] = None,
96 run_llm_correction: bool = True,
97 ) -> None:
98 self.detector = detector or get_detector()
99 self.completeness = completeness or get_completeness()
100 self.dpi = dpi
101 self.run_completeness = run_completeness
102 self.run_recovery = run_recovery
103 self.run_fill = run_fill
104 self.run_box_repair = run_box_repair
105 self.run_grid_recovery = run_grid_recovery
106 # Constructed lazily on first escalation so the extractor can run with
107 # grid recovery disabled (or no API creds) without building the agent.
108 self.grid_locator = grid_locator
109 self.max_concurrent = max_concurrent
110 self.flavor_timeout_s = flavor_timeout_s
111 # Output-stage structure correction. The grid stays the value
112 # source of record; the LLM only restructures presentation.
113 self.llm = llm
114 self.run_llm_correction = run_llm_correction
115 # Bounded-concurrency gate for the correction calls; created lazily
116 # inside the running loop so the extractor stays loop-agnostic.
117 self.correct_sem = None
119 async def extract_tables(self, source: Path) -> List[ExtractedTable]:
120 source = Path(source)
121 if not source.exists():
122 raise FileNotFoundError(source)
124 with tempfile.TemporaryDirectory(prefix="quber-camelot-corr-") as tmpdir:
125 tmp = Path(tmpdir)
126 page_images = render_pages(source, dpi=self.dpi, out_dir=tmp)
128 # Step 1: Camelot (both flavors) on the CPU, concurrent with
129 # per-page detection (network).
130 candidates = await asyncio.to_thread(run_camelot_flavors_parallel, source, self.flavor_timeout_s)
131 populated = [c for c in candidates if not is_content_empty(c.markdown)]
132 logger.info(
133 "camelot: {} chunks ({} lattice, {} stream), {} populated after empty-shell filter",
134 len(candidates),
135 sum(1 for c in candidates if c.flavor == "lattice"),
136 sum(1 for c in candidates if c.flavor == "stream"),
137 len(populated),
138 )
140 by_page_chunks: Dict[int, List[CamelotCandidate]] = {}
141 for c in populated:
142 by_page_chunks.setdefault(c.page, []).append(c)
144 sem = asyncio.Semaphore(self.max_concurrent)
146 async def detect_page(page: int) -> Tuple[int, List[DetectedTable]]:
147 async with sem:
148 res = await self.detector.detect(page_images[page - 1])
149 return page, res.tables
151 detections = await asyncio.gather(*(detect_page(p) for p in range(1, len(page_images) + 1)))
152 by_page_detected: Dict[int, List[DetectedTable]] = dict(detections)
154 page_results = await asyncio.gather(
155 *(
156 self.process_page(
157 page=p,
158 page_image=page_images[p - 1],
159 detected=sorted(by_page_detected.get(p, []), key=lambda d: d.ordinal),
160 chunks=by_page_chunks.get(p, []),
161 source=str(source),
162 )
163 for p in range(1, len(page_images) + 1)
164 )
165 )
166 return [t for page in page_results for t in page]
168 async def process_page(
169 self,
170 page: int,
171 page_image: Path,
172 detected: List[DetectedTable],
173 chunks: List[CamelotCandidate],
174 source: str,
175 ) -> List[ExtractedTable]:
176 if not detected:
177 # No table on the page per the arbiter. Any Camelot chunks
178 # here are stream false positives (footnotes, prose); they
179 # match nothing, so they fall away.
180 if chunks:
181 logger.info(
182 "page {}: detector found 0 tables; dropping {} unmatched chunks", page, len(chunks)
183 )
184 return []
186 lattice = sorted([c for c in chunks if c.flavor == "lattice"], key=chunk_top)
187 stream = sorted([c for c in chunks if c.flavor == "stream"], key=chunk_top)
189 page_w_pts, page_h_pts = page_size_pts(page_image, self.dpi)
191 # Validate/repair detector boxes against the text layer (deterministic,
192 # no second model call). The detector locates tables but its boxes can
193 # land in whitespace or clip a table; downstream matching AND the
194 # focused-stream recovery both aim at these boxes, so a bad box becomes
195 # a miss. Snap each box onto the tabular text it actually covers. When
196 # no band can be found we leave the box UNCHANGED rather than drop the
197 # detection — a real table must never be silently dropped, so a box we
198 # cannot place flows through to detected_not_extracted (reported).
199 if self.run_box_repair and any(d.bbox is not None for d in detected):
200 pw_words, ph_words, words = await asyncio.to_thread(page_words, Path(source), page)
201 bands = tabular_bands(words)
202 for d in detected:
203 if d.bbox is None:
204 continue
205 repaired = repair_box(d.bbox, bands, pw_words, ph_words)
206 if repaired is not None:
207 d.bbox = repaired
209 for d in detected:
210 if d.bbox is None:
211 logger.warning(
212 "page {}: detected table {} has no bbox; geometric matching cannot place it",
213 page,
214 d.ordinal,
215 )
217 covered: set[int] = set() # detected indices already resolved
218 results: List[ExtractedTable] = []
220 # Lattice first (trusted), then stream for whatever lattice missed.
221 sources: List[Tuple[Flavor, List[CamelotCandidate]]] = [
222 ("lattice", lattice),
223 ("stream", stream),
224 ]
225 for flavor, flavor_chunks in sources:
226 remaining = [di for di in range(len(detected)) if di not in covered]
227 if not remaining or not flavor_chunks:
228 continue
230 assignment = assign_chunks(detected, remaining, flavor_chunks, page_w_pts, page_h_pts)
232 # Combined: one chunk owns two-or-more detected tables. Kept
233 # whole, flagged, never split.
234 for ci, dis in assignment.tables_for_chunk.items():
235 if len(dis) >= 2:
236 cand = flavor_chunks[ci]
237 ordinals = [detected[di].ordinal for di in dis]
238 combined_cells = [list(r) for r in cand.cells]
239 correction = await self.apply_structure_correction(
240 combined_cells, page_image, cand.bbox, source, page
241 )
242 results.append(
243 self.make_extracted(
244 cand=cand,
245 cells=combined_cells,
246 detected=detected[dis[0]],
247 flavor=flavor,
248 source=source,
249 source_ids=[cand.candidate_id],
250 status="extracted",
251 combined=True,
252 combined_ordinals=ordinals,
253 correction=correction,
254 )
255 )
256 covered.update(dis)
257 logger.info("page {}: chunk {} combined -> tables {}", page, cand.candidate_id, ordinals)
259 # Orphan chunks (no detected table's best match) are the only
260 # safe source for completeness extension. Shared and consumed
261 # once across this flavor's tables, in page order.
262 orphan_queue = list(assignment.orphans)
264 # Simple: each remaining table takes its own best chunk; a
265 # truncated table is completed from the orphan queue downstream.
266 for di in remaining:
267 if di in covered or di not in assignment.best_chunk:
268 continue
269 ci = assignment.best_chunk[di]
270 if len(assignment.tables_for_chunk[ci]) != 1:
271 continue
272 matched_chunks = [flavor_chunks[ci]]
273 table = await self.finalize(
274 detected=detected[di],
275 region_bbox=neighbor_bounded_bbox(detected, di),
276 primary=matched_chunks[0],
277 matched_chunks=matched_chunks,
278 orphan_queue=orphan_queue,
279 flavor=flavor,
280 source=source,
281 page=page,
282 page_image=page_image,
283 )
284 results.append(table)
285 covered.add(di)
287 # Report, don't drop (with a guarded recovery pass). A detected
288 # table the full-page passes missed gets one targeted, region-
289 # constrained stream attempt at the detector's box. If it yields
290 # data, emit it; otherwise (Camelot crash on an empty region, or
291 # empty result) report it as detected_not_extracted — never worse
292 # than today.
293 # Grid-locator escalation state, computed lazily on first need and
294 # shared across this page's unrecovered tables (one vision call max).
295 grid_located: Optional[List[LocatedTable]] = None
296 consumed_grid: set[int] = set()
298 for di in range(len(detected)):
299 if di in covered:
300 continue
301 d = detected[di]
302 region = neighbor_bounded_bbox(detected, di)
304 recovered: Optional[CamelotCandidate] = None
305 if self.run_recovery and region is not None:
306 area = norm_bbox_to_table_area(region, page_w_pts, page_h_pts)
307 try:
308 cand = await asyncio.to_thread(camelot_targeted, source, page, area, d.ordinal)
309 except Exception as exc:
310 cand = None
311 logger.warning("page {}: targeted recovery raised for table {}: {}", page, d.ordinal, exc)
312 if cand is not None and not is_content_empty(cand.markdown):
313 recovered = cand
315 # Escalation: the detector's box drifted off this table, so the
316 # box-aimed recovery hit empty space. Relocate it with the grid
317 # locator (Set-of-Mark discrete IDs, immune to bbox drift) and retry
318 # the focused stream at the tightened region. Lazy and surgical: at
319 # most one vision call per page, only when the cheap path failed.
320 if recovered is None and self.run_grid_recovery:
321 if grid_located is None:
322 locator = self.grid_locator or get_grid_locator()
323 self.grid_locator = locator
324 grid_located = await locator.locate(page_image, Path(source), page)
325 gi = nearest_grid_region(grid_located, consumed_grid, d.bbox)
326 if gi is not None:
327 consumed_grid.add(gi)
328 grid_region = grid_located[gi].region
329 area = norm_bbox_to_table_area(grid_region, page_w_pts, page_h_pts)
330 try:
331 cand = await asyncio.to_thread(camelot_targeted, source, page, area, d.ordinal)
332 except Exception as exc:
333 cand = None
334 logger.warning(
335 "page {}: grid-relocation recovery raised for table {}: {}", page, d.ordinal, exc
336 )
337 if cand is not None and not is_content_empty(cand.markdown):
338 recovered = cand
339 region = grid_region
340 logger.info(
341 "page {}: table {} RELOCATED via grid locator (grid region {})",
342 page,
343 d.ordinal,
344 gi + 1,
345 )
347 if recovered is not None:
348 logger.info(
349 "page {}: table {} RECOVERED via targeted extraction ({} rows)",
350 page,
351 d.ordinal,
352 len(recovered.markdown.splitlines()),
353 )
354 results.append(
355 await self.finalize(
356 detected=d,
357 region_bbox=region,
358 primary=recovered,
359 matched_chunks=[recovered],
360 orphan_queue=[],
361 flavor="stream",
362 source=source,
363 page=page,
364 page_image=page_image,
365 )
366 )
367 covered.add(di)
368 continue
370 logger.warning(
371 "page {}: table {} '{}' DETECTED-NOT-EXTRACTED (detector found it; "
372 "lattice, stream, and recovery all produced nothing)",
373 page,
374 d.ordinal,
375 d.description[:60],
376 )
377 results.append(
378 ExtractedTable(
379 title=d.description,
380 markdown="",
381 page=page,
382 source=source,
383 extraction_record=ExtractionRecord(
384 status="detected_not_extracted",
385 detected_ordinal=d.ordinal,
386 detected_description=d.description,
387 ),
388 )
389 )
391 return results
393 async def apply_structure_correction(
394 self,
395 cells: List[List[str]],
396 page_image: Optional[Path],
397 bbox: Optional[Tuple[float, float, float, float]],
398 source: str,
399 page: int,
400 ) -> Optional[StructureCorrection]:
401 """Config gate over `correction.correct_structure`: skipped when
402 correction is disabled or no LLM is wired. The shared semaphore is
403 created lazily inside the running loop so the extractor stays
404 loop-agnostic.
405 """
406 if self.llm is None or not self.run_llm_correction:
407 return None
408 if self.correct_sem is None:
409 self.correct_sem = asyncio.Semaphore(self.max_concurrent)
410 return await correct_structure(
411 cells=cells,
412 page_image=page_image,
413 bbox=bbox,
414 source=source,
415 page=page,
416 llm=self.llm,
417 correct_sem=self.correct_sem,
418 dpi=self.dpi,
419 )
421 async def finalize(
422 self,
423 detected: DetectedTable,
424 region_bbox: Optional[Tuple[float, float, float, float]],
425 primary: CamelotCandidate,
426 matched_chunks: List[CamelotCandidate],
427 orphan_queue: List[CamelotCandidate],
428 flavor: Flavor,
429 source: str,
430 page: int,
431 page_image: Optional[Path],
432 ) -> ExtractedTable:
433 status: Literal["extracted", "incomplete"] = "extracted"
434 complete: Optional[bool] = None
435 gap: str = ""
436 used = list(matched_chunks)
437 filled: List[FilledCell] = []
438 # The structured grid is canonical here: assembly, the completeness
439 # audit and the fill all operate on it; markdown is rendered once
440 # by make_extracted.
441 cells = assemble_cells(used)
443 if self.run_completeness:
444 # The audit reads the PDF text layer (the same source Camelot
445 # reads) within `region_bbox`, which is bounded by the neighbor
446 # tables so an adjacent table's rows are not read as a
447 # continuation. `extracted_bboxes` gives the span we actually
448 # captured, so the audit can place any missing figure at the
449 # top or bottom edge.
450 verdict = None
451 while True:
452 verdict = await self.completeness.audit(
453 source=Path(source),
454 page=page,
455 region_bbox=region_bbox,
456 extracted_bboxes=[c.bbox for c in used],
457 assembled_markdown=grid_to_markdown(cells),
458 )
459 complete, gap = verdict.complete, verdict.gap
460 if complete or not orphan_queue:
461 break
462 # Truncated at an edge: extend with the next orphan chunk
463 # (shared, consumed once across the page so no table is
464 # polluted), then re-audit.
465 nxt = orphan_queue.pop(0)
466 used.append(nxt)
467 cells = assemble_cells(used)
468 logger.info(
469 "page {}: table {} truncated; extending with {} and re-auditing",
470 page,
471 detected.ordinal,
472 nxt.candidate_id,
473 )
475 # Text-layer fill (round-off), the last-resort straggler cleanup:
476 # no Camelot chunk could supply the truncated edge row, but the
477 # figures are in the text layer (that is how the audit found the
478 # gap). Insert them into the grid in place, re-audit, and record
479 # the amendment as provenance — never blurred with Camelot's
480 # cells.
481 if not complete and self.run_fill and verdict is not None and verdict.missing_figures:
482 _, page_h_pts, words = await asyncio.to_thread(page_words, Path(source), page)
483 span = extracted_span_pts([c.bbox for c in used], page_h_pts)
484 if span is not None:
485 new_cells, placed = round_off_grid(cells, words, span, verdict.missing_figures)
486 if placed:
487 cells = new_cells
488 filled = [
489 FilledCell(
490 value=p.value,
491 column=p.column,
492 row_label=p.row_label,
493 edge=p.edge,
494 x=p.x,
495 y=p.y,
496 )
497 for p in placed
498 ]
499 reverdict = await self.completeness.audit(
500 source=Path(source),
501 page=page,
502 region_bbox=region_bbox,
503 extracted_bboxes=[c.bbox for c in used],
504 assembled_markdown=grid_to_markdown(cells),
505 )
506 complete, gap = reverdict.complete, reverdict.gap
507 logger.info(
508 "page {}: table {} ROUNDED OFF; filled {} text-layer figure(s); now {}",
509 page,
510 detected.ordinal,
511 len(filled),
512 "complete" if complete else "still incomplete",
513 )
515 if not complete:
516 status = "incomplete"
517 logger.warning(
518 "page {}: table {} INCOMPLETE; {} (gap reported)",
519 page,
520 detected.ordinal,
521 gap or "edge truncated",
522 )
524 correction = await self.apply_structure_correction(cells, page_image, primary.bbox, source, page)
525 return self.make_extracted(
526 cand=primary,
527 cells=cells,
528 detected=detected,
529 flavor=flavor,
530 source=source,
531 source_ids=[c.candidate_id for c in used],
532 status=status,
533 completeness_complete=complete,
534 completeness_gap=gap,
535 filled_cells=filled,
536 correction=correction,
537 )
539 def make_extracted(
540 self,
541 cand: CamelotCandidate,
542 cells: List[List[str]],
543 detected: DetectedTable,
544 flavor: Flavor,
545 source: str,
546 source_ids: List[str],
547 status: Literal["extracted", "incomplete"],
548 completeness_complete: Optional[bool] = None,
549 completeness_gap: str = "",
550 combined: bool = False,
551 combined_ordinals: Optional[List[int]] = None,
552 filled_cells: Optional[List[FilledCell]] = None,
553 correction: Optional[StructureCorrection] = None,
554 ) -> ExtractedTable:
555 filled_cells = filled_cells or []
556 # Default presentation is the deterministic grid render with the
557 # detector's description as title. When structure correction was
558 # accepted, the LLM's restructured markdown and heading metadata take
559 # over; the LLM title wins, falling back to the detector description
560 # when the model returns none.
561 title = detected.description
562 subtitle = ""
563 caption = ""
564 footnotes: List[FootnoteDef] = []
565 markdown = grid_to_markdown(cells)
566 llm_corrected = False
567 if correction is not None:
568 title = correction.title or detected.description
569 caption = correction.caption
570 footnotes = list(correction.footnotes)
571 markdown = correction.markdown
572 llm_corrected = correction.llm_corrected
573 return ExtractedTable(
574 title=title,
575 subtitle=subtitle,
576 caption=caption,
577 markdown=markdown,
578 footnotes=footnotes,
579 page=cand.page,
580 bbox=cand.bbox,
581 flavor=flavor,
582 source=source,
583 camelot_accuracy=cand.accuracy,
584 llm_corrected=llm_corrected,
585 extraction_record=ExtractionRecord(
586 status=status,
587 detected_ordinal=detected.ordinal,
588 detected_description=detected.description,
589 matched_flavor=flavor,
590 source_candidate_ids=source_ids,
591 combined=combined,
592 combined_ordinals=combined_ordinals or [],
593 completeness_complete=completeness_complete,
594 completeness_gap=completeness_gap,
595 filled_from_text_layer=bool(filled_cells),
596 filled_cells=filled_cells,
597 ),
598 )
600 def extract_tables_sync(self, source: Path) -> List[ExtractedTable]:
601 """Sync entry point for callers without an event loop (CLI,
602 scripts). Wraps the async path with `asyncio.run`.
603 """
604 return asyncio.run(self.extract_tables(source))