Coverage for src / quber / core / extractors / camelot / llm / pipeline.py: 96%
164 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""
2The extraction pipeline as a pydantic-graph async graph.
4The pipeline is orchestrated as a pydantic-graph
5[`GraphBuilder`](https://ai.pydantic.dev/graph/) async graph (the current
6API; the ``Graph(nodes=[...])``/``BaseNode`` runner is deprecated).
7Stage-to-step mapping:
9 start -> acquire page render ‖ Camelot flavors (concurrent),
10 content-empty shells dropped
11 -> map -> classify one step run per candidate
12 -> join collect classifier decisions
13 -> group_pages drop rejects, group survivors per page
14 -> map -> unify_page one step run per page (single-candidate
15 pages bypass the unifier)
16 -> join flatten per-page rows
17 -> log_unified log accepted -> unified table counts at
18 the stage boundary
19 -> map -> correct one step run per unified table
20 -> join -> finalize deterministic (page, row) ordering
21 -> end
23Per-stage fan-out happens through graph ``map`` edges (a task per item);
24LLM concurrency stays bounded by per-stage semaphores carried in
25`PipelineDeps`. Joins use ``preferred_parent_fork='closest'`` so each map
26edge pairs with its own join. Keep that one-join-per-map shape: a join
27fed by other joins across the broadcast fork fires before all branches
28have finished, silently dropping results.
29"""
31from __future__ import annotations
33import asyncio
34from dataclasses import dataclass, field
35from functools import cache
36from pathlib import Path
37from typing import List, Tuple
39from loguru import logger
40from pydantic import BaseModel, ConfigDict
41from pydantic_graph import GraphBuilder, StepContext, reduce_list_append, reduce_list_extend
42from pydantic_graph.graph_builder import Graph
43from pydantic_graph.id_types import JoinID, NodeID
45from quber.agents.classifier import ClassifierResult, TableClassifier
46from quber.agents.llm_client import FootnoteDef, LLMClient
47from quber.agents.unifier import CandidateInput, TableUnifier
48from quber.core.extractors.base import ExtractedTable
49from quber.core.extractors.camelot.acquire import (
50 CamelotCandidate,
51 is_content_empty,
52 render_pages,
53 run_camelot_flavors_parallel,
54)
56# --- Graph payload models (frozen: items flowing along graph edges) ---------
59class SequencedCandidate(BaseModel):
60 """A Camelot candidate plus its position in Camelot's output order.
62 `seq` exists because graph joins collect in completion order: without
63 it, the unifier would see each page's candidates in a nondeterministic
64 order, and the unifier LLM's output is sensitive to candidate order
65 (observed on Visa Q2FY25 page 3: shuffled input produced digit-altered
66 merges). Sorting on `seq` restores the sequential pipeline's exact
67 LLM-input ordering.
68 """
70 model_config = ConfigDict(frozen=True)
72 seq: int
73 candidate: CamelotCandidate
76class ClassifiedCandidate(BaseModel):
77 """A Camelot candidate paired with its classifier decision."""
79 model_config = ConfigDict(frozen=True)
81 seq: int
82 candidate: CamelotCandidate
83 decision: ClassifierResult
86class PageGroup(BaseModel):
87 """All classifier-accepted candidates on one physical page."""
89 model_config = ConfigDict(frozen=True)
91 page: int
92 items: Tuple[ClassifiedCandidate, ...]
95class UnifiedTableRow(BaseModel):
96 """One canonical table emitted by the unify stage.
98 `page` + `row_idx` form the deterministic output ordering key — the
99 graph's correction fan-out completes in arbitrary order, so the final
100 step sorts on (page, row_idx) to reproduce the sequential pipeline's
101 output order exactly.
102 """
104 model_config = ConfigDict(frozen=True)
106 page: int
107 row_idx: int
108 candidate: CamelotCandidate
109 decision: ClassifierResult
110 markdown: str
111 source_ids: Tuple[str, ...]
114class CorrectedTableItem(BaseModel):
115 """A finished ExtractedTable plus its ordering key."""
117 model_config = ConfigDict(frozen=True)
119 page: int
120 row_idx: int
121 table: ExtractedTable
124# --- Graph state and dependencies -------------------------------------------
127@dataclass
128class PipelineState:
129 """Mutable run-scoped state. `page_images` is written once by the
130 acquire step (single writer); the counters feed stage-boundary logs.
131 """
133 page_images: Tuple[Path, ...] = ()
134 accepted_count: int = 0
137@dataclass(frozen=True)
138class PipelineDeps:
139 """Run-scoped clients + config carried into every graph step.
141 The per-stage semaphores preserve the pre-graph concurrency contract:
142 the graph fans out one task per item, but LLM calls stay bounded at
143 `max_concurrent` per stage.
144 """
146 llm: LLMClient
147 classifier: TableClassifier
148 unifier: TableUnifier
149 source: Path
150 tmp_dir: Path
151 dpi: int
152 run_llm_correction: bool
153 flavor_timeout_s: float
154 classify_sem: asyncio.Semaphore = field(repr=False)
155 unify_sem: asyncio.Semaphore = field(repr=False)
156 correct_sem: asyncio.Semaphore = field(repr=False)
159# --- Graph steps -------------------------------------------------------------
162async def acquire_candidates(
163 ctx: StepContext[PipelineState, PipelineDeps, Path],
164) -> List[SequencedCandidate]:
165 """Render page images and run both Camelot flavors concurrently, then
166 drop content-empty shells before they reach the classifier.
168 The render and the Camelot subprocesses are independent I/O-bound
169 work; overlapping them takes the render off the critical path (the
170 sequential pipeline ran them back to back).
171 """
172 deps = ctx.deps
173 page_images, candidates = await asyncio.gather(
174 asyncio.to_thread(render_pages, deps.source, deps.dpi, deps.tmp_dir),
175 asyncio.to_thread(run_camelot_flavors_parallel, deps.source, deps.flavor_timeout_s),
176 )
177 ctx.state.page_images = tuple(page_images)
178 logger.info(
179 "camelot produced {} candidates ({} lattice, {} stream) from {}",
180 len(candidates),
181 sum(1 for c in candidates if c.flavor == "lattice"),
182 sum(1 for c in candidates if c.flavor == "stream"),
183 deps.source.name,
184 )
186 populated = [c for c in candidates if not is_content_empty(c.markdown)]
187 dropped = len(candidates) - len(populated)
188 if dropped:
189 logger.info("dropped {} content-empty candidates pre-classifier", dropped)
190 return [SequencedCandidate(seq=seq, candidate=c) for seq, c in enumerate(populated)]
193async def classify_candidate(
194 ctx: StepContext[PipelineState, PipelineDeps, SequencedCandidate],
195) -> ClassifiedCandidate:
196 item = ctx.inputs
197 async with ctx.deps.classify_sem:
198 decision = await ctx.deps.classifier.classify(item.candidate.markdown)
199 return ClassifiedCandidate(seq=item.seq, candidate=item.candidate, decision=decision)
202async def group_pages(
203 ctx: StepContext[PipelineState, PipelineDeps, List[ClassifiedCandidate]],
204) -> List[PageGroup]:
205 """Drop classifier rejects and group survivors per page (ascending).
207 Items are re-sorted on `seq` inside each page: the classify join
208 collects in completion order, and the unifier LLM's behavior is
209 order-sensitive, so each page must present its candidates in
210 Camelot's original output order — exactly what the sequential
211 pipeline fed it.
212 """
213 accepted = sorted((item for item in ctx.inputs if item.decision.is_table), key=lambda item: item.seq)
214 rejected = len(ctx.inputs) - len(accepted)
215 logger.info("classifier accepted {} / rejected {}", len(accepted), rejected)
216 ctx.state.accepted_count = len(accepted)
218 by_page: dict[int, List[ClassifiedCandidate]] = {}
219 for item in accepted:
220 by_page.setdefault(item.candidate.page, []).append(item)
221 return [PageGroup(page=page, items=tuple(items)) for page, items in sorted(by_page.items())]
224async def unify_page(
225 ctx: StepContext[PipelineState, PipelineDeps, PageGroup],
226) -> List[UnifiedTableRow]:
227 """Unify the flavor outputs on one page into canonical tables.
229 Single-candidate pages bypass the unifier (nothing to combine);
230 out-of-range pages are dropped with a warning, exactly like the
231 sequential pipeline.
232 """
233 group = ctx.inputs
234 page_images = ctx.state.page_images
235 if group.page < 1 or group.page > len(page_images):
236 logger.warning("unify: page {} out of range; skipping", group.page)
237 return []
239 if len(group.items) == 1:
240 item = group.items[0]
241 return [
242 UnifiedTableRow(
243 page=group.page,
244 row_idx=0,
245 candidate=item.candidate,
246 decision=item.decision,
247 markdown=item.candidate.markdown,
248 source_ids=(item.candidate.candidate_id,),
249 )
250 ]
252 page_image = page_images[group.page - 1]
253 unifier_input = [
254 CandidateInput(
255 candidate_id=item.candidate.candidate_id,
256 markdown=item.candidate.markdown,
257 bbox=item.candidate.bbox,
258 )
259 for item in group.items
260 ]
261 async with ctx.deps.unify_sem:
262 result = await ctx.deps.unifier.unify(page_image, unifier_input)
264 item_by_id = {item.candidate.candidate_id: item for item in group.items}
265 rows: List[UnifiedTableRow] = []
266 for row_idx, table in enumerate(result.tables):
267 primary_id = table.source_candidate_ids[0] if table.source_candidate_ids else None
268 primary = item_by_id.get(primary_id) if primary_id else None
269 if primary is None:
270 primary = group.items[0]
271 rows.append(
272 UnifiedTableRow(
273 page=group.page,
274 row_idx=row_idx,
275 candidate=primary.candidate,
276 decision=primary.decision,
277 markdown=table.markdown,
278 source_ids=tuple(table.source_candidate_ids),
279 )
280 )
281 return rows
284async def log_unified(
285 ctx: StepContext[PipelineState, PipelineDeps, List[UnifiedTableRow]],
286) -> List[UnifiedTableRow]:
287 """Stage boundary: all pages unified, corrections not yet started."""
288 logger.info(
289 "unifier combined {} accepted -> {} unified tables",
290 ctx.state.accepted_count,
291 len(ctx.inputs),
292 )
293 return ctx.inputs
296async def correct_table(
297 ctx: StepContext[PipelineState, PipelineDeps, UnifiedTableRow],
298) -> CorrectedTableItem:
299 """Per-table LLM structure correction. The unifier's markdown is the
300 input; the page image is reference. Camelot-empty input -> no
301 correction (strict source-of-truth, no image-reading rescue).
302 """
303 row = ctx.inputs
304 deps = ctx.deps
305 page_images = ctx.state.page_images
306 if row.page < 1 or row.page > len(page_images):
307 page_image = None
308 else:
309 page_image = page_images[row.page - 1]
311 title = ""
312 subtitle = ""
313 footnotes: List[FootnoteDef] = []
314 corrected_md = row.markdown
315 llm_corrected = False
317 if deps.run_llm_correction and page_image is not None and not is_content_empty(row.markdown):
318 async with deps.correct_sem:
319 try:
320 correction = await deps.llm.correct_structure(page_image, row.markdown)
321 except Exception as exc:
322 logger.error(
323 "LLM correction failed page={} candidate={}: {}",
324 row.page,
325 row.candidate.candidate_id,
326 exc,
327 )
328 correction = None
329 if correction is not None:
330 title = correction.title
331 subtitle = ""
332 footnotes = list(correction.footnotes)
333 if correction.markdown and correction.markdown.strip():
334 corrected_md = correction.markdown
335 llm_corrected = corrected_md != row.markdown
337 table = ExtractedTable(
338 title=title,
339 subtitle=subtitle,
340 markdown=corrected_md,
341 footnotes=footnotes,
342 page=row.page,
343 bbox=row.candidate.bbox,
344 flavor=row.candidate.flavor,
345 source=str(deps.source),
346 camelot_accuracy=row.candidate.accuracy,
347 classifier_decision=row.decision,
348 llm_corrected=llm_corrected,
349 )
350 return CorrectedTableItem(page=row.page, row_idx=row.row_idx, table=table)
353async def finalize(
354 ctx: StepContext[PipelineState, PipelineDeps, List[CorrectedTableItem]],
355) -> List[ExtractedTable]:
356 """Restore deterministic output order: pages ascending, unifier row
357 order within a page (the corrections join collects in completion
358 order).
359 """
360 ordered = sorted(ctx.inputs, key=lambda item: (item.page, item.row_idx))
361 return [item.table for item in ordered]
364@cache
365def build_pipeline_graph() -> Graph[PipelineState, PipelineDeps, Path, List[ExtractedTable]]:
366 """Build the extraction pipeline graph (cached; the graph is immutable
367 and steps receive everything run-scoped via state/deps).
369 Topology notes: every map edge pairs with its own closest-fork join —
370 see the module docstring for why this exact shape is load-bearing.
371 """
372 g = GraphBuilder(
373 name="camelot.llm_pipeline",
374 state_type=PipelineState,
375 deps_type=PipelineDeps,
376 input_type=Path,
377 output_type=List[ExtractedTable],
378 )
380 acquire_step = g.step(acquire_candidates)
381 classify_step = g.step(classify_candidate)
382 group_step = g.step(group_pages)
383 unify_step = g.step(unify_page)
384 log_unified_step = g.step(log_unified)
385 correct_step = g.step(correct_table)
386 finalize_step = g.step(finalize)
388 j_classified = g.join(
389 reduce_list_append,
390 initial_factory=list,
391 node_id="j_classified",
392 preferred_parent_fork="closest",
393 )
394 j_rows = g.join(
395 reduce_list_extend,
396 initial_factory=list,
397 node_id="j_rows",
398 preferred_parent_fork="closest",
399 )
400 j_corrected = g.join(
401 reduce_list_append,
402 initial_factory=list,
403 node_id="j_corrected",
404 preferred_parent_fork="closest",
405 )
407 g.add_edge(g.start_node, acquire_step)
408 g.add_mapping_edge(acquire_step, classify_step, downstream_join_id=JoinID(NodeID("j_classified")))
409 g.add_edge(classify_step, j_classified)
410 g.add_edge(j_classified, group_step)
411 g.add_mapping_edge(group_step, unify_step, downstream_join_id=JoinID(NodeID("j_rows")))
412 g.add_edge(unify_step, j_rows)
413 g.add_edge(j_rows, log_unified_step)
414 g.add_mapping_edge(log_unified_step, correct_step, downstream_join_id=JoinID(NodeID("j_corrected")))
415 g.add_edge(correct_step, j_corrected)
416 g.add_edge(j_corrected, finalize_step)
417 g.add_edge(finalize_step, g.end_node)
419 return g.build()