Coverage for src / quber / core / fusion / fusion.py: 31%
83 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Fuse the document and table extractions into the two corrected outputs.
3The fusion step: given a `ParseResult` (the document extraction) and the
4Set-of-Mark/Camelot tables (the table extraction), match the two by overlap,
5split any region SoM merged, annotate charts and image tables, and graft the
6Camelot bodies into the docling spine. The result carries both corrected
7outputs: the SoM/Camelot tables (split, annotated) and the unified
8`DoclingDocument`.
10`fuse_artifacts` is the standalone core: it takes already-produced data and the
11source PDF and knows nothing about how the data was produced, so a cloud job can
12run it against artifacts loaded from S3. `DocumentFusion` is the local control
13flow: it runs the document and table extractions in-process (concurrently — they
14are independent), then calls `fuse_artifacts`.
15"""
17from __future__ import annotations
19import asyncio
20from pathlib import Path
21from typing import Dict, List, Optional, Tuple
23from docling_core.types.doc.base import CoordOrigin
24from docling_core.types.doc.document import TableItem
25from loguru import logger
27from quber.agents.llm_client import LLMClient, get_llm_client
28from quber.core.extractors.base import ExtractedTable
29from quber.core.extractors.camelot.correspondence.geometry import camelot_bbox_to_norm
30from quber.core.extractors.set_of_mark import SetOfMarkExtractor
31from quber.core.extractors.set_of_mark.split import split_table
32from quber.core.fusion.graft import build_unified_document
33from quber.core.fusion.heading_review import review_headings
34from quber.core.fusion.matching import match_tables
35from quber.core.fusion.models import FusionResult, RegionMatch
36from quber.core.parsers import Parser, ParseResult
38NormBox = Tuple[float, float, float, float]
41async def fuse_artifacts(
42 parse: ParseResult,
43 som_tables: List[ExtractedTable],
44 source: Path,
45 llm: LLMClient,
46) -> FusionResult:
47 """Fuse already-produced document and table extractions into both outputs.
49 Takes the document extraction (`parse`), the table extraction (`som_tables`),
50 and the source PDF (needed to re-extract sub-regions when a merged region is
51 split). Knows nothing about how the inputs were produced, so the same call
52 serves the local flow and a cloud job loading artifacts from S3.
53 """
54 page_dims = {p.page_no: (p.width, p.height) for p in parse.pages}
55 tables_by_ref = {t.self_ref: t for t in parse.document.tables}
57 matches = match_tables(parse, som_tables)
59 # Split each region SoM merged (docling found more tables than SoM). The
60 # merged SoM table is replaced in place by its splits, then the whole set is
61 # re-matched so the region reads as a 1:1 replace.
62 som_tables, split_any = await _split_merged_regions(
63 matches, som_tables, source, tables_by_ref, page_dims, llm
64 )
65 if split_any:
66 matches = match_tables(parse, som_tables)
68 # Annotate charts and image tables on the SoM output.
69 _annotate(matches, som_tables)
71 # Graft the Camelot bodies into a clone of the docling spine.
72 unified, errors = build_unified_document(parse.document, som_tables, matches, page_dims)
73 for err in errors:
74 logger.error("fuse: {}", err)
76 # Demote page decoration the parser labeled as section headings, so no
77 # consumer carries a repeated banner or a leaked column label as context.
78 heading_flags = await review_headings(unified, llm, source=str(source))
80 return FusionResult(
81 document=unified,
82 tables=som_tables,
83 matches=matches,
84 errors=errors,
85 heading_flags=heading_flags,
86 parse=parse,
87 )
90class DocumentFusion:
91 """Local control flow: run both extractions in-process, then fuse them.
93 Holds the parser, the table extractor, and the LLM client used by the split
94 pass. `fuse` runs the whole flow; `fuse_sync` wraps it for the CLI. The cloud
95 flow does not use this class — it runs the document and table extractions as
96 separate jobs and calls `fuse_artifacts` directly on the loaded artifacts.
97 """
99 def __init__(
100 self,
101 parser: Optional[Parser] = None,
102 extractor: Optional[SetOfMarkExtractor] = None,
103 llm: Optional[LLMClient] = None,
104 preset: str = "tuned-financial",
105 ) -> None:
106 if parser is None:
107 # Engine import deferred: building a default parser needs the
108 # full docling package, but artifact-fed fusion (the CPU cloud
109 # job) always passes no engines and must not import it.
110 from quber.core.parsers import parser_for_preset
112 parser = parser_for_preset(preset)
113 self.parser = parser
114 self.llm = llm or get_llm_client(None)
115 self.extractor = extractor or SetOfMarkExtractor(llm=self.llm)
117 async def fuse(self, source: Path) -> FusionResult:
118 # docling parse and SoM extraction are independent; run them together.
119 parse_task = asyncio.to_thread(self.parser.parse, source)
120 som_tables = await self.extractor.extract_tables(source)
121 parse: ParseResult = await parse_task
122 return await fuse_artifacts(parse, som_tables, source, self.llm)
124 def fuse_sync(self, source: Path) -> FusionResult:
125 return asyncio.run(self.fuse(source))
128async def _split_merged_regions(
129 matches: List[RegionMatch],
130 som_tables: List[ExtractedTable],
131 source: Path,
132 tables_by_ref: Dict[str, TableItem],
133 page_dims: Dict[int, Tuple[float, float]],
134 llm: LLMClient,
135) -> Tuple[List[ExtractedTable], bool]:
136 """Replace each SoM-merged table with the tables `split_table` recovers.
138 Only the single-SoM-table case is handled (one fused region docling sees as
139 several); a region with more than one SoM table is too ambiguous to split and
140 is left for the graft pass to report.
141 """
142 out = list(som_tables)
143 split_any = False
144 for match in matches:
145 if match.kind != "som_merged" or len(match.som_indices) != 1:
146 continue
147 idx = match.som_indices[0]
148 count = len(match.docling_table_refs)
149 boundaries = [
150 _docling_norm_box(tables_by_ref[r], page_dims, match.page)
151 for r in match.docling_table_refs
152 if r in tables_by_ref
153 ]
154 boundaries = [b for b in boundaries if b is not None]
155 if len(boundaries) != count:
156 continue
157 # split_table reaches the source PDF through the table's own `source`
158 # field; keep it pointed at the resolved local path.
159 out[idx].source = str(source)
160 splits = await split_table(out[idx], count, boundaries, llm)
161 if len(splits) > 1:
162 out[idx] = splits[0]
163 # Insert the remaining splits right after, preserving order.
164 for offset, extra in enumerate(splits[1:], start=1):
165 out.insert(idx + offset, extra)
166 split_any = True
167 logger.info(
168 "fuse: split SoM-merged region on page {} into {} tables",
169 match.page,
170 len(splits),
171 )
172 return out, split_any
175def _annotate(matches: List[RegionMatch], som_tables: List[ExtractedTable]) -> None:
176 """Set `ExtractedTable.kind` for image tables and charts from the matches."""
177 for match in matches:
178 if match.kind == "image_table":
179 for i in match.som_indices:
180 som_tables[i].kind = "image_table"
181 elif match.kind == "chart":
182 for i in match.som_indices:
183 som_tables[i].kind = "chart"
186def _docling_norm_box(
187 table: TableItem, page_dims: Dict[int, Tuple[float, float]], page: int
188) -> Optional[NormBox]:
189 """A docling table's box in the normalized top-left frame split_table expects."""
190 if not table.prov:
191 return None
192 width, height = page_dims.get(page, (612.0, 792.0))
193 bbox = table.prov[0].bbox
194 if bbox.coord_origin == CoordOrigin.TOPLEFT:
195 return (bbox.l / width, bbox.t / height, bbox.r / width, bbox.b / height)
196 return camelot_bbox_to_norm((bbox.l, bbox.b, bbox.r, bbox.t), width, height)