""" TableUnifier Protocol — unifies the Camelot flavor outputs for a page into one canonical table per physical table. Camelot is run in two flavors (lattice + stream) on purpose: each captures structure the other misses. When both produce an extraction of the same physical table, the unifier combines them into one canonical table (most complete header from one, full row set from the other), using the page image as ground truth. Candidates that are genuinely distinct tables pass through untouched. This is not deduplication — the second flavor's output is a deliberate, complementary source, not an unwanted duplicate. For each page with multiple surviving classifier-positive candidates, the unifier receives the page image and all candidates and returns the canonical set. Pure-LLM (no geometric heuristics): the LLM judges which candidates describe the same physical table and emits the unified best version. Kept separate from `LLMClient` for the same reasons as `TableClassifier`: single-responsibility, potentially different model than structure-correction, isolated retry semantics. """ from __future__ import annotations import json from pathlib import Path from typing import Any, Dict, List, Literal, Optional, Protocol, Tuple, runtime_checkable from loguru import logger from pydantic import BaseModel, Field from quber.agents.langsmith_tracer import LangSmithTracer, usage_metadata_from from quber.settings import DEFAULT_LLM_MODEL, get_settings class CandidateInput(BaseModel): candidate_id: str = Field(description="Stable id used to reference this candidate in the unifier output") markdown: str = Field(description="Camelot-extracted markdown for this candidate") bbox: Optional[Tuple[float, float, float, float]] = Field( default=None, description="Camelot's geometric box (x1, y1, x2, y2); informational, " "not load-bearing for the LLM's unify decision", ) class UnifiedTable(BaseModel): markdown: str = Field(description="Canonical markdown for this logical table") source_candidate_ids: List[str] = Field( description="candidate_ids the unifier combined into this output. " "Single-id list means the candidate stood alone; multi-id means an actual unify." ) class UnifierResult(BaseModel): tables: List[UnifiedTable] = Field( default_factory=list, description="Canonical tables on this page after unifying the flavor outputs", ) # Schema embedded so the model is the single source of truth for the # response shape. The page image is sent alongside this prompt as the # ground-truth reference for unify decisions. UNIFY_FLAVORS_PROMPT = f"""\ You are reconciling candidate tables extracted from a single PDF page by different table-extraction strategies. Multiple candidates may represent the same logical table (extracted twice with slight differences) or distinct tables that happen to be on the same page. You are given: 1. A rendered image of the PDF page (ground truth). 2. A list of candidate tables, each with an id, markdown, and optional bbox. Your job: return the canonical set of tables present on this page. Rules: - *Same logical table* — two candidates describe the same physical table on the page (even if their bboxes, row counts, or cell content differ slightly). Return ONE unified entry combining the best parts of each (most complete header, all data rows, correct column count). List ALL contributing candidate_ids in `source_candidate_ids`. - *Distinct tables* — two candidates are different physical tables that happen to be on the same page. Return each as its own entry with `source_candidate_ids` containing only its own id. - *Use the page image* as the source of truth for which physical tables exist and what their structure should be. - *Preserve numeric values exactly* from whichever candidate(s) you draw them from. Do not invent or alter numbers. Return a JSON object conforming exactly to this schema: {json.dumps(UnifierResult.model_json_schema(), indent=2)} Return only the JSON object — no prose, no markdown code fences. """ @runtime_checkable class TableUnifier(Protocol): async def unify( self, page_image: Path, candidates: List[CandidateInput], ) -> UnifierResult: ... def candidates_payload(candidates: List[CandidateInput]) -> str: return json.dumps( [c.model_dump() for c in candidates], indent=2, ) class PydanticAIUnifier: """TableUnifier backed by pydantic-ai with an Anthropic model. Sends the page image as `BinaryContent` alongside the candidate payload — the image is the ground-truth reference for unify decisions. """ 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 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( "PydanticAIUnifier: 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=UnifierResult, system_prompt=UNIFY_FLAVORS_PROMPT, ) # LangSmith tracer is a no-op if `TRACE_TO_LANGSMITH` is unset. self.tracer = LangSmithTracer(run_name="quber-unifier") def trace_inputs(self, user_text: str) -> Dict[str, Any]: return { "messages": [ {"role": "system", "content": UNIFY_FLAVORS_PROMPT}, { "role": "user", "content": [ {"type": "text", "text": user_text}, {"type": "image", "text": "[page image attached]"}, ], }, ] } async def unify( self, page_image: Path, candidates: List[CandidateInput], ) -> UnifierResult: from pydantic_ai import BinaryContent if not candidates: return UnifierResult(tables=[]) user_text = f"Candidates on this page:\n{candidates_payload(candidates)}" image = BinaryContent(data=Path(page_image).read_bytes(), media_type="image/png") inputs = self.trace_inputs(user_text) try: async with self.tracer.llm_run("unify_flavors", inputs, model=self.model) as run: result = await self.agent.run([user_text, image]) output = result.output run.outputs = { "messages": [{"role": "assistant", "content": output.model_dump_json()}], "table_count": len(output.tables), "usage_metadata": usage_metadata_from(result.usage), } return output except Exception as exc: logger.error( "unify: LLM call failed; returning candidates as-is page_image={} count={} exc={}", page_image.name, len(candidates), exc, ) return UnifierResult( tables=[ UnifiedTable(markdown=c.markdown, source_candidate_ids=[c.candidate_id]) for c in candidates ] ) class MockUnifier: """Returns a canned UnifierResult, or pass-through if none injected. Pass-through behavior: emit one UnifiedTable per candidate, each referencing only its own candidate_id (no unifying). For tests. """ def __init__(self, result: Optional[UnifierResult] = None) -> None: self.result = result async def unify( self, page_image: Path, candidates: List[CandidateInput], ) -> UnifierResult: _ = page_image if self.result is not None: return self.result return UnifierResult( tables=[ UnifiedTable(markdown=c.markdown, source_candidate_ids=[c.candidate_id]) for c in candidates ] ) UnifierBackend = Literal["api", "mock"] def get_unifier(backend: Optional[UnifierBackend] = None) -> TableUnifier: selected = backend or get_settings().llm.unifier_backend if selected == "api": return PydanticAIUnifier() if selected == "mock": return MockUnifier() raise ValueError(f"Unknown QUBER_UNIFIER_BACKEND: {selected!r}. Expected api|mock.")