"""Fuse the document and table extractions into the two corrected outputs. The fusion step: given a `ParseResult` (the document extraction) and the Set-of-Mark/Camelot tables (the table extraction), match the two by overlap, split any region SoM merged, annotate charts and image tables, and graft the Camelot bodies into the docling spine. An LLM heading review then demotes page decoration the parser labeled as section headings in the unified document, and reports each demotion in `heading_flags`. The result carries both corrected outputs: the SoM/Camelot tables (split, annotated) and the unified `DoclingDocument`. `fuse_artifacts` is the standalone core: it takes already-produced data and the source PDF and knows nothing about how the data was produced, so a cloud job can run it against artifacts loaded from S3. `DocumentFusion` is the local control flow: it runs the document and table extractions in-process, one after the other, then calls `fuse_artifacts`. """ from __future__ import annotations import asyncio from pathlib import Path from typing import Dict, List, Optional, Tuple from docling_core.types.doc.base import CoordOrigin from docling_core.types.doc.document import TableItem from loguru import logger from quber.agents.llm_client import LLMClient, get_llm_client from quber.core.extractors.base import ExtractedTable from quber.core.extractors.camelot.correspondence.geometry import camelot_bbox_to_norm from quber.core.extractors.set_of_mark import SetOfMarkExtractor from quber.core.extractors.set_of_mark.split import split_table from quber.core.fusion.graft import build_unified_document from quber.core.fusion.heading_review import review_headings from quber.core.fusion.matching import match_tables from quber.core.fusion.models import FusionResult, RegionMatch from quber.core.parsers import Parser, ParseResult NormBox = Tuple[float, float, float, float] async def fuse_artifacts( parse: ParseResult, som_tables: List[ExtractedTable], source: Path, llm: LLMClient, ) -> FusionResult: """Fuse already-produced document and table extractions into both outputs. Takes the document extraction (`parse`), the table extraction (`som_tables`), and the source PDF (needed to re-extract sub-regions when a merged region is split). Knows nothing about how the inputs were produced, so the same call serves the local flow and a cloud job loading artifacts from S3. """ page_dims = {p.page_no: (p.width, p.height) for p in parse.pages} tables_by_ref = {t.self_ref: t for t in parse.document.tables} matches = match_tables(parse, som_tables) # Split each region SoM merged (docling found more tables than SoM). The # merged SoM table is replaced in place by its splits, then the whole set is # re-matched so the region reads as a 1:1 replace. som_tables, split_any = await _split_merged_regions( matches, som_tables, source, tables_by_ref, page_dims, llm ) if split_any: matches = match_tables(parse, som_tables) # Annotate charts and image tables on the SoM output. _annotate(matches, som_tables) # Graft the Camelot bodies into a clone of the docling spine. unified, errors = build_unified_document(parse.document, som_tables, matches, page_dims) for err in errors: logger.error("fuse: {}", err) # Demote page decoration the parser labeled as section headings, so no # consumer carries a repeated banner or a leaked column label as context. heading_flags = await review_headings(unified, llm, source=str(source)) return FusionResult( document=unified, tables=som_tables, matches=matches, errors=errors, heading_flags=heading_flags, parse=parse, ) class DocumentFusion: """Local control flow: run both extractions in-process, then fuse them. Holds the parser, the table extractor, and the LLM client. The client serves the split pass, the heading review, and the default `SetOfMarkExtractor` built when no extractor is passed. `fuse` runs the whole flow; `fuse_sync` wraps it for the CLI. The cloud flow does not use this class — it runs the document and table extractions as separate jobs and calls `fuse_artifacts` directly on the loaded artifacts. """ def __init__( self, parser: Optional[Parser] = None, extractor: Optional[SetOfMarkExtractor] = None, llm: Optional[LLMClient] = None, preset: str = "tuned-financial", ) -> None: if parser is None: # Engine import deferred: building a default parser needs the # full docling package. Importing this module for `fuse_artifacts` # (the artifact-fed path, which never constructs this class) must # not pull it in. Constructing this class without a parser does. from quber.core.parsers import parser_for_preset parser = parser_for_preset(preset) self.parser = parser self.llm = llm or get_llm_client(None) self.extractor = extractor or SetOfMarkExtractor(llm=self.llm) async def fuse(self, source: Path) -> FusionResult: # The two extractions run one after the other: SoM extraction first, # then the docling parse in a worker thread. `asyncio.to_thread` only # builds a coroutine, so the parse starts when it is awaited. parse_task = asyncio.to_thread(self.parser.parse, source) som_tables = await self.extractor.extract_tables(source) parse: ParseResult = await parse_task return await fuse_artifacts(parse, som_tables, source, self.llm) def fuse_sync(self, source: Path) -> FusionResult: return asyncio.run(self.fuse(source)) async def _split_merged_regions( matches: List[RegionMatch], som_tables: List[ExtractedTable], source: Path, tables_by_ref: Dict[str, TableItem], page_dims: Dict[int, Tuple[float, float]], llm: LLMClient, ) -> Tuple[List[ExtractedTable], bool]: """Replace each SoM-merged table with the tables `split_table` recovers. Only the single-SoM-table case is handled (one fused region docling sees as several); a region with more than one SoM table is too ambiguous to split and is left for the graft pass to report. """ out = list(som_tables) split_any = False for match in matches: if match.kind != "som_merged" or len(match.som_indices) != 1: continue idx = match.som_indices[0] count = len(match.docling_table_refs) boundaries = [ _docling_norm_box(tables_by_ref[r], page_dims, match.page) for r in match.docling_table_refs if r in tables_by_ref ] boundaries = [b for b in boundaries if b is not None] if len(boundaries) != count: continue # split_table reaches the source PDF through the table's own `source` # field; keep it pointed at the resolved local path. out[idx].source = str(source) splits = await split_table(out[idx], count, boundaries, llm) if len(splits) > 1: out[idx] = splits[0] # Insert the remaining splits right after, preserving order. for offset, extra in enumerate(splits[1:], start=1): out.insert(idx + offset, extra) split_any = True logger.info( "fuse: split SoM-merged region on page {} into {} tables", match.page, len(splits), ) return out, split_any def _annotate(matches: List[RegionMatch], som_tables: List[ExtractedTable]) -> None: """Set `ExtractedTable.kind` for image tables and charts from the matches.""" for match in matches: if match.kind == "image_table": for i in match.som_indices: som_tables[i].kind = "image_table" elif match.kind == "chart": for i in match.som_indices: som_tables[i].kind = "chart" def _docling_norm_box( table: TableItem, page_dims: Dict[int, Tuple[float, float]], page: int ) -> Optional[NormBox]: """A docling table's box in the normalized top-left frame split_table expects.""" if not table.prov: return None width, height = page_dims.get(page, (612.0, 792.0)) bbox = table.prov[0].bbox if bbox.coord_origin == CoordOrigin.TOPLEFT: return (bbox.l / width, bbox.t / height, bbox.r / width, bbox.b / height) return camelot_bbox_to_norm((bbox.l, bbox.b, bbox.r, bbox.t), width, height)