""" CaptureAdvisor — recommends a Camelot retry adjustment for a flagged capture drop. Camelot occasionally drops printed values during cell assignment even though they sit inside the table region and its text extraction captured them. The measured example: a totals row typeset on a raised baseline splits into two row bands at the default grouping tolerance, and the numbers' vertical center lands exactly on the band boundary, failing the strict containment test on both sides — the row's label and dollar signs survive, its numbers vanish. The drop is detected deterministically (numeric tokens in the region's text layer that appear in no grid cell) and the retry is accepted deterministically (it must recover the missing tokens and lose none). This agent sits between those two checks: given the cropped table image, the extracted grid, and the dropped tokens with their printed lines, it diagnoses the failure and recommends ONE bounded retry adjustment. The agent only proposes; the acceptance gate holds authority, so a wrong recommendation costs one extra Camelot pass and nothing else. """ from __future__ import annotations from typing import Any, Dict, Literal, Optional, Protocol, Sequence, 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 class RetryAdvice(BaseModel): """One bounded Camelot retry adjustment, with the diagnosis behind it.""" diagnosis: str = Field( description=( "One or two sentences: WHY were these tokens dropped from the grid, " "grounded in what the image and the provided evidence show." ) ) row_tol: int = Field( ge=1, le=15, description=( "Vertical grouping tolerance in points for the retry. The failed pass " "used 2. Raise it when one printed row was split into two bands (e.g. " "raised or offset baselines); keep 2 if row grouping is not the problem." ), ) column_tol: int = Field( default=0, ge=-10, le=15, description=( "Column grouping tolerance for the retry. The failed pass used 0. " "Change only if the drop is a column-assignment problem." ), ) flavor: Literal["stream", "lattice"] = Field( default="stream", description=( "Parser for the retry. Use lattice only if the table has ruled cell borders drawn on the page." ), ) ADVISE_PROMPT = """\ You diagnose a table-extraction capture failure and recommend ONE retry adjustment. A region-constrained Camelot pass extracted the grid below, but the listed numeric tokens — visible on the page and present in the PDF text layer inside the table region — did NOT land in any grid cell. You get: 1. The cropped table image (what the page really shows). 2. The extracted grid (markdown). 3. The dropped tokens, each with the full printed line it sits on. 4. The knobs the failed pass used: flavor=stream, row_tol=2, column_tol=0. Recommend the single adjustment most likely to capture the dropped tokens without corrupting the rest of the grid. Ground your diagnosis in the evidence; do not guess beyond it. """ @runtime_checkable class CaptureAdvisor(Protocol): async def recommend( self, crop_png: bytes, grid_markdown: str, dropped_lines: Sequence[str], page: int ) -> Optional[RetryAdvice]: ... class PydanticAICaptureAdvisor: """CaptureAdvisor backed by pydantic-ai with an Anthropic vision model. Mirrors `PydanticAIGridLocator`: same auth resolution (`ANTHROPIC_AUTH_TOKEN` preferred over `ANTHROPIC_API_KEY`), greedy decoding so the same failure yields the same recommendation run to run. A failed call returns None — the caller keeps the original grid with its warned gap; a capture repair must never take a table down with it. """ DEFAULT_TEMPERATURE = 0.0 def __init__( self, model: Optional[str] = None, auth_token: Optional[str] = None, api_key: Optional[str] = None, ) -> 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.model or DEFAULT_LLM_MODEL 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( "PydanticAICaptureAdvisor: neither ANTHROPIC_AUTH_TOKEN nor " "ANTHROPIC_API_KEY is set. Provide one via env or constructor." ) self.model = model self.agent = Agent( anth_model, output_type=RetryAdvice, system_prompt=ADVISE_PROMPT, model_settings=ModelSettings(temperature=self.DEFAULT_TEMPERATURE) if supports_sampling_temperature(model) else None, ) self.tracer = LangSmithTracer(run_name="quber-capture-advisor") def trace_inputs(self, user_text: str) -> Dict[str, Any]: return { "messages": [ {"role": "system", "content": ADVISE_PROMPT}, {"role": "user", "content": [{"type": "text", "text": user_text}]}, ] } async def recommend( self, crop_png: bytes, grid_markdown: str, dropped_lines: Sequence[str], page: int ) -> Optional[RetryAdvice]: from pydantic_ai import BinaryContent user_text = ( f"=== Extracted grid ===\n{grid_markdown}\n\n" "=== Dropped tokens ===\n" + "\n".join(dropped_lines) + "\n\n" "knobs used: flavor=stream, row_tol=2, column_tol=0" ) inputs = self.trace_inputs(user_text) try: async with self.tracer.llm_run("capture_advise", inputs, model=self.model) as run: result = await self.agent.run( [user_text, BinaryContent(data=crop_png, media_type="image/png")] ) advice: RetryAdvice = result.output run.outputs = { "messages": [{"role": "assistant", "content": advice.model_dump_json()}], "usage_metadata": usage_metadata_from(result.usage), } return advice except Exception as exc: logger.warning("page {}: capture advisor call failed; keeping original grid: {}", page, exc) return None class MockCaptureAdvisor: """Returns a canned RetryAdvice. For tests.""" def __init__(self, advice: Optional[RetryAdvice] = None) -> None: self.advice = advice self.calls: list[tuple[str, list[str], int]] = [] async def recommend( self, crop_png: bytes, grid_markdown: str, dropped_lines: Sequence[str], page: int ) -> Optional[RetryAdvice]: _ = crop_png self.calls.append((grid_markdown, list(dropped_lines), page)) return self.advice def get_capture_advisor(backend: Optional[str] = None) -> Optional[CaptureAdvisor]: """Resolve the capture advisor, or None when disabled (`off`).""" selected = backend or get_settings().llm.capture_advisor_backend if selected == "api": return PydanticAICaptureAdvisor() if selected == "mock": return MockCaptureAdvisor() if selected == "off": return None raise ValueError(f"Unknown QUBER_CAPTURE_ADVISOR_BACKEND: {selected!r}. Expected api|mock|off.")