"""FigureCorrector Protocol — what a figure the scan has read leaves on the page. Three things, from one look at the page image: which of the text around the figure is its furniture, which footnote markers the figure's own labels carry, and the note lines printed for each figure, as marker and text. The caller records the markers and notes on the figure's picture. Once the scan has read a figure, the fragments the parse lifted off the same image are that figure read worse, and they are removed. Removing them by tree position only reaches the ones the parse filed under the picture. It files inconsistently: on one page of a quarterly deck it put twenty-two of a page's twenty-six texts beneath the two pictures, and on another it left forty-six of seventy-eight attached to the page body, where nothing looking at the picture's children can see them. Those forty-six included forty axis tick labels. Position on the page does not settle it either, and neither does the text. Both were measured on three decks and both fail, in opposite directions. Sweeping everything printed inside the figure's box deletes real content. On one page the scan bounded two charts as a single figure spanning the width of the page, and inside that box sat a four-line sentence about the portfolio's sensitivity to rate changes that the scan did not transcribe. Sweeping only what the scan's own description repeats deletes nothing useful. A chart's axis runs 0.00 to 5.00 in half steps and the description never mentions those numbers, because a gridline is a ruler marking rather than a fact. All forty ticks would stay. What separates the two is what the text is, which is visible on the page. A tick label sits at the edge of the plot in an evenly spaced series. An annotation sits over the plot and reads as language. So the page image goes to an agent with the description and the candidate fragments, and it names the furniture. It is told to keep whatever it is unsure of. Leaving a tick label in costs a chunk that answers nothing; removing a sentence costs a statement the document no longer makes anywhere. A failed call removes nothing at all. The smaller model is not sufficient here. Asked the same question on the same pages it removed the four-line sentence while its own stated reason called it an annotation, and missed seven of the forty ticks. The stronger model returned all forty and none of the sentence, identically on three runs. The markers come from the same call because they are the same act of reading. A chart's label prints its qualification as a raised digit — "Undepreciated Book Equity Value(1,2)" — and the marker alone says only that a qualification exists, not what it says. Read on its own, that chart states a figure the document elsewhere excludes a noncontrolling interest from and calls non-GAAP. Which markers a figure carries is what a table's correction agent reports for a table, off the table image, for the same reason and against the same distinctions: a raised digit is a marker, a parenthesized negative value is not, and a reference naming a section of the document points somewhere a footnote lookup will never find. A figure had nobody reporting it, and this agent is already looking at the page the markers are printed on. """ from __future__ import annotations import json from typing import Any, Dict, List, Literal, Optional, Protocol, Sequence, runtime_checkable from loguru import logger from pydantic import BaseModel, Field, field_validator from quber.agents.factory import supports_sampling_temperature from quber.agents.langsmith_tracer import LangSmithTracer, usage_metadata_from from quber.agents.llm_client import COMPOUND_MARKER_RE, MARKER_GROUP_RE from quber.settings import DEFAULT_LLM_MODEL, get_settings class Figure(BaseModel): """One figure the scan read on this page, numbered so markers name it.""" index: int = Field(description="Position in the list the agent was shown") description: str = Field(description="What the scanner read the figure as saying") class Fragment(BaseModel): """One candidate, numbered so the agent answers by index rather than by text.""" index: int = Field(description="Position in the list the agent was shown") text: str = Field(description="What the fragment says") class Furniture(BaseModel): index: int = Field(description="The index of a fragment that is the figure's furniture") reason: str = Field(description="Why it is furniture, in a few words") class Marker(BaseModel): """One footnote reference marker printed on one figure's labels.""" figure: int = Field(description="The index of the figure carrying it, from the list shown") marker: str = Field( description=( "ONE reference marker exactly as printed: '1', '(1)', a letter 'a', an " "asterisk '*', a dagger. A label printed with several markers at once — " "'(1,2)', '(b)(c)' — carries several: report one entry per marker, never " "the combined string." ) ) label: str = Field( default="", description="The label the marker is printed on, copied as it appears on the figure", ) kind: Literal["footnote", "section"] = Field( default="footnote", description=( "What the marker points at, judged from the PAGE, not from the marker's " "wording. 'footnote': it points at a note printed for this figure — a note " "line below it or nearby that defines it, EVEN WHEN that line's own text " "sends the reader onward ('see Selected Definitions on page 40' is a " "printed note and the answer is that line). 'section': it is a " "cross-reference to a named part of the document (a numbered note in the " "financial statements, a schedule, an appendix) whose content lives " "elsewhere, not in a note line printed for this figure — and the label " "names that part, so the marker is written WITH its word ('Note 16', " "'Schedule II'), never as a bare number. A bare marker is never 'section'." ), ) class Note(BaseModel): """One note printed on the page for a figure: its own marker plus its text. The marker is read off the image, where the note line prints it. The image's formatting decides what is a marker; an unmarked general note has none. """ figure: int = Field(description="The index of the figure the note is printed for") marker: str = Field( description=( "The note's OWN marker exactly as printed at the start of its line: '1', " "'(1)', 'a', '*', a dagger. EMPTY for an unmarked general note that opens " "with no marker." ) ) text: str = Field(description="The note's text, without the leading marker") class FigureCorrection(BaseModel): furniture: List[Furniture] = Field( default_factory=list, description="Only the furniture; every fragment not listed stays" ) markers: List[Marker] = Field( default_factory=list, description="Every footnote reference marker printed on a figure's labels; empty if none", ) @field_validator("furniture", mode="after") @classmethod def one_verdict_per_fragment(cls, v: List[Furniture]) -> List[Furniture]: """A fragment named twice is named once. The agent sometimes returns the same index in two entries with different reasons. Both name one text item, and the caller deletes what it is given, so a repeat reaches the document model as the same element queued for removal twice — which it rejects, failing the whole sweep. The first reason is kept. """ seen: set[int] = set() return [f for f in v if not (f.index in seen or seen.add(f.index))] @field_validator("markers", mode="after") @classmethod def split_compound_markers(cls, v: List[Marker]) -> List[Marker]: """A stacked group reported as one marker becomes one marker per group. The prompt asks for one entry per marker and compliance is not certain, so a label printed '(b)(c)(d)' that comes back as a single string is split here — the same normalization a table's markers get, against the same pattern, which is strict enough that a parenthesized negative value can never match it. """ out: List[Marker] = [] for mark in v: s = mark.marker.strip() if COMPOUND_MARKER_RE.fullmatch(s): out.extend(mark.model_copy(update={"marker": g}) for g in MARKER_GROUP_RE.findall(s)) else: out.append(mark) return out notes: List[Note] = Field( default_factory=list, description="Every note line printed on the page for a figure, as marker and text", ) FIGURE_CORRECTION_PROMPT = f"""\ You are shown one page of a financial document, a numbered list of the figures a scanner read on that page with the description it wrote for each, and a numbered list of text fragments a parser extracted from inside those figures' regions. You have two jobs. JOB 1 — NAME THE FURNITURE. Decide which fragments are the figure's FURNITURE: text that exists only to mark its scale or label its parts, and carries no fact on its own. Axis tick values, gridline labels, the numbers running up a y-axis or along an x-axis, legend keys and unit markers are furniture. They are already accounted for by the description, which reports what the figure shows. Everything else STAYS. A sentence, a callout, an annotation, a note, a heading, a stated figure that is not a position on an axis — all of these say something and must be kept even when they sit inside the figure. Look at the image. A tick label sits at the edge of the plot in a regular series, evenly spaced, ascending or descending. An annotation sits over or beside the plot and reads as language. That difference is visible and it is what you are judging. Default to keeping. If you are not sure a fragment is furniture, leave it out of your answer. Removing a real statement is far worse than leaving a tick label in. When the fragment list is EMPTY there is nothing to judge: return an empty furniture list and do job 2 only. Never name an index that was not shown to you. JOB 2 — READ THE FOOTNOTE MARKERS. Report every FOOTNOTE REFERENCE MARKER printed on a figure's labels — its title, its axis names, its series names, its data labels — read off the IMAGE, because a raised digit is often lost from extracted text. A marker is a superscript or parenthetical number, letter or symbol pointing at a note: '1', '(1)', a letter 'a', an asterisk '*', a dagger '†' or '‡', a section sign '§'. Write each as it appears, and name the figure it belongs to by its index. These are NOT markers. A parenthesized NEGATIVE VALUE like '(84)'. A unit like '(%)' or '($m)'. A year, a count, or any number that is the thing being stated rather than a pointer away from it. A label naming a NAMED SECTION of the document — '(Note 16)', 'Schedule II' — is written as printed WITH its word, never as a bare number. The kind says what the marker POINTS AT, judged from the page. A marker pointing at a note printed for this figure is 'footnote' — including when that note's own text sends the reader onward. 'For a description of this non-GAAP financial measure, see Selected Definitions on page 40' is a printed note, and the answer is that line, so its marker is 'footnote'. 'section' is for a label that CROSS-REFERENCES a named part of the document whose content lives elsewhere, not in a note printed here. Such a label names its target in words, so its marker is written with that word — 'Note 16', 'Schedule II'. A bare number or letter is never 'section'; a marker with no note printed for it is still 'footnote', and reporting it that way is what says the note is missing. One entry per marker. A label reading 'Book Value(1,2)' carries two markers and gets two entries, never one reading '1,2'. Empty list if the page's figures carry no markers, which is the common case. JOB 3 — READ THE NOTES THEMSELVES. Report every note line printed on the page for a figure, as a PAIR: the note's OWN marker exactly as printed at the start of its line, and its text without that marker. Read them off the IMAGE. The image's formatting decides what is a marker. A general note that opens with no marker — a basis-of-presentation line, a line beginning 'Note:' — gets an EMPTY marker and keeps its full text. A line whose text sends the reader onward is still a note: report it as printed. Notes are printed at the foot of the page, beneath the figure, or beside it. They often run together on one physical line: 'Note: As of 03/31/2026. 1. For a description of this measure, see Selected Definitions on page 40.' is a general note AND note 1 — report them separately. Every marker you reported in job 2 should have a note here unless no line on the page defines it. Report a note even when no marker points at it. Return a JSON object conforming exactly to this schema: {json.dumps(FigureCorrection.model_json_schema(), indent=2)} Return only the JSON object — no prose, no markdown code fences. """ @runtime_checkable class FigureCorrector(Protocol): async def correct_figure( self, page_image: bytes, figures: Sequence[Figure], fragments: Sequence[Fragment] ) -> FigureCorrection: ... def figure_block(figures: Sequence[Figure]) -> str: """The page's figures, numbered the way a marker refers to them.""" return "\n\n".join(f"FIGURE {f.index}:\n{f.description}" for f in figures) def fragment_block(fragments: Sequence[Fragment]) -> str: """The candidates, one per line, numbered the way the answer refers to them.""" return "\n".join(f"{f.index}: {f.text}" for f in fragments) class PydanticAIFigureCorrector: """FigureCorrector backed by pydantic-ai with an Anthropic model. Auth resolution mirrors the other agents: `ANTHROPIC_AUTH_TOKEN` (OAuth subscription) preferred over `ANTHROPIC_API_KEY`. """ # Which text on a page is a figure's furniture does not vary between runs on # the same page. Greedy decoding removes the sampling variance. 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, ) -> 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.figure_correction_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( "PydanticAIFigureCorrector: neither ANTHROPIC_AUTH_TOKEN nor " "ANTHROPIC_API_KEY is set. Provide one via env or constructor." ) self.model = model self.temperature = temperature model_settings = None if temperature is None else ModelSettings(temperature=temperature) self.agent = Agent( anth_model, output_type=FigureCorrection, system_prompt=FIGURE_CORRECTION_PROMPT, model_settings=model_settings, ) # LangSmith tracer is a no-op if `TRACE_TO_LANGSMITH` is unset. self.tracer = LangSmithTracer(run_name="quber-figure-correction") def trace_inputs(self, figures: Sequence[Figure], fragments: Sequence[Fragment]) -> Dict[str, Any]: return { "messages": [ {"role": "system", "content": FIGURE_CORRECTION_PROMPT}, { "role": "user", "content": [ {"type": "text", "text": figure_block(figures)}, {"type": "text", "text": fragment_block(fragments)}, {"type": "image", "text": "[page image attached]"}, ], }, ] } async def correct_figure( self, page_image: bytes, figures: Sequence[Figure], fragments: Sequence[Fragment] ) -> FigureCorrection: from pydantic_ai import BinaryContent image = BinaryContent(data=page_image, media_type="image/png") described = f"Figures the scanner read on this page:\n{figure_block(figures)}" listed = f"Fragments extracted from inside the figure region:\n{fragment_block(fragments)}" try: async with self.tracer.llm_run( "figure_correction", self.trace_inputs(figures, fragments), model=self.model ) as run: result = await self.agent.run([described, listed, image]) output = result.output run.outputs = { "messages": [{"role": "assistant", "content": output.model_dump_json()}], "furniture": len(output.furniture), "markers": len(output.markers), "notes": len(output.notes), "usage_metadata": usage_metadata_from(result.usage), } asked = {f.index for f in fragments} kept = [f for f in output.furniture if f.index in asked] if len(kept) != len(output.furniture): # Shown no fragments, the agent sometimes answers the furniture # question anyway, naming indices from the figure's own labels. # Nothing can be removed from an empty list, so that costs # nothing and is not worth an alarm; a verdict off the end of a # list that WAS shown is a real miss. report = logger.debug if not fragments else logger.warning report( "figure_correction: ignoring {} verdict(s) naming a fragment not in the list of {}", len(output.furniture) - len(kept), len(fragments), ) shown = {f.index for f in figures} marks = [m for m in output.markers if m.figure in shown and m.marker.strip()] if len(marks) != len(output.markers): logger.warning( "figure_correction: ignoring {} marker(s) naming a figure not in the list", len(output.markers) - len(marks), ) notes = [n for n in output.notes if n.figure in shown and n.text.strip()] return FigureCorrection(furniture=kept, markers=marks, notes=notes) except Exception as exc: logger.error( "figure_correction: LLM call failed; all {} fragment(s) are kept and no markers " "are recorded exc={}", len(fragments), exc, ) return FigureCorrection(furniture=[]) class MockFigureCorrector: """Names a fixed set of fragment texts as furniture, and fixed markers. For tests.""" def __init__( self, furniture: Optional[Sequence[str]] = None, markers: Optional[Sequence[Marker]] = None, notes: Optional[Sequence[Note]] = None, ) -> None: self.furniture = set(furniture or ()) self.markers = list(markers or ()) self.notes = list(notes or ()) self.calls: List[List[Fragment]] = [] self.figures: List[List[Figure]] = [] async def correct_figure( self, page_image: bytes, figures: Sequence[Figure], fragments: Sequence[Fragment] ) -> FigureCorrection: _ = page_image self.calls.append(list(fragments)) self.figures.append(list(figures)) shown = {f.index for f in figures} return FigureCorrection( furniture=[ Furniture(index=f.index, reason="mock") for f in fragments if f.text in self.furniture ], markers=[m for m in self.markers if m.figure in shown], notes=[n for n in self.notes if n.figure in shown], ) FigureCorrectionBackend = Literal["api", "mock", "off"] def get_figure_corrector(backend: Optional[FigureCorrectionBackend] = None) -> Optional[FigureCorrector]: """The figure corrector for this run, or None when the sweep is switched off.""" selected = backend or get_settings().llm.figure_correction_backend if selected == "api": return PydanticAIFigureCorrector() if selected == "mock": return MockFigureCorrector() if selected == "off": return None raise ValueError(f"Unknown QUBER_FIGURE_CORRECTION_BACKEND: {selected!r}. Expected api|mock|off.")