""" GridLocator — Set-of-Mark table flagging. The grid locator finds every table and its region for the default Set-of-Mark engine and for the dual engine. The camelot correspondence engine also calls it lazily, when a detector box has drifted off its table. The raw detector bbox drifts: a vision model asked for continuous coordinates emits a discretized, often evenly-spaced grid that walks off the actual tables on repetitive/stacked layouts (measured on stacked4 and real TMUS/Visa pages). This module sidesteps that failure mode. Instead of asking the model for coordinates, it overlays a *labeled* grid on the rasterized page (numbered rows, lettered columns) and asks which grid cells each table occupies. The model reports discrete row/column IDs — reading printed labels, which it does reliably — and the grid, being ours, maps each ID back to an exact coordinate. The coarse cell region is then tightened to the page's text layer for a precise box. This follows the Set-of-Mark / coordinate-scaffold line of visual prompting: - Yang, Zhang, Li, Zou, Li, Gao, "Set-of-Mark Prompting Unleashes Extraordinary Visual Grounding in GPT-4V," arXiv:2310.11441 (2023): overlay labelled marks on the image so the model answers by reading the marks rather than predicting continuous coordinates. - Lei, Yang, et al., "Scaffolding Coordinates to Promote Vision-Language Coordination in Large Multi-Modal Models," arXiv:2402.12058 (2024): overlay a labelled coordinate grid as positional anchors the model reads. Overlaying a labelled row/column grid and reading back grid IDs in place of continuous coordinates is the direct application of that work here. The model is asked to bound the WHOLE table, not just its data grid: the title/caption, units line, multi-level column headers, the row-label stub column, the body and total rows, and any footnotes attached to the table. After the text-layer tighten, each box is grown by a small fixed margin so a downstream crop does not clip an edge, and two declash passes settle the seam between overlapping neighbours: a stacked pair gets a horizontal seam through the widest blank band near the model's row boundary, a side-by-side pair a vertical seam through a blank column in the overlap. Which pass a pair gets is the model's call, read off the grid column spans it reported. A side-by-side pair whose overlap holds no blank column keeps the overlap, so a shared column can land in both boxes. When tightening finds no text (an image-only region), the coarse grid region is kept; a region is never dropped. """ from __future__ import annotations import json import string import tempfile from pathlib import Path from typing import Any, Dict, List, Optional, Protocol, Sequence, Tuple, runtime_checkable from loguru import logger from pydantic import BaseModel, Field from quber.agents.factory import supports_sampling_temperature from quber.agents.langsmith_tracer import LangSmithTracer, usage_metadata_from from quber.settings import DEFAULT_LLM_MODEL, get_settings # Defaults tuned for a US-letter page rendered at 200 DPI and then resized # by the API to <=1568px on the long edge: 36 rows ~= 22pt each and 12 # columns ~= 51pt each. The finer row pitch resolves a table's title from # its column header (often less than a coarse row apart) and keeps the # bottom of one stacked table off the top of the next, while staying legible # enough at the resized resolution that the printed row labels read without # off-by-one errors. Measured run-to-run stable across the review corpus. DEFAULT_ROWS = 36 DEFAULT_COLS = 12 # Symmetric margin (PDF points) grown onto every located box. The box is # guidance for later stages, where a clipped edge loses data but a little # slack does not; 6pt covers residual sub-row edge misses without reaching # into adjacent prose. Overlaps it creates between stacked tables are removed # by the declash pass that runs after it. BOX_PAD_PTS = 6.0 # Tightening pads (PDF points): how far outside the coarse grid region a # word's center may sit and still be pulled into the tight box. TIGHTEN_Y_PAD_PTS = 6.0 TIGHTEN_X_PAD_PTS = 10.0 def column_labels(cols: int) -> str: """Column letters A, B, C, ... for `cols` columns (cols <= 26).""" return string.ascii_uppercase[:cols] def column_index(letter: str, cols: int = DEFAULT_COLS) -> int: """Map a column letter back to its 0-based index. A blank letter or one past the last column maps to 0, the left edge. It is not clamped to the nearest valid column. """ labels = column_labels(cols) if not letter.strip(): return 0 ch = letter.strip().upper()[0] return labels.index(ch) if ch in labels else 0 class LocatedTable(BaseModel): """A table flagged by the grid locator, with a tightened region. `region` is normalized 0..1 with the page's top-left as origin (x1, y1, x2, y2) — the same frame as `DetectedTable.bbox`, so it drops straight into the correspondence/recovery geometry. `tightened` records whether the region was snapped to the text layer (False means the coarse grid cell was kept because no text was found there). """ ordinal: int = Field(ge=1, description="1-based reading-order position on the page") title: str = Field(default="", description="Visible caption or one-line column summary") region: Tuple[float, float, float, float] = Field( description="Tightened region, normalized 0..1 top-left origin (x1, y1, x2, y2)" ) grid_rows: Tuple[int, int] = Field(description="Grid row span the model reported (start, end), 1-based") grid_cols: Tuple[int, int] = Field(description="Grid column span as 0-based indices (start, end)") tightened: bool = Field( description="True if snapped to the text layer; False if coarse cell kept (no text)" ) # --- Model I/O schema (the grid-ID flags the LLM returns) -------------------- class GridFlag(BaseModel): ordinal: int = Field(ge=1, description="1-based reading order (top-to-bottom, then left-to-right)") title: str = Field(default="", description="Visible caption, else a one-line column summary") row_start: int = Field(description="Printed grid row number at the table's TOP edge") row_end: int = Field(description="Printed grid row number at the table's BOTTOM edge") col_start: str = Field(description="Printed grid column LETTER at the table's LEFT edge") col_end: str = Field(description="Printed grid column LETTER at the table's RIGHT edge") class GridFlagResult(BaseModel): tables: List[GridFlag] = Field(default_factory=list) def build_prompt(rows: int, cols: int) -> str: last = column_labels(cols)[-1] return f"""\ The page image has a grid overlaid: {rows} numbered horizontal rows (1 at the top to {rows} at the bottom, labeled on BOTH side margins) and {cols} lettered columns (A at the left to {last} at the right, labeled on the top and bottom margins), with thin gridlines. WHAT COUNTS AS A TABLE: a block where values are arranged in a GRID — the same two or more columns line up vertically across multiple rows. These are tables. The following are NOT tables, even when they contain dollar amounts or numbers, and must never be flagged: a bulleted or numbered list; a single column of figures; a heading followed by prose sentences; a paragraph. If the numbers do not line up into repeating vertical columns, it is not a table. Find every TABLE on the page and report the grid region that bounds each one COMPLETELY. Bound the WHOLE table, including all of its parts: - TITLE / CAPTION: the table's heading line(s) and any units line directly above the columns (e.g. "Condensed Consolidated Statements of Income", "(In millions)", "(UNAUDITED)"). Include them in the region. - COLUMN HEADERS: every header row, including multi-level / spanning headers (a label like "Three Months Ended" that sits above several sub-columns, with the period and year rows beneath it). Include all header levels. - ROW LABELS (the stub column): the left-hand column of row names. The region's LEFT edge must include it — do not start at the first number. - BODY: every data row, including subtotal and total rows. Never cut a table off partway; blank lines inside a table do NOT end it. - FOOTNOTES: short notes tied to the table by a marker (e.g. "(1)", "(2)") sitting just below it. Include them inside that table's region. Footnotes are NEVER a table of their own — never report them as a separate table. Counting tables — judge by what you see, not by repeated text: - Two tables stacked vertically are separate ONLY when each has its own title/caption and its own column-header block. Report them separately. - Two tables printed SIDE BY SIDE, one to the left of the other in the same rows, are separate tables. Report each with its own column span. - One table may have section sub-headings inside it (e.g. "Cash flows from operating activities:", "Current assets:") and blank lines between sections. These do NOT start a new table — keep the whole thing as ONE region. - Repeated "$" signs or a repeated subtotal style inside one table do NOT make a new table. When unsure, prefer ONE region over splitting. Do NOT merge two genuinely separate tables, and do NOT split one table. Exclude page furniture (running header timestamp, page number, document URL) and free-standing prose paragraphs / lists that are not a table. For each table, in reading order (top-to-bottom, then left-to-right), report: - ordinal (1 = first) - title (its visible caption, else a one-line summary of its columns) - row_start: printed row number at the TOP of the title line (one row ABOVE the column headers when a title/caption is present) — include the title row - row_end: printed row number at the BOTTOM-most part (last total, or footnote) - col_start: printed column LETTER at the left edge, INCLUDING the row labels - col_end: printed column LETTER at the right edge (last data column) Read the printed labels off the grid — do not estimate coordinates. Return ONLY a JSON object conforming to this schema (no prose, no code fences): {json.dumps(GridFlagResult.model_json_schema(), indent=2)} """ # --- Deterministic geometry (no model; unit-tested directly) ----------------- def grid_region_norm( row_start: int, row_end: int, col_start_idx: int, col_end_idx: int, rows: int, cols: int ) -> Tuple[float, float, float, float]: """Coarse grid cell span -> normalized 0..1 top-left box. Rows are 1-based and inclusive (a table on rows 4..6 spans the top of row 4 to the bottom of row 6). Columns are 0-based inclusive indices. """ y0 = max(0.0, (row_start - 1) / rows) y1 = min(1.0, row_end / rows) x0 = max(0.0, col_start_idx / cols) x1 = min(1.0, (col_end_idx + 1) / cols) return (x0, y0, x1, y1) def tighten_region( region_norm: Tuple[float, float, float, float], words: Sequence[Tuple[Any, ...]], page_w: float, page_h: float, shares_rows: bool, y_pad: float = TIGHTEN_Y_PAD_PTS, x_pad: float = TIGHTEN_X_PAD_PTS, ) -> Optional[Tuple[float, float, float, float]]: """Snap a coarse grid region to the page's text layer. The grid region's vertical span is reliable (rows separate stacked tables); within it, take the words and tight-box them. When the table is vertically alone (`shares_rows` False) the full horizontal extent of those words is used — recovering a label column the model clipped. When it shares its rows with a side-by-side neighbour, selection is confined to the flagged column range so neighbours do not merge. `words` are PyMuPDF word tuples (x0, y0, x1, y1, text, ...) in PDF points, top-left origin. Returns the tightened normalized box, or None when no word falls in the region (an image-only region — the caller keeps the coarse region rather than drop it). """ x0 = min(region_norm[0], region_norm[2]) * page_w x1 = max(region_norm[0], region_norm[2]) * page_w y0 = min(region_norm[1], region_norm[3]) * page_h y1 = max(region_norm[1], region_norm[3]) * page_h in_y = [w for w in words if y0 - y_pad <= (w[1] + w[3]) / 2.0 <= y1 + y_pad] if shares_rows: sel = [w for w in in_y if x0 - x_pad <= (w[0] + w[2]) / 2.0 <= x1 + x_pad] else: sel = in_y if not sel: return None tx0 = min(w[0] for w in sel) ty0 = min(w[1] for w in sel) tx1 = max(w[2] for w in sel) ty1 = max(w[3] for w in sel) return (tx0 / page_w, ty0 / page_h, tx1 / page_w, ty1 / page_h) def overlay_grid( image_path: Path, out_path: Path, rows: int = DEFAULT_ROWS, cols: int = DEFAULT_COLS ) -> None: """Draw a labeled row/column grid on a rendered page image.""" from PIL import Image, ImageDraw, ImageFont img = Image.open(image_path).convert("RGB") draw = ImageDraw.Draw(img) try: font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 18) except Exception: font = ImageFont.load_default() labels = column_labels(cols) rstep = img.height / rows cstep = img.width / cols for r in range(rows + 1): y = r * rstep draw.line([(0, y), (img.width, y)], fill=(80, 120, 255), width=1) if r < rows: draw.text((2, y + 2), str(r + 1), fill=(40, 80, 230), font=font) draw.text((img.width - 26, y + 2), str(r + 1), fill=(40, 80, 230), font=font) for ci in range(cols + 1): x = ci * cstep draw.line([(x, 0), (x, img.height)], fill=(255, 140, 0), width=1) if ci < cols: draw.text((x + 3, 2), labels[ci], fill=(210, 110, 0), font=font) draw.text((x + 3, img.height - 20), labels[ci], fill=(210, 110, 0), font=font) img.save(out_path) def flags_to_located( flags: Sequence[GridFlag], words: Sequence[Tuple[Any, ...]], page_w: float, page_h: float, rows: int, cols: int, ) -> List[LocatedTable]: """Turn the model's grid-ID flags into tightened LocatedTables. `shares_rows` for each flag is computed against the other flags: a flag whose grid-row span overlaps another's is treated as side-by-side and tightened within its columns; otherwise it snaps to the full word extent in its rows. """ located: List[LocatedTable] = [] for f in flags: shares = any(o is not f and not (f.row_end < o.row_start or f.row_start > o.row_end) for o in flags) c0 = column_index(f.col_start, cols) c1 = column_index(f.col_end, cols) if c1 < c0: c0, c1 = c1, c0 coarse = grid_region_norm(f.row_start, f.row_end, c0, c1, rows, cols) tight = tighten_region(coarse, words, page_w, page_h, shares) located.append( LocatedTable( ordinal=f.ordinal, title=f.title, region=tight if tight is not None else coarse, grid_rows=(f.row_start, f.row_end), grid_cols=(c0, c1), tightened=tight is not None, ) ) return located def pad_boxes( located: Sequence[LocatedTable], page_w: float, page_h: float, pad_pts: float ) -> List[LocatedTable]: """Grow every region by a fixed margin on all sides, clamped to the page. A guidance box should over-include slightly rather than clip, so a plain symmetric pad covers residual sub-row edge misses without any content-dependent heuristic. Any overlap this creates between stacked tables is resolved by `declash_stacked`, which runs after it. """ px, py = pad_pts / page_w, pad_pts / page_h out: List[LocatedTable] = [] for t in located: x0, y0, x1, y1 = t.region region = ( max(0.0, min(x0, x1) - px), max(0.0, min(y0, y1) - py), min(1.0, max(x0, x1) + px), min(1.0, max(y0, y1) + py), ) out.append(t.model_copy(update={"region": region})) return out # How far, in grid rows, above and below the model's row-span boundary to look # for the blank band that separates two stacked tables. The model gives a coarse # boundary; the true gap sits within about a row of it, so a narrow window keeps # the search off section breaks deeper inside either table. SEAM_WINDOW_ROWS = 1.5 def cluster_text_rows( words: Sequence[Tuple[Any, ...]], x0: float, x1: float, y_lo: float, y_hi: float, tol: float = 3.0, ) -> List[Dict[str, Any]]: """Group page words into printed text rows within a window. Words whose centers fall outside the window [x0, x1] x [y_lo, y_hi] are discarded. The rest, taken top to bottom by vertical center, join the current row when their center is within `tol` points of the last word added to it. Otherwise they start a new row. Each returned row carries its top edge, bottom edge and left-to-right joined text. Rows come back ordered top to bottom. `words` are PyMuPDF word tuples (x0, y0, x1, y1, text, ...) in PDF points, top-left origin. """ sel = [w for w in words if y_lo <= (w[1] + w[3]) / 2.0 <= y_hi and x0 <= (w[0] + w[2]) / 2.0 <= x1] sel.sort(key=lambda w: (w[1] + w[3]) / 2.0) rows: List[Dict[str, Any]] = [] for w in sel: center = (w[1] + w[3]) / 2.0 if rows and center - rows[-1]["center"] <= tol: row = rows[-1] row["top"] = min(row["top"], w[1]) row["bottom"] = max(row["bottom"], w[3]) row["center"] = center row["words"].append(w) else: rows.append({"center": center, "top": w[1], "bottom": w[3], "words": [w]}) for row in rows: row["text"] = " ".join(w[4] for w in sorted(row["words"], key=lambda w: w[0])) return rows def widest_gap_seam( words: Sequence[Tuple[Any, ...]], x0: float, x1: float, y_lo: float, y_hi: float, ) -> Optional[float]: """Midpoint of the widest blank band between consecutive text rows in the window [y_lo, y_hi] (PDF points), or None when fewer than two rows fall in it. Two stacked tables are separated by a blank band wider than the spacing between rows inside either table, so the widest gap in a window centred on their boundary is the seam — independent of whether the rows on each side are totals, titles, headers or plain data. """ rows = cluster_text_rows(words, x0, x1, y_lo, y_hi) if len(rows) < 2: return None best_gap, seam = -1.0, None for upper, lower in zip(rows, rows[1:], strict=False): gap = lower["top"] - upper["bottom"] if gap > best_gap: best_gap, seam = gap, (upper["bottom"] + lower["top"]) / 2.0 return seam def fallback_seam( upper_region: Tuple[float, float, float, float], lower_top: float, words: Sequence[Tuple[Any, ...]], page_w: float, page_h: float, ) -> float: """Pull the upper box's bottom up to its own last line of text above the lower box's top. `declash_stacked` uses this when `widest_gap_seam` finds no band near the model's row boundary, or finds one outside the two boxes' combined vertical span. Returns the new upper-box bottom in PDF points (top-origin); never reaches below the lower box's top, so the boxes still tile. """ u_top = min(upper_region[1], upper_region[3]) * page_h ux0 = min(upper_region[0], upper_region[2]) * page_w ux1 = max(upper_region[0], upper_region[2]) * page_w inside = [ w for w in words if u_top < (w[1] + w[3]) / 2.0 < lower_top - 0.5 and ux0 - 2.0 <= (w[0] + w[2]) / 2.0 <= ux1 + 2.0 ] return min(max((w[3] for w in inside), default=lower_top), lower_top) # A pair of boxes is treated as stacked only when they share at least this # fraction of the narrower span's grid columns. The locator's answer is the # authority on what is a table and where it sits; whether two of its tables # are stacked is read off that same answer. Genuinely stacked tables occupy # the same column band, so the narrower span nests inside the wider one; # side-by-side tables share at most an edge column of quantization slop. # Forcing a horizontal seam through a side-by-side pair amputates the bottom # rows of one table and the top rows of the other. STACKED_MIN_COL_OVERLAP = 0.5 def col_span_overlap_fraction(a: Tuple[int, int], b: Tuple[int, int]) -> float: """Fraction of the narrower grid-column span shared by the two spans. Spans are the locator's 0-based inclusive (start, end) column indices. """ a0, a1 = min(a), max(a) b0, b1 = min(b), max(b) shared = min(a1, b1) - max(a0, b0) + 1 narrower = min(a1 - a0, b1 - b0) + 1 return shared / narrower if shared > 0 else 0.0 def declash_stacked( located: Sequence[LocatedTable], words: Sequence[Tuple[Any, ...]], page_w: float, page_h: float, min_gap_pts: float = 2.0, rows: int = DEFAULT_ROWS, ) -> List[LocatedTable]: """Set the seam between two stacked boxes to the blank band between the two tables, so neither box trades a row with its neighbour. When an upper box's bottom runs into the next box's top — grid quantization, padding, or an upward tighten can pull one box over a row belonging to the other — take the model's row-span boundary between the two tables as a coarse anchor, find the widest word-free band within about a grid row of it, and snap both the upper box's bottom and the lower box's top to its midpoint. Each table keeps every one of its own rows, whatever they are. When no gap is found in the window, fall back to pulling the upper box's bottom up to its own last line of text. Only cleans the seam; it never merges two boxes or splits one. Only pairs the locator placed in the same column band get a seam: whether two tables are stacked is the locator's call, read off the grid column spans it reported. A pair whose spans share less than `STACKED_MIN_COL_OVERLAP` of the narrower span is side-by-side, and a seam through it would cut real rows from both tables. Such pairs pass through untouched. `words` are PyMuPDF word tuples (x0, y0, x1, y1, text, ...) in PDF points, top-left origin — the same frame `tighten_region` uses. """ order = sorted(range(len(located)), key=lambda i: min(located[i].region[1], located[i].region[3])) regions = {i: located[i].region for i in range(len(located))} row_h = page_h / rows for upper, lower in zip(order, order[1:], strict=False): ru, rl = regions[upper], regions[lower] cols_u, cols_l = located[upper].grid_cols, located[lower].grid_cols if col_span_overlap_fraction(cols_u, cols_l) < STACKED_MIN_COL_OVERLAP: continue # the locator put these side by side, not stacked: no seam to clean u_top_norm = min(ru[1], ru[3]) u_bottom = max(ru[1], ru[3]) * page_h l_top = min(rl[1], rl[3]) * page_h l_bottom_norm = max(rl[1], rl[3]) if u_bottom <= l_top - min_gap_pts: continue # already separated by a clear gap # Coarse boundary from the model's row spans, then the precise blank band near it. boundary = ( (located[upper].grid_rows[1] / rows + (located[lower].grid_rows[0] - 1) / rows) / 2.0 * page_h ) wx0 = min(ru[0], ru[2], rl[0], rl[2]) * page_w - 2.0 wx1 = max(ru[0], ru[2], rl[0], rl[2]) * page_w + 2.0 seam = widest_gap_seam( words, wx0, wx1, boundary - SEAM_WINDOW_ROWS * row_h, boundary + SEAM_WINDOW_ROWS * row_h ) if seam is not None and u_top_norm * page_h < seam < l_bottom_norm * page_h: seam_norm = seam / page_h regions[upper] = (ru[0], u_top_norm, ru[2], seam_norm) regions[lower] = (rl[0], seam_norm, rl[2], l_bottom_norm) else: new_bottom = fallback_seam(ru, l_top, words, page_w, page_h) regions[upper] = (ru[0], u_top_norm, ru[2], new_bottom / page_h) return [located[i].model_copy(update={"region": regions[i]}) for i in range(len(located))] # A blank column may be crossed by lines that belong to neither table: the # running header centred over the page, a units line, the prose sentence # introducing both tables. When the boxes reach up to the page header, those # lines can be a third of the rows in the window. Up to this fraction of the # rows may cross a band and it still counts as blank; a table's own column is # crossed by every row, so it never qualifies. Which blank band is the seam is # then decided by what follows it, not by how few rows cross it. SEAM_MAX_CROSSING_ROWS = 0.5 # Slack (PDF points) added on either side of two boxes' overlap when searching # it for the blank column between them, so a seam whose edge coincides with the # overlap's edge is still found whole. SEAM_OVERLAP_PAD_PTS = 3.0 def blank_column_bands( words: Sequence[Tuple[Any, ...]], x0: float, x1: float, y_lo: float, y_hi: float, min_width_pts: float = 4.0, ) -> List[Tuple[float, float]]: """Every ink-free vertical band at least `min_width_pts` wide in the window [x0, x1] x [y_lo, y_hi] (PDF points), as (left edge, right edge) pairs in left-to-right order; empty when the window holds fewer than two text rows. A band is measured per printed text row: each row's words are projected onto x, and a band counts as blank when at most `SEAM_MAX_CROSSING_ROWS` of the rows have ink in it. An introductory sentence printed across the whole page crosses the seam between two tables, and every row of both tables respects it; a gap between two columns inside one table is blank in the same way, which is why the caller, not this function, decides which band is the seam. """ rows = cluster_text_rows(words, x0, x1, y_lo, y_hi) if len(rows) < 2: return [] step = 0.5 n = int((x1 - x0) / step) + 1 coverage = [0] * n for row in rows: for w in row["words"]: a = max(0, int((w[0] - x0) / step)) b = min(n - 1, int((w[2] - x0) / step)) for i in range(a, b + 1): coverage[i] += 1 allowed = SEAM_MAX_CROSSING_ROWS * len(rows) bands: List[Tuple[float, float]] = [] i = 0 while i < n: if coverage[i] > allowed: i += 1 continue j = i while j + 1 < n and coverage[j + 1] <= allowed: j += 1 if (j - i + 1) * step >= min_width_pts: bands.append((x0 + i * step, x0 + (j + 1) * step)) i = j + 1 return bands # How far right of a candidate band to look for the column that follows it, # and the share of that column's words that must start with a letter for the # column to read as a table's row-label column. STUB_LOOKAHEAD_PTS = 90.0 STUB_MIN_ALPHA_FRACTION = 0.6 def stub_column_follows( words: Sequence[Tuple[Any, ...]], band_right: float, y_lo: float, y_hi: float ) -> bool: """Whether the column of text just right of `band_right` reads as a row-label column: on most rows its first word starts with a letter. A financial table's first column names its rows; its other columns hold figures. So of the blank bands inside two side-by-side boxes' overlap, the one followed by a column of words is the seam between the tables, and one followed by a column of figures is a gap inside the left table. """ rows = cluster_text_rows(words, band_right, band_right + STUB_LOOKAHEAD_PTS, y_lo, y_hi) firsts = [sorted(r["words"], key=lambda w: w[0])[0][4] for r in rows if r["words"]] if len(firsts) < 2: return False alpha = sum(1 for t in firsts if t[:1].isalpha()) return alpha / len(firsts) >= STUB_MIN_ALPHA_FRACTION def declash_side_by_side( located: Sequence[LocatedTable], words: Sequence[Tuple[Any, ...]], page_w: float, page_h: float, min_gap_pts: float = 2.0, ) -> List[LocatedTable]: """Set the seam between two side-by-side boxes to the blank column between the two tables, so neither box trades a column with its neighbour. The counterpart of `declash_stacked` for pairs the locator placed in different column bands of the same rows. The grid columns are coarse, so a left table's span and a right table's span usually share an edge column; the tighten then admits words from the neighbour's stub column into both boxes, and Camelot reads that column into both tables. Search the two boxes' overlap for blank columns and snap the left box's right edge and the right box's left edge to the midpoint of one. A blank column that a row-label column follows is preferred. When none qualifies, the blank column nearest the model's coarse column boundary is used, and it can be a gap inside the left table. When the overlap holds no blank column at all, the pair passes through untouched: a wrong seam would cut a real column, and the overlap is at worst a duplicated one. Only pairs whose regions overlap vertically and whose column spans share less than `STACKED_MIN_COL_OVERLAP` of the narrower span are considered; stacked pairs are `declash_stacked`'s business. """ regions = {i: located[i].region for i in range(len(located))} order = sorted(range(len(located)), key=lambda i: min(located[i].region[0], located[i].region[2])) for a, b in zip(order, order[1:], strict=False): left, right = regions[a], regions[b] if col_span_overlap_fraction(located[a].grid_cols, located[b].grid_cols) >= STACKED_MIN_COL_OVERLAP: continue # stacked, or one span nested in the other: not a side-by-side pair l_x1 = max(left[0], left[2]) * page_w r_x0 = min(right[0], right[2]) * page_w if l_x1 <= r_x0 - min_gap_pts: continue # already separated by a clear gap y_lo = max(min(left[1], left[3]), min(right[1], right[3])) * page_h y_hi = min(max(left[1], left[3]), max(right[1], right[3])) * page_h if y_hi <= y_lo: continue # no shared rows: nothing to trade # The overlap is the tighten's doing: each box reached to the words it # admitted from the neighbour, so it holds the left table's last column # or two, the blank column between the tables, and the right table's # first column or two. Of the blank bands in it, the seam is the one a # row-label column follows; with none such, the one nearest the model's # coarse column boundary. bands = blank_column_bands( words, r_x0 - SEAM_OVERLAP_PAD_PTS, l_x1 + SEAM_OVERLAP_PAD_PTS, y_lo, y_hi ) l_x0 = min(left[0], left[2]) * page_w r_x1 = max(right[0], right[2]) * page_w bands = [(b0, b1) for b0, b1 in bands if l_x0 < (b0 + b1) / 2.0 < r_x1] if not bands: logger.debug( "declash_side_by_side: no blank column in the overlap x={:.0f}-{:.0f}pt; pair left as is", r_x0, l_x1, ) continue boundary = ( ((located[a].grid_cols[1] + 1) / DEFAULT_COLS + located[b].grid_cols[0] / DEFAULT_COLS) / 2.0 * page_w ) stub = [band for band in bands if stub_column_follows(words, band[1], y_lo, y_hi)] band = min(stub or bands, key=lambda bd: abs((bd[0] + bd[1]) / 2.0 - boundary)) seam = (band[0] + band[1]) / 2.0 seam_norm = seam / page_w regions[a] = (min(left[0], left[2]), min(left[1], left[3]), seam_norm, max(left[1], left[3])) regions[b] = (seam_norm, min(right[1], right[3]), max(right[0], right[2]), max(right[1], right[3])) return [located[i].model_copy(update={"region": regions[i]}) for i in range(len(located))] # Transient LLM failures worth retrying before giving up: rate limiting, server # overload, and connection/timeout blips. Matched by HTTP status when the client # exposes one, otherwise by exception class name, so detection holds whether the # error surfaces as an Anthropic SDK error or a pydantic-ai wrapper around it. TRANSIENT_LLM_STATUS = frozenset({408, 409, 429, 500, 502, 503, 504, 529}) GRID_LOCATE_MAX_ATTEMPTS = 4 GRID_LOCATE_BACKOFF_BASE_S = 1.0 def is_transient_llm_error(exc: BaseException) -> bool: """True for failures that a short backoff is likely to clear (rate limit, overload, timeout) — as opposed to a malformed request or auth error, which will fail identically on retry.""" status = getattr(exc, "status_code", None) if isinstance(status, int) and status in TRANSIENT_LLM_STATUS: return True name = type(exc).__name__.lower() return any( token in name for token in ( "ratelimit", "overload", "timeout", "apiconnection", "internalserver", "serviceunavailable", ) ) @runtime_checkable class GridLocator(Protocol): async def locate(self, page_image: Path, source: Path, page: int) -> List[LocatedTable]: ... class PydanticAIGridLocator: """GridLocator backed by pydantic-ai with an Anthropic vision model. Mirrors `PydanticAITableDetector`: same auth resolution (`ANTHROPIC_AUTH_TOKEN` preferred over `ANTHROPIC_API_KEY`), greedy decoding for run-to-run stability of the discrete grid IDs. """ DEFAULT_TEMPERATURE = 0.0 def __init__( self, model: Optional[str] = None, auth_token: Optional[str] = None, api_key: Optional[str] = None, temperature: Optional[float] = DEFAULT_TEMPERATURE, rows: int = DEFAULT_ROWS, cols: int = DEFAULT_COLS, dpi: int = 200, ) -> None: from pydantic_ai import Agent from pydantic_ai.models.anthropic import AnthropicModel from pydantic_ai.providers.anthropic import AnthropicProvider from pydantic_ai.settings import ModelSettings llm_settings = get_settings().llm model = model or llm_settings.grid_locator_model or llm_settings.model or DEFAULT_LLM_MODEL if temperature is not None and not supports_sampling_temperature(model): temperature = None resolved_auth = auth_token or llm_settings.anthropic_auth_token resolved_key = api_key or llm_settings.anthropic_api_key if resolved_auth: from quber.agents._oauth_gate import make_oauth_anthropic_model anth_model = make_oauth_anthropic_model(model, resolved_auth) elif resolved_key: provider = AnthropicProvider(api_key=resolved_key) anth_model = AnthropicModel(model, provider=provider) else: raise RuntimeError( "PydanticAIGridLocator: neither ANTHROPIC_AUTH_TOKEN nor " "ANTHROPIC_API_KEY is set. Provide one via env or constructor." ) self.model = model self.rows = rows self.cols = cols self.dpi = dpi model_settings = None if temperature is None else ModelSettings(temperature=temperature) self.system_prompt = build_prompt(rows, cols) self.agent = Agent( anth_model, output_type=GridFlagResult, system_prompt=self.system_prompt, model_settings=model_settings, ) # LangSmith tracer is a no-op if `TRACE_TO_LANGSMITH` is unset. self.tracer = LangSmithTracer(run_name="quber-grid-locator") def trace_inputs(self) -> Dict[str, Any]: return { "messages": [ {"role": "system", "content": self.system_prompt}, { "role": "user", "content": [ {"type": "text", "text": "Flag the tables using the grid."}, {"type": "image", "text": "[gridded page image attached]"}, ], }, ] } async def locate(self, page_image: Path, source: Path, page: int) -> List[LocatedTable]: import asyncio from pydantic_ai import BinaryContent from quber.agents.completeness import page_words with tempfile.TemporaryDirectory(prefix="quber-grid-") as tmp: gridded = Path(tmp) / "gridded.png" # Drawing the grid is CPU-bound PIL work; run it off the event loop # so it cannot serialize the dispatch of concurrent locate calls. await asyncio.to_thread(overlay_grid, Path(page_image), gridded, self.rows, self.cols) image = BinaryContent(data=gridded.read_bytes(), media_type="image/png") inputs = self.trace_inputs() async with self.tracer.llm_run("grid_locate", inputs, model=self.model) as run: result = await self.run_locator_with_retry(image, page) run.outputs = { "messages": [{"role": "assistant", "content": result.output.model_dump_json()}], "table_count": len(result.output.tables), "usage_metadata": usage_metadata_from(result.usage), } page_w, page_h, words = await asyncio.to_thread(page_words, Path(source), page) located = flags_to_located(result.output.tables, words, page_w, page_h, self.rows, self.cols) located = pad_boxes(located, page_w, page_h, BOX_PAD_PTS) located = declash_stacked(located, words, page_w, page_h, rows=self.rows) located = declash_side_by_side(located, words, page_w, page_h) return located async def run_locator_with_retry(self, image: Any, page: int) -> Any: """Run the grid-locate call, retrying transient LLM failures with exponential backoff. A page that still fails after the retry budget is raised, never swallowed: returning no flags would silently drop every table on the page with no record in the output -- on a financial document an undetected page is a liability, not an acceptable omission. TODO: offer a graceful-degradation mode that records a failed-page marker and continues with a non-zero exit, instead of aborting the whole run. Failing hard for now. """ import asyncio for attempt in range(1, GRID_LOCATE_MAX_ATTEMPTS + 1): try: return await self.agent.run(["Flag the tables using the grid.", image]) except Exception as exc: if attempt == GRID_LOCATE_MAX_ATTEMPTS or not is_transient_llm_error(exc): logger.error( "grid_locate: LLM call failed page={} attempt={}/{}; aborting " "(refusing to silently drop the page): {}", page, attempt, GRID_LOCATE_MAX_ATTEMPTS, exc, ) raise delay = GRID_LOCATE_BACKOFF_BASE_S * 2 ** (attempt - 1) logger.warning( "grid_locate: transient LLM error page={} attempt={}/{}; retrying in {}s: {}", page, attempt, GRID_LOCATE_MAX_ATTEMPTS, delay, exc, ) await asyncio.sleep(delay) class MockGridLocator: """Returns canned LocatedTables. For tests.""" def __init__(self, located: Optional[List[LocatedTable]] = None) -> None: self.located = located or [] async def locate(self, page_image: Path, source: Path, page: int) -> List[LocatedTable]: _ = (page_image, source, page) return list(self.located) def get_grid_locator(backend: Optional[str] = None) -> GridLocator: selected = backend or get_settings().llm.grid_locator_backend if selected == "api": return PydanticAIGridLocator() if selected == "mock": return MockGridLocator() raise ValueError(f"Unknown QUBER_GRID_LOCATOR_BACKEND: {selected!r}. Expected api|mock.")