"""Deterministic-core tests for the GridLocator (no LLM). Covers the grid-cell -> normalized-box mapping, the text-layer tightening (alone-in-rows recovers the full word extent; side-by-side stays confined to its columns; no text -> None so the caller keeps the coarse region), and the shares-rows detection that drives which branch a flag takes. """ import asyncio from typing import Any import pytest from quber.agents.grid_locator import ( DEFAULT_COLS, DEFAULT_ROWS, GRID_LOCATE_MAX_ATTEMPTS, STACKED_MIN_COL_OVERLAP, GridFlag, LocatedTable, PydanticAIGridLocator, blank_column_bands, build_prompt, col_span_overlap_fraction, column_index, column_labels, declash_side_by_side, declash_stacked, flags_to_located, grid_region_norm, is_transient_llm_error, pad_boxes, stub_column_follows, tighten_region, ) PAGE_W = 612.0 PAGE_H = 792.0 def word(x0: float, y0: float, x1: float, y1: float, text: str = "x") -> tuple[Any, ...]: """A minimal PyMuPDF-style word tuple.""" return (x0, y0, x1, y1, text, 0, 0, 0) def test_column_labels_and_index(): assert column_labels(12) == "ABCDEFGHIJKL" assert column_index("A", 12) == 0 assert column_index("L", 12) == 11 # blank or out-of-range letters clamp to 0, never raise assert column_index("", 12) == 0 assert column_index("Z", 12) == 0 def test_grid_region_norm_maps_rows_and_cols(): # rows 4..6 of 24, full width cols 0..11 of 12 box = grid_region_norm(4, 6, 0, 11, 24, 12) assert box[0] == 0.0 # x0 assert box[2] == 1.0 # x1 (col 11 inclusive -> 12/12) assert abs(box[1] - 3 / 24) < 1e-9 # y0 = top of row 4 assert abs(box[3] - 6 / 24) < 1e-9 # y1 = bottom of row 6 def test_grid_region_norm_clamps_to_unit_square(): box = grid_region_norm(1, 24, 0, 11, 24, 12) assert box == (0.0, 0.0, 1.0, 1.0) def test_tighten_alone_recovers_clipped_label_column(): # Model flagged cols E-I (idx 4-8) but the row-label column sits at x~46, # left of the flag. Alone in its rows -> snap to the full word extent. region = grid_region_norm(13, 16, 4, 8, DEFAULT_ROWS, DEFAULT_COLS) words = [ word(46, 280, 90, 290, "Region"), # left label, outside flagged cols word(210, 280, 260, 290, "1,234"), word(360, 282, 400, 292, "5.1%"), word(46, 320, 90, 330, "US"), ] tight = tighten_region(region, words, PAGE_W, PAGE_H, shares_rows=False) assert tight is not None # left edge pulled back to the label column (~46pt), not the flag (~204pt) assert abs(tight[0] * PAGE_W - 46) < 1.0 assert abs(tight[2] * PAGE_W - 400) < 1.0 def test_tighten_side_by_side_stays_within_columns(): # Two tables share rows 5-8; tightening the LEFT one must not grab the # RIGHT one's words. region = grid_region_norm(5, 8, 1, 3, DEFAULT_ROWS, DEFAULT_COLS) # cols B-D words = [ word(46, 120, 187, 130, "left"), # left table word(120, 150, 180, 160, "456"), word(306, 120, 447, 130, "right"), # right table, same rows word(380, 150, 447, 160, "789"), ] tight = tighten_region(region, words, PAGE_W, PAGE_H, shares_rows=True) assert tight is not None # right edge stays left of the right-hand table (which starts at x=306) assert tight[2] * PAGE_W < 250 def test_tighten_no_text_returns_none(): # An image-only region: no words inside -> None so the caller keeps the # coarse grid region (never drops a flagged table). region = grid_region_norm(2, 4, 0, 11, DEFAULT_ROWS, DEFAULT_COLS) words = [word(46, 600, 90, 610, "far below")] assert tighten_region(region, words, PAGE_W, PAGE_H, shares_rows=False) is None def test_flags_to_located_stacked_are_not_shared(): # Two vertically stacked tables (disjoint rows) -> each tightens to full # word extent in its own rows. flags = [ GridFlag(ordinal=1, title="t1", row_start=4, row_end=6, col_start="A", col_end="E"), GridFlag(ordinal=2, title="t2", row_start=13, row_end=15, col_start="A", col_end="E"), ] words = [ word(46, 100, 300, 110, "row1"), # in table 1's rows (y ~66-132) word(46, 290, 300, 300, "row2"), # in table 2's rows (y ~264-330) ] located = flags_to_located(flags, words, PAGE_W, PAGE_H, DEFAULT_ROWS, DEFAULT_COLS) assert len(located) == 2 assert all(t.tightened for t in located) # ordinals preserved, regions separated vertically assert located[0].region[3] < located[1].region[1] def test_flags_to_located_side_by_side_share_rows(): flags = [ GridFlag(ordinal=1, title="left", row_start=5, row_end=8, col_start="B", col_end="D"), GridFlag(ordinal=2, title="right", row_start=5, row_end=8, col_start="G", col_end="I"), ] words = [ word(46, 100, 187, 160, "left"), word(306, 100, 447, 160, "right"), ] located = flags_to_located(flags, words, PAGE_W, PAGE_H, DEFAULT_ROWS, DEFAULT_COLS) assert len(located) == 2 # separated horizontally, neither swallowing the other assert located[0].region[2] < located[1].region[0] def located( region: tuple[float, float, float, float], ordinal: int = 1, grid_cols: tuple[int, int] = (0, 1), ) -> LocatedTable: return LocatedTable( ordinal=ordinal, title="t", region=region, grid_rows=(1, 1), grid_cols=grid_cols, tightened=True ) def test_pad_boxes_grows_and_clamps(): grown = pad_boxes([located((0.1, 0.1, 0.9, 0.9))], PAGE_W, PAGE_H, 6.0)[0].region assert abs(grown[0] - (0.1 - 6 / PAGE_W)) < 1e-9 assert abs(grown[1] - (0.1 - 6 / PAGE_H)) < 1e-9 assert abs(grown[2] - (0.9 + 6 / PAGE_W)) < 1e-9 # a box hard against the top-left edge clamps to 0, never negative edge = pad_boxes([located((0.001, 0.001, 0.5, 0.5))], PAGE_W, PAGE_H, 6.0)[0].region assert edge[0] == 0.0 and edge[1] == 0.0 def test_declash_pulls_upper_bottom_off_lower_top(): # upper box bottom (y=0.40) overruns lower box top (y=0.34); the upper # table's last word ends at y=250 -> upper bottom snaps up to 250. upper = located((0.05, 0.05, 0.95, 0.40), ordinal=1) lower = located((0.05, 0.34, 0.95, 0.60), ordinal=2) words = [word(50, 240, 300, 250, "last-upper-row")] # center y=245, above lower top out = declash_stacked([upper, lower], words, PAGE_W, PAGE_H) assert abs(out[0].region[3] * PAGE_H - 250) < 1.0 # no overlap left between the two boxes assert max(out[0].region[1], out[0].region[3]) <= min(out[1].region[1], out[1].region[3]) def test_declash_seam_lands_in_blank_band_between_tables(): # Two stacked tables whose boxes overlap (upper bottom 0.40 > lower top 0.32). # The model's row boundary is rows 12|13 (y~264pt). The upper table's last row # sits at y~255-264 and the lower table's first row at y~275-284, with a blank # band between. The seam must land in that band, keeping the upper last row # (center ~259) with the upper box — with no title and no "total" assumption. upper = LocatedTable( ordinal=1, title="", region=(0.05, 0.05, 0.95, 0.40), grid_rows=(3, 12), grid_cols=(0, 11), tightened=True, ) lower = LocatedTable( ordinal=2, title="", region=(0.05, 0.32, 0.95, 0.60), grid_rows=(13, 23), grid_cols=(0, 11), tightened=True, ) words = [ word(50, 255, 300, 264, "upper-last-row"), word(50, 275, 400, 284, "lower-first-row"), ] out = declash_stacked([upper, lower], words, PAGE_W, PAGE_H) seam = out[0].region[3] * PAGE_H assert abs(seam - (264 + 275) / 2) < 1.0 # midpoint of the blank band assert out[1].region[1] * PAGE_H == seam # boxes meet at the seam assert seam > 259.5 # upper last row (center ~259) stays with the upper box def test_declash_leaves_side_by_side_pair_untouched(): # Two tables the locator flagged in side-by-side column bands overlap # vertically by construction, but they are not stacked: a seam through # them would cut the left table's bottom rows and the right table's top # rows. Geometry and grid spans taken from a real 10-Q page where a # geography table (17 rows, cols B-F) sits left of a property type table # (9 rows, cols G-K). left = located((0.0711, 0.0903, 0.5567, 0.4422), ordinal=1, grid_cols=(1, 5)) right = located((0.4609, 0.0903, 0.9277, 0.3135), ordinal=2, grid_cols=(6, 10)) words = [ word(50, 200, 300, 210, "left-row"), word(320, 200, 500, 210, "right-row"), ] out = declash_stacked([left, right], words, PAGE_W, PAGE_H) assert out[0].region == left.region assert out[1].region == right.region def test_col_span_overlap_classifies_stacked_vs_side_by_side(): # A narrow table under a full-width one nests inside its column band. assert col_span_overlap_fraction((0, 11), (0, 4)) == 1.0 # Disjoint spans (the side-by-side case) share nothing. assert col_span_overlap_fraction((1, 5), (6, 10)) == 0.0 # One edge column of quantization slop stays below the threshold. assert col_span_overlap_fraction((1, 5), (5, 10)) < STACKED_MIN_COL_OVERLAP def test_declash_no_op_when_disjoint(): upper = located((0.05, 0.05, 0.95, 0.30), ordinal=1) lower = located((0.05, 0.50, 0.95, 0.70), ordinal=2) out = declash_stacked([upper, lower], [word(50, 200, 300, 210, "x")], PAGE_W, PAGE_H) assert out[0].region == upper.region assert out[1].region == lower.region class _StatusError(Exception): """Stand-in for an LLM client error carrying an HTTP status_code.""" def __init__(self, status_code: int) -> None: super().__init__(f"status {status_code}") self.status_code = status_code class _RateLimitError(Exception): """Stand-in matched by class name rather than status_code.""" class _FakeAgent: """Minimal stand-in for the pydantic-ai Agent: fails the first `fail_times` calls with `exc`, then returns a sentinel.""" def __init__(self, exc: Exception, fail_times: int, sentinel: object = "OK") -> None: self.exc = exc self.fail_times = fail_times self.sentinel = sentinel self.calls = 0 async def run(self, _inputs: object) -> object: self.calls += 1 if self.calls <= self.fail_times: raise self.exc return self.sentinel def _locator_with_agent(agent: _FakeAgent) -> PydanticAIGridLocator: # Bypass __init__ (which needs real Anthropic credentials); only `agent` is # exercised by run_locator_with_retry. loc = object.__new__(PydanticAIGridLocator) loc.agent = agent # type: ignore[attr-defined] return loc def test_is_transient_llm_error_by_status(): assert is_transient_llm_error(_StatusError(429)) assert is_transient_llm_error(_StatusError(529)) assert not is_transient_llm_error(_StatusError(400)) assert not is_transient_llm_error(_StatusError(401)) def test_is_transient_llm_error_by_class_name(): assert is_transient_llm_error(_RateLimitError()) assert not is_transient_llm_error(ValueError("malformed request")) def test_run_locator_retries_transient_then_succeeds(monkeypatch: pytest.MonkeyPatch): async def _no_sleep(_seconds: float) -> None: return None monkeypatch.setattr(asyncio, "sleep", _no_sleep) agent = _FakeAgent(_StatusError(429), fail_times=2, sentinel="located") loc = _locator_with_agent(agent) out = asyncio.run(loc.run_locator_with_retry(image=None, page=5)) assert out == "located" assert agent.calls == 3 # two failures retried, third succeeds def test_run_locator_raises_after_budget_never_silent(monkeypatch: pytest.MonkeyPatch): async def _no_sleep(_seconds: float) -> None: return None monkeypatch.setattr(asyncio, "sleep", _no_sleep) agent = _FakeAgent(_StatusError(429), fail_times=99) loc = _locator_with_agent(agent) # A page that never recovers must raise -- not return [] and silently drop # every table on the page. with pytest.raises(_StatusError): asyncio.run(loc.run_locator_with_retry(image=None, page=7)) assert agent.calls == GRID_LOCATE_MAX_ATTEMPTS def test_run_locator_non_transient_raises_immediately(): agent = _FakeAgent(ValueError("malformed request"), fail_times=99) loc = _locator_with_agent(agent) with pytest.raises(ValueError): asyncio.run(loc.run_locator_with_retry(image=None, page=1)) assert agent.calls == 1 # no retry on a non-transient error def test_prompt_counts_side_by_side_tables_separately(): # Two tables printed beside each other share a row band; without a rule # of their own the counting guidance only speaks of stacked tables and # ends by preferring one region, which merges the pair into one span. prompt = build_prompt(DEFAULT_ROWS, DEFAULT_COLS) stacked = prompt.index("Two tables stacked vertically") side = prompt.index("printed SIDE BY SIDE") assert stacked < side < prompt.index("When unsure, prefer ONE region") assert "one to the left of the other in the same\n rows, are separate tables" in prompt def _side_by_side_words() -> list[tuple[Any, ...]]: # Two tables beside each other in rows y=130..220: a left table whose # values end at x=295 and a right table whose stub starts at x=314, with # an introductory sentence at y=78 printed across the seam. words = [word(230, 78, 300, 88, "collateral"), word(302, 78, 380, 88, "outstanding")] for k, y in enumerate(range(130, 221, 13)): words += [ word(60, y, 120, y + 9, f"Left{k}"), word(210, y, 240, y + 9, "1.0"), word(271, y, 295, y + 9, "2.0"), word(314, y, 360, y + 9, f"Right{k}"), word(470, y, 500, y + 9, "3.0"), ] return words def test_blank_column_bands_ignore_a_prose_row_crossing_them(): words = _side_by_side_words() # The window spans the left table's last two columns, the blank column # between the tables, and the right table's stub. Two bands come back: the # gap between the left table's own value columns and the seam. The prose # row at y=78 crosses the seam and does not count against it. bands = blank_column_bands(words, 230.0, 344.0, 70.0, 230.0) assert len(bands) == 2 assert abs(bands[0][1] - 271.0) < 1.0 # ends where the "2.0" column starts assert abs(bands[1][0] - 295.0) < 1.0 and abs(bands[1][1] - 314.0) < 1.0 def test_stub_column_follows_reads_labels_not_figures(): words = _side_by_side_words() assert stub_column_follows(words, 314.0, 120.0, 230.0) # "Right0", "Right1", ... assert not stub_column_follows(words, 271.0, 120.0, 230.0) # "2.0" figures def test_blank_column_bands_none_inside_one_table(): # Every row has ink across the whole window: no band. words = [word(100, y, 300, y + 9, "wide") for y in range(130, 221, 13)] assert blank_column_bands(words, 150.0, 250.0, 120.0, 230.0) == [] def test_declash_side_by_side_snaps_the_shared_edge_column_to_the_blank_column(): words = _side_by_side_words() # The locator flagged cols A-G for the left table and F-L for the right # (one shared edge column); the tighten let each box reach into the # neighbour's stub, so the boxes overlap between x=282 and x=341. left = LocatedTable( ordinal=1, title="", region=(44 / PAGE_W, 120 / PAGE_H, 341 / PAGE_W, 230 / PAGE_H), grid_rows=(6, 11), grid_cols=(0, 6), tightened=True, ) right = LocatedTable( ordinal=2, title="", region=(282 / PAGE_W, 120 / PAGE_H, 568 / PAGE_W, 230 / PAGE_H), grid_rows=(6, 11), grid_cols=(5, 11), tightened=True, ) out = declash_side_by_side([left, right], words, PAGE_W, PAGE_H) l_x1 = out[0].region[2] * PAGE_W r_x0 = out[1].region[0] * PAGE_W assert abs(l_x1 - r_x0) < 1e-6 # the boxes now tile at the seam assert 295.0 < l_x1 < 314.0 assert out[0].region[0] == left.region[0] and out[1].region[2] == right.region[2] # outer edges untouched assert out[0].region[1] == left.region[1] and out[0].region[3] == left.region[3] # rows untouched def test_declash_side_by_side_leaves_stacked_and_separated_pairs_alone(): words = _side_by_side_words() upper = LocatedTable( ordinal=1, title="", region=(0.1, 0.10, 0.9, 0.40), grid_rows=(4, 14), grid_cols=(1, 10), tightened=True, ) lower = LocatedTable( ordinal=2, title="", region=(0.1, 0.38, 0.9, 0.70), grid_rows=(14, 25), grid_cols=(1, 10), tightened=True, ) assert [t.region for t in declash_side_by_side([upper, lower], words, PAGE_W, PAGE_H)] == [ upper.region, lower.region, ] left = LocatedTable( ordinal=1, title="", region=(0.07, 0.15, 0.48, 0.30), grid_rows=(6, 11), grid_cols=(0, 5), tightened=True, ) right = LocatedTable( ordinal=2, title="", region=(0.52, 0.15, 0.93, 0.30), grid_rows=(6, 11), grid_cols=(6, 11), tightened=True, ) assert [t.region for t in declash_side_by_side([left, right], words, PAGE_W, PAGE_H)] == [ left.region, right.region, ]