""" The extraction pipeline as a pydantic-graph async graph. The pipeline is orchestrated as a pydantic-graph [`GraphBuilder`](https://ai.pydantic.dev/graph/) async graph (the current API; the ``Graph(nodes=[...])``/``BaseNode`` runner is deprecated). Stage-to-step mapping: start -> acquire page render ‖ Camelot flavors (concurrent), content-empty shells dropped -> map -> classify one step run per candidate -> join collect classifier decisions -> group_pages drop rejects, group survivors per page -> map -> unify_page one step run per page (single-candidate pages bypass the unifier) -> join flatten per-page rows -> log_unified log accepted -> unified table counts at the stage boundary -> map -> correct one step run per unified table -> join -> finalize deterministic (page, row) ordering -> end Per-stage fan-out happens through graph ``map`` edges (a task per item); LLM concurrency stays bounded by per-stage semaphores carried in `PipelineDeps`. Joins use ``preferred_parent_fork='closest'`` so each map edge pairs with its own join. Keep that one-join-per-map shape: a join fed by other joins across the broadcast fork fires before all branches have finished, silently dropping results. """ from __future__ import annotations import asyncio from dataclasses import dataclass, field from functools import cache from pathlib import Path from typing import List, Tuple from loguru import logger from pydantic import BaseModel, ConfigDict from pydantic_graph import GraphBuilder, StepContext, reduce_list_append, reduce_list_extend from pydantic_graph.graph_builder import Graph from pydantic_graph.id_types import JoinID, NodeID from quber.agents.classifier import ClassifierResult, TableClassifier from quber.agents.llm_client import FootnoteDef, LLMClient from quber.agents.unifier import CandidateInput, TableUnifier from quber.core.extractors.base import ExtractedTable from quber.core.extractors.camelot.acquire import ( CamelotCandidate, is_content_empty, render_pages, run_camelot_flavors_parallel, ) # --- Graph payload models (frozen: items flowing along graph edges) --------- class SequencedCandidate(BaseModel): """A Camelot candidate plus its position in Camelot's output order. `seq` exists because graph joins collect in completion order: without it, the unifier would see each page's candidates in a nondeterministic order, and the unifier LLM's output is sensitive to candidate order (observed on Visa Q2FY25 page 3: shuffled input produced digit-altered merges). Sorting on `seq` restores the sequential pipeline's exact LLM-input ordering. """ model_config = ConfigDict(frozen=True) seq: int candidate: CamelotCandidate class ClassifiedCandidate(BaseModel): """A Camelot candidate paired with its classifier decision.""" model_config = ConfigDict(frozen=True) seq: int candidate: CamelotCandidate decision: ClassifierResult class PageGroup(BaseModel): """All classifier-accepted candidates on one physical page.""" model_config = ConfigDict(frozen=True) page: int items: Tuple[ClassifiedCandidate, ...] class UnifiedTableRow(BaseModel): """One canonical table emitted by the unify stage. `page` + `row_idx` form the deterministic output ordering key — the graph's correction fan-out completes in arbitrary order, so the final step sorts on (page, row_idx) to reproduce the sequential pipeline's output order exactly. """ model_config = ConfigDict(frozen=True) page: int row_idx: int candidate: CamelotCandidate decision: ClassifierResult markdown: str source_ids: Tuple[str, ...] class CorrectedTableItem(BaseModel): """A finished ExtractedTable plus its ordering key.""" model_config = ConfigDict(frozen=True) page: int row_idx: int table: ExtractedTable # --- Graph state and dependencies ------------------------------------------- @dataclass class PipelineState: """Mutable run-scoped state. `page_images` is written once by the acquire step (single writer); the counters feed stage-boundary logs. """ page_images: Tuple[Path, ...] = () accepted_count: int = 0 @dataclass(frozen=True) class PipelineDeps: """Run-scoped clients + config carried into every graph step. The per-stage semaphores bound LLM concurrency: the graph fans out one task per item, but LLM calls stay bounded at `max_concurrent` per stage. """ llm: LLMClient classifier: TableClassifier unifier: TableUnifier source: Path tmp_dir: Path dpi: int run_llm_correction: bool flavor_timeout_s: float classify_sem: asyncio.Semaphore = field(repr=False) unify_sem: asyncio.Semaphore = field(repr=False) correct_sem: asyncio.Semaphore = field(repr=False) # --- Graph steps ------------------------------------------------------------- async def acquire_candidates( ctx: StepContext[PipelineState, PipelineDeps, Path], ) -> List[SequencedCandidate]: """Render page images and run both Camelot flavors concurrently, then drop content-empty shells before they reach the classifier. The render and the Camelot subprocesses are independent I/O-bound work; overlapping them takes the render off the critical path. """ deps = ctx.deps page_images, candidates = await asyncio.gather( asyncio.to_thread(render_pages, deps.source, deps.dpi, deps.tmp_dir), asyncio.to_thread(run_camelot_flavors_parallel, deps.source, deps.flavor_timeout_s), ) ctx.state.page_images = tuple(page_images) logger.info( "camelot produced {} candidates ({} lattice, {} stream) from {}", len(candidates), sum(1 for c in candidates if c.flavor == "lattice"), sum(1 for c in candidates if c.flavor == "stream"), deps.source.name, ) populated = [c for c in candidates if not is_content_empty(c.markdown)] dropped = len(candidates) - len(populated) if dropped: logger.info("dropped {} content-empty candidates pre-classifier", dropped) return [SequencedCandidate(seq=seq, candidate=c) for seq, c in enumerate(populated)] async def classify_candidate( ctx: StepContext[PipelineState, PipelineDeps, SequencedCandidate], ) -> ClassifiedCandidate: item = ctx.inputs async with ctx.deps.classify_sem: decision = await ctx.deps.classifier.classify(item.candidate.markdown) return ClassifiedCandidate(seq=item.seq, candidate=item.candidate, decision=decision) async def group_pages( ctx: StepContext[PipelineState, PipelineDeps, List[ClassifiedCandidate]], ) -> List[PageGroup]: """Drop classifier rejects and group survivors per page (ascending). Items are re-sorted on `seq` inside each page: the classify join collects in completion order, and the unifier LLM's behavior is order-sensitive, so each page must present its candidates in Camelot's output order. """ accepted = sorted((item for item in ctx.inputs if item.decision.is_table), key=lambda item: item.seq) rejected = len(ctx.inputs) - len(accepted) logger.info("classifier accepted {} / rejected {}", len(accepted), rejected) ctx.state.accepted_count = len(accepted) by_page: dict[int, List[ClassifiedCandidate]] = {} for item in accepted: by_page.setdefault(item.candidate.page, []).append(item) return [PageGroup(page=page, items=tuple(items)) for page, items in sorted(by_page.items())] async def unify_page( ctx: StepContext[PipelineState, PipelineDeps, PageGroup], ) -> List[UnifiedTableRow]: """Unify the flavor outputs on one page into canonical tables. Single-candidate pages bypass the unifier (nothing to combine); out-of-range pages are dropped with a warning. """ group = ctx.inputs page_images = ctx.state.page_images if group.page < 1 or group.page > len(page_images): logger.warning("unify: page {} out of range; skipping", group.page) return [] if len(group.items) == 1: item = group.items[0] return [ UnifiedTableRow( page=group.page, row_idx=0, candidate=item.candidate, decision=item.decision, markdown=item.candidate.markdown, source_ids=(item.candidate.candidate_id,), ) ] page_image = page_images[group.page - 1] unifier_input = [ CandidateInput( candidate_id=item.candidate.candidate_id, markdown=item.candidate.markdown, bbox=item.candidate.bbox, ) for item in group.items ] async with ctx.deps.unify_sem: result = await ctx.deps.unifier.unify(page_image, unifier_input) item_by_id = {item.candidate.candidate_id: item for item in group.items} rows: List[UnifiedTableRow] = [] for row_idx, table in enumerate(result.tables): primary_id = table.source_candidate_ids[0] if table.source_candidate_ids else None primary = item_by_id.get(primary_id) if primary_id else None if primary is None: primary = group.items[0] rows.append( UnifiedTableRow( page=group.page, row_idx=row_idx, candidate=primary.candidate, decision=primary.decision, markdown=table.markdown, source_ids=tuple(table.source_candidate_ids), ) ) return rows async def log_unified( ctx: StepContext[PipelineState, PipelineDeps, List[UnifiedTableRow]], ) -> List[UnifiedTableRow]: """Stage boundary: all pages unified, corrections not yet started.""" logger.info( "unifier combined {} accepted -> {} unified tables", ctx.state.accepted_count, len(ctx.inputs), ) return ctx.inputs async def correct_table( ctx: StepContext[PipelineState, PipelineDeps, UnifiedTableRow], ) -> CorrectedTableItem: """Per-table LLM structure correction. The unifier's markdown is the input; the page image is reference. Camelot-empty input -> no correction (strict source-of-truth, no image-reading rescue). """ row = ctx.inputs deps = ctx.deps page_images = ctx.state.page_images if row.page < 1 or row.page > len(page_images): page_image = None else: page_image = page_images[row.page - 1] title = "" subtitle = "" footnotes: List[FootnoteDef] = [] corrected_md = row.markdown llm_corrected = False if deps.run_llm_correction and page_image is not None and not is_content_empty(row.markdown): async with deps.correct_sem: try: correction = await deps.llm.correct_structure(page_image, row.markdown) except Exception as exc: logger.error( "LLM correction failed page={} candidate={}: {}", row.page, row.candidate.candidate_id, exc, ) correction = None if correction is not None: title = correction.title subtitle = "" footnotes = list(correction.footnotes) if correction.markdown and correction.markdown.strip(): corrected_md = correction.markdown llm_corrected = corrected_md != row.markdown table = ExtractedTable( title=title, subtitle=subtitle, markdown=corrected_md, footnotes=footnotes, page=row.page, bbox=row.candidate.bbox, flavor=row.candidate.flavor, source=str(deps.source), camelot_accuracy=row.candidate.accuracy, classifier_decision=row.decision, llm_corrected=llm_corrected, ) return CorrectedTableItem(page=row.page, row_idx=row.row_idx, table=table) async def finalize( ctx: StepContext[PipelineState, PipelineDeps, List[CorrectedTableItem]], ) -> List[ExtractedTable]: """Restore deterministic output order: pages ascending, unifier row order within a page (the corrections join collects in completion order). """ ordered = sorted(ctx.inputs, key=lambda item: (item.page, item.row_idx)) return [item.table for item in ordered] @cache def build_pipeline_graph() -> Graph[PipelineState, PipelineDeps, Path, List[ExtractedTable]]: """Build the extraction pipeline graph (cached; the graph is immutable and steps receive everything run-scoped via state/deps). Topology notes: every map edge pairs with its own closest-fork join — see the module docstring for why this exact shape is load-bearing. """ g = GraphBuilder( name="camelot.llm_pipeline", state_type=PipelineState, deps_type=PipelineDeps, input_type=Path, output_type=List[ExtractedTable], ) acquire_step = g.step(acquire_candidates) classify_step = g.step(classify_candidate) group_step = g.step(group_pages) unify_step = g.step(unify_page) log_unified_step = g.step(log_unified) correct_step = g.step(correct_table) finalize_step = g.step(finalize) j_classified = g.join( reduce_list_append, initial_factory=list, node_id="j_classified", preferred_parent_fork="closest", ) j_rows = g.join( reduce_list_extend, initial_factory=list, node_id="j_rows", preferred_parent_fork="closest", ) j_corrected = g.join( reduce_list_append, initial_factory=list, node_id="j_corrected", preferred_parent_fork="closest", ) g.add_edge(g.start_node, acquire_step) g.add_mapping_edge(acquire_step, classify_step, downstream_join_id=JoinID(NodeID("j_classified"))) g.add_edge(classify_step, j_classified) g.add_edge(j_classified, group_step) g.add_mapping_edge(group_step, unify_step, downstream_join_id=JoinID(NodeID("j_rows"))) g.add_edge(unify_step, j_rows) g.add_edge(j_rows, log_unified_step) g.add_mapping_edge(log_unified_step, correct_step, downstream_join_id=JoinID(NodeID("j_corrected"))) g.add_edge(correct_step, j_corrected) g.add_edge(j_corrected, finalize_step) g.add_edge(finalize_step, g.end_node) return g.build()