""" Table Inference with LLM Uses PydanticAI with various model providers to analyze tables and generate meaningful titles and descriptions based on document context. """ import asyncio import json import os from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Optional import logfire from loguru import logger from pydantic import BaseModel, Field if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument, NodeItem from quber.agents import AgentFactory from quber.agents.factory import ModelProvider from quber.core.consolidation import HeaderConsolidator from quber.core.models import TableMetadata from quber.core.parsers import parser_for_preset from quber.prompts import load_prompt from quber.settings import Settings, get_settings from quber.utils import setup_logging def provenance_page(item: "NodeItem") -> Optional[int]: """The 1-indexed page an item sits on, taken from its docling provenance. docling records an item's location in `item.prov` (a list of provenance spans, each with a `page_no`); the page lives there, not on a `_page` attribute. An item may carry several spans when it straddles a page break; the first span's page is used. Structural group nodes have no provenance and return None. `prov` is read via getattr because the base node type does not declare it (only laid-out items do).""" prov = getattr(item, "prov", None) if not prov: return None return prov[0].page_no class TableContext(BaseModel): """Input context for table analysis.""" executive_summary: str = Field(description="3-4 sentence document-level summary") procedural_title: str = Field(description="Title generated by procedural rules") headers: List[str] = Field(description="Individual headers found") descriptive_text: Optional[str] = Field(default=None, description="Text that describes the table") table_preview: str = Field(description="First few rows of the table") preceding_text: Optional[str] = Field(default=None, description="Text before the table") page_context: Optional[str] = Field( default=None, description="Extended page-level context including headers" ) page_number: int = Field(description="Page where table appears") class TableInfo(BaseModel): """LLM-inferred information about a table.""" title: str = Field(description="Clear, concise title for the table") description: str = Field( description="Functional description of what the table contains and how it might be used" ) class ExecutiveSummary(BaseModel): """Executive summary of a document.""" summary: str = Field(description="3-4 sentence executive summary of the document") class TableInferenceProcessor: """Process tables with LLM inference for enhanced metadata.""" def __init__( self, provider: Optional[ModelProvider] = None, model: Optional[str] = None, settings: Optional[Settings] = None, preset: str = "tuned-financial", ): """ Initialize the processor with specified configuration. Args: provider: Model provider to use (auto-detected if not specified) model: Specific model to use (uses default if not specified) settings: Settings object (uses the cached get_settings() if not provided). The optional-injection seam lets callers/tests pass an overridden Settings (e.g. CLI flags). preset: DoclingParser preset to use ("tuned-financial" canonical default, or "legacy" for docling's default backend and pipeline options with OCR off). """ self.settings = settings or get_settings() self.provider = provider self.model = model self.preset = preset # Set up logging setup_logging( service_name="table-inference", enable_logfire=self.settings.obs.enable_logfire, log_level=self.settings.log_level, ) # Load system prompt from TOML file system_prompt = load_prompt("table_inference", "system") # Create agent factory self.agent_factory = AgentFactory(enable_logfire=self.settings.obs.enable_logfire) # Create the agent self.agent = self.agent_factory.create_agent( output_type=TableInfo, system_prompt=system_prompt, provider=provider, model=model, ) logger.debug( f"Initialized TableInferenceProcessor with provider: {provider or 'auto'}, model: {model or 'default'}" ) async def analyze_table( self, metadata: TableMetadata, page_context: Optional[str] = None, executive_summary: Optional[str] = None, ) -> Dict[str, Any]: """ Analyze a single table with LLM inference. Args: metadata: Table metadata built by HeaderConsolidator.consolidate page_context: Extended context from the page executive_summary: Document-level executive summary Returns: Dictionary with both procedural and LLM-inferred information """ # Prepare context for LLM context = TableContext( executive_summary=executive_summary or "No document summary available", procedural_title=metadata.consolidated_header or "Untitled Table", headers=metadata.headers or [], descriptive_text=metadata.descriptive_text, table_preview=metadata.table_markdown or "", # Send full table for accurate analysis preceding_text=metadata.preceding_text, page_context=page_context, page_number=metadata.page_number, ) # Get LLM inference with telemetry try: prompt = self.format_prompt(context) if self.settings.obs.enable_logfire: with logfire.span("analyze_table", table_id=metadata.table_index, page=metadata.page_number): result = await self.agent.run(prompt) llm_info = result.output else: result = await self.agent.run(prompt) llm_info = result.output logger.debug(f"Successfully analyzed table {metadata.table_index}") except Exception as e: logger.error(f"LLM inference failed for table {metadata.table_index}: {e}") llm_info = TableInfo( title=metadata.consolidated_header or f"Table {metadata.table_index + 1}", description="Unable to generate description", ) # Combine all information return { "table_id": metadata.table_index, "page": metadata.page_number, "procedural_title": metadata.consolidated_header, "headers": metadata.headers, "descriptive_text": metadata.descriptive_text, "preceding_text": metadata.preceding_text, "llm_title": llm_info.title, "llm_description": llm_info.description, "table_markdown": metadata.table_markdown, "metadata": { "rows": metadata.rows, "cols": metadata.cols, "first_cell": metadata.first_cell_content, }, } async def analyze_tables_batch( self, metadata_list: List[TableMetadata], page_contexts: Optional[Dict[int, str]] = None, executive_summary: Optional[str] = None, max_concurrent: Optional[int] = None, ) -> List[Dict[str, Any]]: """ Analyze multiple tables concurrently with rate limiting. Args: metadata_list: List of table metadata to analyze page_contexts: Optional dict mapping table indices to page contexts executive_summary: Document-level executive summary max_concurrent: Maximum concurrent analyses (uses config default if None) Returns: List of analysis results """ max_concurrent = max_concurrent or self.settings.max_concurrent_tables page_contexts = page_contexts or {} # Create semaphore for rate limiting semaphore = asyncio.Semaphore(max_concurrent) async def analyze_with_limit(metadata: TableMetadata): async with semaphore: page_context = page_contexts.get(metadata.table_index) return await self.analyze_table(metadata, page_context, executive_summary) # Process all tables concurrently with rate limiting tasks = [analyze_with_limit(metadata) for metadata in metadata_list] with logfire.span( "analyze_tables_batch", total_tables=len(metadata_list), max_concurrent=max_concurrent ): results = await asyncio.gather(*tasks) return results def format_prompt(self, context: TableContext) -> str: """Format the context as a prompt for the agent.""" return f""" Analyze this table from page {context.page_number}: === DOCUMENT EXECUTIVE SUMMARY === {context.executive_summary} === EXTENDED PAGE CONTEXT === {context.page_context or "No extended page context available"} === IMMEDIATE TABLE CONTEXT === Procedural Title: {context.procedural_title} Headers Found: {", ".join(context.headers) if context.headers else "None"} Preceding Text (chronological order, furthest to closest): {context.preceding_text or "None"} Note: The LAST element in preceding text may be a footnote from the previous table. [Table] markers indicate preceding tables. === TABLE DATA (may be truncated at page boundary) === {context.table_preview} Based on ALL context levels (document summary, page context, procedural headers, and table data), extract the most appropriate title and provide a functional description for what this COMPLETE table represents. """ def get_table_preview(self, markdown: Optional[str], max_rows: int = 10) -> str: """Get preview of table (first N rows).""" if not markdown: return "No table content available" lines = markdown.strip().split("\n") preview_lines = lines[: min(max_rows + 2, len(lines))] # +2 for header and separator return "\n".join(preview_lines) def extract_page_context( self, document: "DoclingDocument", page_number: int, char_limit: int = 2000 ) -> str: """ Extract extended context from the page containing the table. Args: document: The Docling document object page_number: The page number (1-indexed) where the table appears char_limit: Maximum characters to extract for context Returns: String containing page-level context including headers and surrounding text """ context_parts: List[str] = [] char_count = 0 # Text on the table's own page. iterate_items() yields (item, level) # tuples; the page of each item comes from its provenance, and headers # carry a level. Text and level are read with getattr because only some # item types (text, headers) have them. current_page_text: List[str] = [] for item, _level in document.iterate_items(): if provenance_page(item) != page_number: continue text = (getattr(item, "text", "") or "").strip() if not text: continue level = getattr(item, "level", None) if level is not None: current_page_text.append(f"[Header L{level}] {text}") else: current_page_text.append(text) char_count += len(text) if char_count >= char_limit: break # If we have room and the table is not on the first page, add a tail of # the previous page for lead-in context. if char_count < char_limit * 0.7 and page_number > 1: # Use 70% to leave room prev_page = page_number - 1 prev_page_text: List[str] = [] for item, _level in document.iterate_items(): if provenance_page(item) != prev_page: continue text = (getattr(item, "text", "") or "").strip() if text: prev_page_text.append(text) # Take last ~500 chars from previous page if prev_page_text: prev_text = "\n".join(prev_page_text[-5:]) # Last 5 text elements if len(prev_text) > 500: prev_text = "..." + prev_text[-500:] context_parts.append("=== Previous Page Context ===\n" + prev_text) # Add current page context if current_page_text: context_parts.append( "=== Current Page Content ===\n" + "\n".join(current_page_text[:20]) ) # First 20 elements return "\n\n".join(context_parts) if context_parts else "No page context available" async def generate_executive_summary(self, document: "DoclingDocument") -> str: """ Generate a 3-4 sentence executive summary of the entire document. Args: document: The Docling document object Returns: Executive summary string """ # Export full document to markdown document_markdown = document.export_to_markdown() # Create a specialized agent for executive summary generation summary_agent = self.agent_factory.create_agent( output_type=ExecutiveSummary, system_prompt="You are analyzing business documents to create executive summaries.", provider=self.provider, model=self.model, ) # Create prompt for executive summary prompt = f""" Analyze this complete document and provide a 3-4 sentence executive summary. The summary should describe: 1. What this document is (type of document, purpose) 2. Who it is about (company, organization, or subject) 3. What are the major points being made (key metrics, themes, or findings) === DOCUMENT === {document_markdown} Provide a concise executive summary that captures the essential nature and content of this document. """ try: if self.settings.obs.enable_logfire: with logfire.span("generate_executive_summary"): result = await summary_agent.run(prompt) return result.output.summary else: result = await summary_agent.run(prompt) return result.output.summary except Exception as e: logger.error(f"Failed to generate executive summary: {e}") return "Unable to generate document summary" async def process_document(self, pdf_path: "str | os.PathLike[str]") -> Dict[str, Any]: """ Process all tables in a document. Args: pdf_path: Path to PDF document. Accepts a local path string or any os.PathLike (e.g. a cloudpathlib S3Path); Path(pdf_path) materializes the latter to a local cached file before parsing. Returns: Dictionary with document and table information """ # str() keeps the label cheap for S3Path (its str is the s3:// URI, not # a download); os.path.basename(S3Path) would force materialization. doc_name = os.path.basename(str(pdf_path)) # Log start of document processing logger.info(f"Starting document processing: {doc_name}") # Parse into the canonical DoclingDocument, then consolidate headers # as a separate post-processing step over that document. parser = parser_for_preset(self.preset) consolidator = HeaderConsolidator(separator=" ^ ", enable_consolidation=True, debug=False) with logfire.span("parse_document", document=str(pdf_path)): document = parser.parse(Path(pdf_path)).document table_metadata = consolidator.consolidate(document) # Log completion of resource-intensive document parsing logger.info(f"Document parsed successfully: {doc_name}") # Generate executive summary of the document logger.info(f"Generating executive summary for {doc_name}") executive_summary = await self.generate_executive_summary(document) logger.debug(f"Executive summary: {executive_summary[:100]}...") # Log table analysis details if len(table_metadata) > 0: model_info = self.model or "default" if self.provider: model_info = f"{self.provider.value}:{model_info}" logger.info(f"Analyzing {len(table_metadata)} tables using {model_info}") # Extract page contexts for all tables page_contexts = {} for metadata in table_metadata: page_contexts[metadata.table_index] = self.extract_page_context(document, metadata.page_number) # Process all tables concurrently tables = await self.analyze_tables_batch(table_metadata, page_contexts, executive_summary) if len(tables) > 0: logger.info(f"Analysis complete: {len(tables)} tables processed") # Build final output return { "document": doc_name, "extraction_date": datetime.now().isoformat(), "total_pages": len(document.pages), "total_tables": len(tables), "model_provider": self.provider.value if self.provider else "auto-detected", "model": self.model or "default", "tables": tables, } def save_json(self, data: Dict[str, Any], output_path: str): """Save results as JSON.""" with open(output_path, "w", encoding="utf-8") as f: json.dump(data, f, indent=2, ensure_ascii=False) logger.debug(f"JSON results saved to: {output_path}") def save_summary(self, data: Dict[str, Any], output_path: str): """Save human-readable summary.""" lines = [] lines.append("# Table Analysis Report") lines.append("") lines.append(f"**Document:** {data['document']}") lines.append("") lines.append(f"**Date:** {data['extraction_date']}") lines.append("") lines.append(f"**Model:** {data.get('model_provider', 'unknown')}/{data.get('model', 'unknown')}") lines.append("") lines.append(f"**Total Tables:** {data['total_tables']}") lines.append("") for table in data["tables"]: lines.append(f"\n## Table {table['table_id'] + 1} (Page {table['page']})") # Show procedural information lines.append("\n### Procedural Extraction:") lines.append(f"- **Consolidated Headers:** {table['procedural_title'] or 'None'}") if table.get("headers"): lines.append(f"- **Individual Headers:** {', '.join(table['headers'])}") if table.get("descriptive_text"): lines.append(f"- **Descriptive Text:** {table['descriptive_text'][:200]}...") # Show LLM analysis lines.append("\n### LLM Analysis:") lines.append(f"- **Title:** {table['llm_title']}") lines.append(f"- **Description:** {table['llm_description']}") # Show metadata lines.append("\n### Metadata:") lines.append( f"- **Size:** {table['metadata']['rows']} rows × {table['metadata']['cols']} columns" ) if table["metadata"].get("first_cell"): lines.append(f"- **First Cell:** {table['metadata']['first_cell'][:50]}...") # Show table preview (first few rows) if table.get("table_markdown"): lines.append("\n### Table Preview:") preview_lines = table["table_markdown"].split("\n")[:7] # First 7 lines for line in preview_lines: lines.append(line) if len(table["table_markdown"].split("\n")) > 7: lines.append("| ... | ... | (additional rows omitted) |") lines.append("\n---") with open(output_path, "w", encoding="utf-8") as f: f.write("\n".join(lines)) logger.debug(f"Summary saved to: {output_path}")