# pyright: ignore """ Table Inference with LLM Uses PydanticAI with Anthropic to analyze tables and generate meaningful titles and descriptions based on document context. """ import asyncio import json import os from typing import List, Optional from datetime import datetime from pydantic import BaseModel, Field from pydantic_ai import Agent from dotenv import load_dotenv import logfire from table_parser_enhanced import EnhancedTableParser, TableMetadata # Load environment variables load_dotenv() # Configure Logfire for telemetry (optional) if os.getenv("LOGFIRE_TOKEN"): logfire.configure(service_name="table-inference") class TableContext(BaseModel): """Input context for table analysis""" procedural_title: str = Field(description="Title generated by procedural rules") headers: List[str] = Field(description="Individual headers found") descriptive_text: Optional[str] = Field(description="Text that describes the table") table_preview: str = Field(description="First few rows of the table") preceding_text: Optional[str] = Field(description="Text before the table") page_context: Optional[str] = Field(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 TableInferenceProcessor: """Process tables with LLM inference for enhanced metadata""" def __init__(self, model: str = 'claude-3-haiku-20240307'): """ Initialize the processor with specified model. Args: model: Anthropic model to use (haiku, sonnet, or opus) """ self.model = model # Create PydanticAI agent self.agent = Agent( model, output_type=TableInfo, # Specify the output type here system_prompt="""You are analyzing tables from business documents. You will be given both immediate context and extended page-level context. CONTEXT HIERARCHY: 1. Page Context: May include document headers, company names, report titles, section names, and other page-level information 2. Procedural Headers: Headers found immediately before the table 3. Table Preview: First few rows of the actual table TITLE EXTRACTION GUIDELINES: - Look for the most specific and informative title by examining ALL context levels - Prioritize headers that include: * Company or entity names * Document/report type (e.g., "Reconciliation of Non-GAAP Financial Measures") * Audit status (e.g., "Unaudited", "Audited") * Specific calculation or data description - Combine relevant headers hierarchically when appropriate - If page context shows "continued" sections, include the main section title IMPORTANT WARNINGS: - You are only seeing a preview of the table's first few rows - The complete table may contain additional sections not shown - Base your analysis on ALL available context, not just the visible rows - Avoid assumptions about partial data (e.g., don't say "Assets Section" when it might be a full balance sheet) DESCRIPTION GUIDELINES: - Describe what the COMPLETE table represents - Include the business purpose and potential uses - Note time periods, metrics, or comparisons shown - Be concise but comprehensive""" ) async def analyze_table(self, metadata: TableMetadata, page_context: Optional[str] = None) -> dict: """ Analyze a single table with LLM inference. Args: metadata: Table metadata from enhanced parser Returns: Dictionary with both procedural and LLM-inferred information """ # Prepare context for LLM context = TableContext( procedural_title=metadata.consolidated_header or "Untitled Table", headers=metadata.headers or [], descriptive_text=metadata.descriptive_text, table_preview=self._get_table_preview(metadata.table_markdown), preceding_text=metadata.preceding_text, page_context=page_context, page_number=metadata.page_number ) # Get LLM inference try: # Format the context as a prompt for the agent prompt = f""" Analyze this table from page {context.page_number}: === 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'} Descriptive Text: {context.descriptive_text or 'None'} === TABLE PREVIEW === (FIRST FEW ROWS ONLY - complete table continues beyond this) {context.table_preview} Based on ALL context levels (page, procedural, and table), extract the most informative title and provide a functional description for what this COMPLETE table represents. Prioritize finding headers that include company names, report types, and specific descriptions. """ result = await self.agent.run(prompt) # Get the output from the result llm_info = result.output except Exception as e: import traceback print(f"LLM inference failed for table {metadata.table_index}: {e}") if os.getenv("DEBUG"): traceback.print_exc() 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, "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 } } 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(' ') preview_lines = lines[:min(max_rows + 2, len(lines))] # +2 for header and separator return ' '.join(preview_lines) def _extract_page_context(self, document, 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 """ page_idx = page_number - 1 # Convert to 0-indexed context_parts = [] char_count = 0 # First, get text from the current page current_page_text = [] for element in document.iterate_items(): if hasattr(element, '_page') and element._page == page_idx: if hasattr(element, 'text') and element.text: text = element.text.strip() if text: # Include element type for better context if hasattr(element, 'level'): current_page_text.append(f"[Header L{element.level}] {text}") else: current_page_text.append(text) char_count += len(text) if char_count >= char_limit: break # If we have room and table is not at top of page, get some previous page context if char_count < char_limit * 0.7 and page_idx > 0: # Use 70% to leave room prev_page_text = [] prev_page_idx = page_idx - 1 # Get last portion of previous page for element in document.iterate_items(): if hasattr(element, '_page') and element._page == prev_page_idx: if hasattr(element, 'text') and element.text: text = element.text.strip() if text: prev_page_text.append(text) # Take last ~500 chars from previous page if prev_page_text: prev_text = ' '.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 === " + prev_text) # Add current page context if current_page_text: context_parts.append("=== Current Page Content === " + ' '.join(current_page_text[:20])) # First 20 elements return ' '.join(context_parts) if context_parts else "No page context available" async def process_document(self, pdf_path: str) -> dict: """ Process all tables in a document. Args: pdf_path: Path to PDF document Returns: Dictionary with document and table information """ print(f"Processing document: {pdf_path}") # Parse document with enhanced parser parser = EnhancedTableParser( separator=" ^ ", enable_consolidation=True, debug=False ) document, table_metadata = parser.parse_document(pdf_path) print(f"Found {len(table_metadata)} tables") # Process each table with LLM tables = [] for i, metadata in enumerate(table_metadata): print(f" Analyzing table {i + 1}/{len(table_metadata)}...") # Extract extended page context for this table page_context = self._extract_page_context(document, metadata.page_number) table_info = await self.analyze_table(metadata, page_context) tables.append(table_info) # Build final output return { "document": os.path.basename(pdf_path), "extraction_date": datetime.now().isoformat(), "total_pages": len(document.pages), "total_tables": len(tables), "tables": tables } def save_json(self, data: dict, 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) print(f"Results saved to: {output_path}") def save_summary(self, data: dict, output_path: str): """Save human-readable summary""" lines = [] lines.append("# Table Analysis Report") lines.append(f" **Document:** {data['document']}") lines.append(f"**Date:** {data['extraction_date']}") lines.append(f"**Total Tables:** {data['total_tables']}") lines.append("") for table in data['tables']: lines.append(f" ## Table {table['table_id'] + 1} (Page {table['page']})") # Show procedural information lines.append(" ### 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(" ### LLM Analysis:") lines.append(f"- **Title:** {table['llm_title']}") lines.append(f"- **Description:** {table['llm_description']}") # Show metadata lines.append(" ### 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(" ### Table Preview:") preview_lines = table['table_markdown'].split(' ')[:7] # First 7 lines for line in preview_lines: lines.append(line) if len(table['table_markdown'].split(' ')) > 7: lines.append("| ... | ... | (additional rows omitted) |") lines.append(" ---") with open(output_path, 'w', encoding='utf-8') as f: f.write(' '.join(lines)) print(f"Summary saved to: {output_path}") async def main(): """Main function for command-line usage""" import argparse parser = argparse.ArgumentParser( description="Analyze document tables with LLM inference" ) parser.add_argument("pdf_path", help="Path to PDF document") parser.add_argument( "--output", default="table_analysis.json", help="Output JSON file (default: table_analysis.json)" ) parser.add_argument( "--summary", default="table_analysis_summary.md", help="Output summary file (default: table_analysis_summary.md)" ) parser.add_argument( "--model", default="claude-3-haiku-20240307", help="Anthropic model to use (default: claude-3-haiku)" ) args = parser.parse_args() # Check for API key if not os.getenv("ANTHROPIC_API_KEY"): print("Error: ANTHROPIC_API_KEY not found in environment") print("Please set it in .env file or export it") return 1 # Process document processor = TableInferenceProcessor(model=args.model) results = await processor.process_document(args.pdf_path) # Save outputs processor.save_json(results, args.output) processor.save_summary(results, args.summary) print(" Analysis complete!") print(f"- JSON: {args.output}") print(f"- Summary: {args.summary}") return 0 if __name__ == "__main__": exit(asyncio.run(main()))