""" CamelotCorrespondenceExtractor: orchestrates the per-page correspondence flow — Camelot, then detection, chunk assignment, guarded recovery, the completeness/extend/fill loop, and output assembly. The deterministic logic each step rests on lives in the sibling modules (`geometry`, `matching`, `recovery`, `correction`). DEPRECATED and dormant: superseded by `SetOfMarkExtractor` (vision-guided, in-region Camelot, grounded correction), which drops this flow's detector/match/recovery escalation. Do not build new work on it. It stays reachable through `quber table --engine correspondence`. The sibling `correction`, `geometry` and `recovery` modules are live dependencies of the Set-of-Mark pipeline. Only this module and `matching` are unused by it. """ from __future__ import annotations import asyncio import tempfile from pathlib import Path from typing import Dict, List, Literal, Optional, Tuple from loguru import logger from quber.agents.completeness import ( CompletenessAuditor, extracted_span_pts, get_completeness, page_words, repair_box, round_off_grid, tabular_bands, ) from quber.agents.detector import DetectedTable, TableDetector, get_detector from quber.agents.grid_locator import GridLocator, LocatedTable, get_grid_locator from quber.agents.llm_client import FootnoteDef, LLMClient from quber.core.extractors.base import ExtractedTable, ExtractionRecord, FilledCell from quber.core.extractors.camelot.acquire import ( CAMELOT_FLAVOR_TIMEOUT_S, CamelotCandidate, Flavor, grid_to_markdown, is_content_empty, render_pages, run_camelot_flavors_parallel, ) from quber.core.extractors.camelot.correspondence.correction import ( StructureCorrection, correct_structure, ) from quber.core.extractors.camelot.correspondence.geometry import ( neighbor_bounded_bbox, norm_bbox_to_table_area, page_size_pts, ) from quber.core.extractors.camelot.correspondence.matching import ( assemble_cells, assign_chunks, chunk_top, ) from quber.core.extractors.camelot.correspondence.recovery import ( camelot_targeted, nearest_grid_region, ) class CamelotCorrespondenceExtractor: detector: TableDetector completeness: CompletenessAuditor dpi: int run_completeness: bool run_recovery: bool run_fill: bool run_box_repair: bool run_grid_recovery: bool grid_locator: Optional[GridLocator] max_concurrent: int flavor_timeout_s: float llm: Optional[LLMClient] run_llm_correction: bool correct_sem: Optional[asyncio.Semaphore] def __init__( self, detector: Optional[TableDetector] = None, completeness: Optional[CompletenessAuditor] = None, dpi: int = 200, run_completeness: bool = True, run_recovery: bool = True, run_fill: bool = True, run_box_repair: bool = True, run_grid_recovery: bool = True, grid_locator: Optional[GridLocator] = None, max_concurrent: int = 5, flavor_timeout_s: float = CAMELOT_FLAVOR_TIMEOUT_S, llm: Optional[LLMClient] = None, run_llm_correction: bool = True, ) -> None: self.detector = detector or get_detector() self.completeness = completeness or get_completeness() self.dpi = dpi self.run_completeness = run_completeness self.run_recovery = run_recovery self.run_fill = run_fill self.run_box_repair = run_box_repair self.run_grid_recovery = run_grid_recovery # Constructed lazily on first escalation so the extractor can run with # grid recovery disabled (or no API creds) without building the agent. self.grid_locator = grid_locator self.max_concurrent = max_concurrent self.flavor_timeout_s = flavor_timeout_s # Output-stage structure correction. The grid stays the value # source of record; the LLM only restructures presentation. self.llm = llm self.run_llm_correction = run_llm_correction # Bounded-concurrency gate for the correction calls; created lazily # inside the running loop so the extractor stays loop-agnostic. self.correct_sem = None async def extract_tables(self, source: Path) -> List[ExtractedTable]: source = Path(source) if not source.exists(): raise FileNotFoundError(source) with tempfile.TemporaryDirectory(prefix="quber-camelot-corr-") as tmpdir: tmp = Path(tmpdir) page_images = render_pages(source, dpi=self.dpi, out_dir=tmp) # Step 1: Camelot (both flavors) on the CPU runs to completion # first. Per-page detection (network) runs after it, concurrent # across pages only. candidates = await asyncio.to_thread(run_camelot_flavors_parallel, source, self.flavor_timeout_s) populated = [c for c in candidates if not is_content_empty(c.markdown)] logger.info( "camelot: {} chunks ({} lattice, {} stream), {} populated after empty-shell filter", len(candidates), sum(1 for c in candidates if c.flavor == "lattice"), sum(1 for c in candidates if c.flavor == "stream"), len(populated), ) by_page_chunks: Dict[int, List[CamelotCandidate]] = {} for c in populated: by_page_chunks.setdefault(c.page, []).append(c) sem = asyncio.Semaphore(self.max_concurrent) async def detect_page(page: int) -> Tuple[int, List[DetectedTable]]: async with sem: res = await self.detector.detect(page_images[page - 1]) return page, res.tables detections = await asyncio.gather(*(detect_page(p) for p in range(1, len(page_images) + 1))) by_page_detected: Dict[int, List[DetectedTable]] = dict(detections) page_results = await asyncio.gather( *( self.process_page( page=p, page_image=page_images[p - 1], detected=sorted(by_page_detected.get(p, []), key=lambda d: d.ordinal), chunks=by_page_chunks.get(p, []), source=str(source), ) for p in range(1, len(page_images) + 1) ) ) return [t for page in page_results for t in page] async def process_page( self, page: int, page_image: Path, detected: List[DetectedTable], chunks: List[CamelotCandidate], source: str, ) -> List[ExtractedTable]: if not detected: # No table on the page per the arbiter. Any Camelot chunks # here are stream false positives (footnotes, prose); they # match nothing, so they fall away. if chunks: logger.info( "page {}: detector found 0 tables; dropping {} unmatched chunks", page, len(chunks) ) return [] lattice = sorted([c for c in chunks if c.flavor == "lattice"], key=chunk_top) stream = sorted([c for c in chunks if c.flavor == "stream"], key=chunk_top) page_w_pts, page_h_pts = page_size_pts(page_image, self.dpi) # Validate/repair detector boxes against the text layer (deterministic, # no second model call). The detector locates tables but its boxes can # land in whitespace or clip a table; downstream matching AND the # focused-stream recovery both aim at these boxes, so a bad box becomes # a miss. Snap each box onto the tabular text it actually covers. When # no band can be found we leave the box UNCHANGED rather than drop the # detection — a real table must never be silently dropped, so a box we # cannot place flows through to detected_not_extracted (reported). if self.run_box_repair and any(d.bbox is not None for d in detected): pw_words, ph_words, words = await asyncio.to_thread(page_words, Path(source), page) bands = tabular_bands(words) for d in detected: if d.bbox is None: continue repaired = repair_box(d.bbox, bands, pw_words, ph_words) if repaired is not None: d.bbox = repaired for d in detected: if d.bbox is None: logger.warning( "page {}: detected table {} has no bbox; geometric matching cannot place it", page, d.ordinal, ) covered: set[int] = set() # detected indices already resolved results: List[ExtractedTable] = [] # Lattice first (trusted), then stream for whatever lattice missed. sources: List[Tuple[Flavor, List[CamelotCandidate]]] = [ ("lattice", lattice), ("stream", stream), ] for flavor, flavor_chunks in sources: remaining = [di for di in range(len(detected)) if di not in covered] if not remaining or not flavor_chunks: continue assignment = assign_chunks(detected, remaining, flavor_chunks, page_w_pts, page_h_pts) # Combined: one chunk owns two-or-more detected tables. Kept # whole, flagged, never split. for ci, dis in assignment.tables_for_chunk.items(): if len(dis) >= 2: cand = flavor_chunks[ci] ordinals = [detected[di].ordinal for di in dis] combined_cells = [list(r) for r in cand.cells] correction = await self.apply_structure_correction( combined_cells, page_image, cand.bbox, source, page ) results.append( self.make_extracted( cand=cand, cells=combined_cells, detected=detected[dis[0]], flavor=flavor, source=source, source_ids=[cand.candidate_id], status="extracted", combined=True, combined_ordinals=ordinals, correction=correction, ) ) covered.update(dis) logger.info("page {}: chunk {} combined -> tables {}", page, cand.candidate_id, ordinals) # Orphan chunks (no detected table's best match) are the only # safe source for completeness extension. Shared and consumed # once across this flavor's tables, in page order. orphan_queue = list(assignment.orphans) # Simple: each remaining table takes its own best chunk; a # truncated table is completed from the orphan queue downstream. for di in remaining: if di in covered or di not in assignment.best_chunk: continue ci = assignment.best_chunk[di] if len(assignment.tables_for_chunk[ci]) != 1: continue matched_chunks = [flavor_chunks[ci]] table = await self.finalize( detected=detected[di], region_bbox=neighbor_bounded_bbox(detected, di), primary=matched_chunks[0], matched_chunks=matched_chunks, orphan_queue=orphan_queue, flavor=flavor, source=source, page=page, page_image=page_image, ) results.append(table) covered.add(di) # Report, don't drop (with a guarded recovery pass). A detected # table the full-page passes missed gets one targeted, region- # constrained stream attempt at the detector's box. If it yields # data, emit it; otherwise (Camelot crash on an empty region, or # empty result) report it as detected_not_extracted — never worse # than today. # Grid-locator escalation state, computed lazily on first need and # shared across this page's unrecovered tables (one vision call max). grid_located: Optional[List[LocatedTable]] = None consumed_grid: set[int] = set() for di in range(len(detected)): if di in covered: continue d = detected[di] region = neighbor_bounded_bbox(detected, di) recovered: Optional[CamelotCandidate] = None if self.run_recovery and region is not None: area = norm_bbox_to_table_area(region, page_w_pts, page_h_pts) try: cand = await asyncio.to_thread(camelot_targeted, source, page, area, d.ordinal) except Exception as exc: cand = None logger.warning("page {}: targeted recovery raised for table {}: {}", page, d.ordinal, exc) if cand is not None and not is_content_empty(cand.markdown): recovered = cand # Escalation: the detector's box drifted off this table, so the # box-aimed recovery hit empty space. Relocate it with the grid # locator (Set-of-Mark discrete IDs, immune to bbox drift) and retry # the focused stream at the tightened region. Lazy and surgical: at # most one vision call per page, only when the cheap path failed. if recovered is None and self.run_grid_recovery: if grid_located is None: locator = self.grid_locator or get_grid_locator() self.grid_locator = locator grid_located = await locator.locate(page_image, Path(source), page) gi = nearest_grid_region(grid_located, consumed_grid, d.bbox) if gi is not None: consumed_grid.add(gi) grid_region = grid_located[gi].region area = norm_bbox_to_table_area(grid_region, page_w_pts, page_h_pts) try: cand = await asyncio.to_thread(camelot_targeted, source, page, area, d.ordinal) except Exception as exc: cand = None logger.warning( "page {}: grid-relocation recovery raised for table {}: {}", page, d.ordinal, exc ) if cand is not None and not is_content_empty(cand.markdown): recovered = cand region = grid_region logger.info( "page {}: table {} RELOCATED via grid locator (grid region {})", page, d.ordinal, gi + 1, ) if recovered is not None: logger.info( "page {}: table {} RECOVERED via targeted extraction ({} rows)", page, d.ordinal, len(recovered.markdown.splitlines()), ) results.append( await self.finalize( detected=d, region_bbox=region, primary=recovered, matched_chunks=[recovered], orphan_queue=[], flavor="stream", source=source, page=page, page_image=page_image, ) ) covered.add(di) continue logger.warning( "page {}: table {} '{}' DETECTED-NOT-EXTRACTED (detector found it; " "lattice, stream, and recovery all produced nothing)", page, d.ordinal, d.description[:60], ) results.append( ExtractedTable( title=d.description, markdown="", page=page, source=source, extraction_record=ExtractionRecord( status="detected_not_extracted", detected_ordinal=d.ordinal, detected_description=d.description, ), ) ) return results async def apply_structure_correction( self, cells: List[List[str]], page_image: Optional[Path], bbox: Optional[Tuple[float, float, float, float]], source: str, page: int, ) -> Optional[StructureCorrection]: """Config gate over `correction.correct_structure`: skipped when correction is disabled or no LLM is wired. The shared semaphore is created lazily inside the running loop so the extractor stays loop-agnostic. """ if self.llm is None or not self.run_llm_correction: return None if self.correct_sem is None: self.correct_sem = asyncio.Semaphore(self.max_concurrent) return await correct_structure( cells=cells, page_image=page_image, bbox=bbox, source=source, page=page, llm=self.llm, correct_sem=self.correct_sem, dpi=self.dpi, ) async def finalize( self, detected: DetectedTable, region_bbox: Optional[Tuple[float, float, float, float]], primary: CamelotCandidate, matched_chunks: List[CamelotCandidate], orphan_queue: List[CamelotCandidate], flavor: Flavor, source: str, page: int, page_image: Optional[Path], ) -> ExtractedTable: status: Literal["extracted", "incomplete"] = "extracted" complete: Optional[bool] = None gap: str = "" used = list(matched_chunks) filled: List[FilledCell] = [] # The structured grid is canonical here: assembly, the completeness # audit and the fill all operate on it; markdown is rendered once # by make_extracted. cells = assemble_cells(used) if self.run_completeness: # The audit reads the PDF text layer (the same source Camelot # reads) within `region_bbox`, which is bounded by the neighbor # tables so an adjacent table's rows are not read as a # continuation. `extracted_bboxes` gives the span we actually # captured, so the audit can place any missing figure at the # top or bottom edge. verdict = None while True: verdict = await self.completeness.audit( source=Path(source), page=page, region_bbox=region_bbox, extracted_bboxes=[c.bbox for c in used], assembled_markdown=grid_to_markdown(cells), ) complete, gap = verdict.complete, verdict.gap if complete or not orphan_queue: break # Truncated at an edge: extend with the next orphan chunk # (shared, consumed once across the page so no table is # polluted), then re-audit. nxt = orphan_queue.pop(0) used.append(nxt) cells = assemble_cells(used) logger.info( "page {}: table {} truncated; extending with {} and re-auditing", page, detected.ordinal, nxt.candidate_id, ) # Text-layer fill (round-off), the last-resort straggler cleanup: # no Camelot chunk could supply the truncated edge row, but the # figures are in the text layer (that is how the audit found the # gap). Insert them into the grid in place, re-audit, and record # the amendment as provenance — never blurred with Camelot's # cells. if not complete and self.run_fill and verdict is not None and verdict.missing_figures: _, page_h_pts, words = await asyncio.to_thread(page_words, Path(source), page) span = extracted_span_pts([c.bbox for c in used], page_h_pts) if span is not None: new_cells, placed = round_off_grid(cells, words, span, verdict.missing_figures) if placed: cells = new_cells filled = [ FilledCell( value=p.value, column=p.column, row_label=p.row_label, edge=p.edge, x=p.x, y=p.y, ) for p in placed ] reverdict = await self.completeness.audit( source=Path(source), page=page, region_bbox=region_bbox, extracted_bboxes=[c.bbox for c in used], assembled_markdown=grid_to_markdown(cells), ) complete, gap = reverdict.complete, reverdict.gap logger.info( "page {}: table {} ROUNDED OFF; filled {} text-layer figure(s); now {}", page, detected.ordinal, len(filled), "complete" if complete else "still incomplete", ) if not complete: status = "incomplete" logger.warning( "page {}: table {} INCOMPLETE; {} (gap reported)", page, detected.ordinal, gap or "edge truncated", ) correction = await self.apply_structure_correction(cells, page_image, primary.bbox, source, page) return self.make_extracted( cand=primary, cells=cells, detected=detected, flavor=flavor, source=source, source_ids=[c.candidate_id for c in used], status=status, completeness_complete=complete, completeness_gap=gap, filled_cells=filled, correction=correction, ) def make_extracted( self, cand: CamelotCandidate, cells: List[List[str]], detected: DetectedTable, flavor: Flavor, source: str, source_ids: List[str], status: Literal["extracted", "incomplete"], completeness_complete: Optional[bool] = None, completeness_gap: str = "", combined: bool = False, combined_ordinals: Optional[List[int]] = None, filled_cells: Optional[List[FilledCell]] = None, correction: Optional[StructureCorrection] = None, ) -> ExtractedTable: filled_cells = filled_cells or [] # Default presentation is the deterministic grid render with the # detector's description as title. When structure correction was # accepted, the LLM's restructured markdown and heading metadata take # over; the LLM title wins, falling back to the detector description # when the model returns none. title = detected.description subtitle = "" caption = "" footnotes: List[FootnoteDef] = [] markdown = grid_to_markdown(cells) llm_corrected = False if correction is not None: title = correction.title or detected.description caption = correction.caption footnotes = list(correction.footnotes) markdown = correction.markdown llm_corrected = correction.llm_corrected return ExtractedTable( title=title, subtitle=subtitle, caption=caption, markdown=markdown, footnotes=footnotes, page=cand.page, bbox=cand.bbox, flavor=flavor, source=source, camelot_accuracy=cand.accuracy, llm_corrected=llm_corrected, extraction_record=ExtractionRecord( status=status, detected_ordinal=detected.ordinal, detected_description=detected.description, matched_flavor=flavor, source_candidate_ids=source_ids, combined=combined, combined_ordinals=combined_ordinals or [], completeness_complete=completeness_complete, completeness_gap=completeness_gap, filled_from_text_layer=bool(filled_cells), filled_cells=filled_cells, ), ) def extract_tables_sync(self, source: Path) -> List[ExtractedTable]: """Sync entry point for callers without an event loop (CLI, scripts). Wraps the async path with `asyncio.run`. """ return asyncio.run(self.extract_tables(source))