Coverage for src / quber / processors / table_inference.py: 66%
196 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""
2Table Inference with LLM
4Uses PydanticAI with various model providers to analyze tables and generate
5meaningful titles and descriptions based on document context.
6"""
8import asyncio
9import json
10import os
11from datetime import datetime
12from pathlib import Path
13from typing import TYPE_CHECKING, Any, Dict, List, Optional
15import logfire
16from loguru import logger
17from pydantic import BaseModel, Field
19if TYPE_CHECKING:
20 from docling_core.types.doc.document import DoclingDocument, NodeItem
22from quber.agents import AgentFactory
23from quber.agents.factory import ModelProvider
24from quber.core.consolidation import HeaderConsolidator
25from quber.core.models import TableMetadata
26from quber.core.parsers import parser_for_preset
27from quber.prompts import load_prompt
28from quber.settings import Settings, get_settings
29from quber.utils import setup_logging
32def provenance_page(item: "NodeItem") -> Optional[int]:
33 """The 1-indexed page an item sits on, taken from its docling provenance.
35 docling records an item's location in `item.prov` (a list of provenance
36 spans, each with a `page_no`); the page lives there, not on a `_page`
37 attribute. An item may carry several spans when it straddles a page break;
38 the first span's page is used. Structural group nodes have no provenance
39 and return None. `prov` is read via getattr because the base node type does
40 not declare it (only laid-out items do)."""
41 prov = getattr(item, "prov", None)
42 if not prov:
43 return None
44 return prov[0].page_no
47class TableContext(BaseModel):
48 """Input context for table analysis."""
50 executive_summary: str = Field(description="3-4 sentence document-level summary")
51 procedural_title: str = Field(description="Title generated by procedural rules")
52 headers: List[str] = Field(description="Individual headers found")
53 descriptive_text: Optional[str] = Field(default=None, description="Text that describes the table")
54 table_preview: str = Field(description="First few rows of the table")
55 preceding_text: Optional[str] = Field(default=None, description="Text before the table")
56 page_context: Optional[str] = Field(
57 default=None, description="Extended page-level context including headers"
58 )
59 page_number: int = Field(description="Page where table appears")
62class TableInfo(BaseModel):
63 """LLM-inferred information about a table."""
65 title: str = Field(description="Clear, concise title for the table")
66 description: str = Field(
67 description="Functional description of what the table contains and how it might be used"
68 )
71class ExecutiveSummary(BaseModel):
72 """Executive summary of a document."""
74 summary: str = Field(description="3-4 sentence executive summary of the document")
77class TableInferenceProcessor:
78 """Process tables with LLM inference for enhanced metadata."""
80 def __init__(
81 self,
82 provider: Optional[ModelProvider] = None,
83 model: Optional[str] = None,
84 settings: Optional[Settings] = None,
85 preset: str = "tuned-financial",
86 ):
87 """
88 Initialize the processor with specified configuration.
90 Args:
91 provider: Model provider to use (auto-detected if not specified)
92 model: Specific model to use (uses default if not specified)
93 settings: Settings object (uses the cached get_settings() if not
94 provided). The optional-injection seam lets callers/tests pass
95 an overridden Settings (e.g. CLI flags).
96 preset: DoclingParser preset to use ("tuned-financial" canonical
97 default, or "legacy" for the pre-QUE-218 configuration).
98 """
99 self.settings = settings or get_settings()
100 self.provider = provider
101 self.model = model
102 self.preset = preset
104 # Set up logging
105 setup_logging(
106 service_name="table-inference",
107 enable_logfire=self.settings.obs.enable_logfire,
108 log_level=self.settings.log_level,
109 )
111 # Load system prompt from TOML file
112 system_prompt = load_prompt("table_inference", "system")
114 # Create agent factory
115 self.agent_factory = AgentFactory(enable_logfire=self.settings.obs.enable_logfire)
117 # Create the agent
118 self.agent = self.agent_factory.create_agent(
119 output_type=TableInfo,
120 system_prompt=system_prompt,
121 provider=provider,
122 model=model,
123 )
125 logger.debug(
126 f"Initialized TableInferenceProcessor with provider: {provider or 'auto'}, model: {model or 'default'}"
127 )
129 async def analyze_table(
130 self,
131 metadata: TableMetadata,
132 page_context: Optional[str] = None,
133 executive_summary: Optional[str] = None,
134 ) -> Dict[str, Any]:
135 """
136 Analyze a single table with LLM inference.
138 Args:
139 metadata: Table metadata from enhanced parser
140 page_context: Extended context from the page
141 executive_summary: Document-level executive summary
143 Returns:
144 Dictionary with both procedural and LLM-inferred information
145 """
146 # Prepare context for LLM
147 context = TableContext(
148 executive_summary=executive_summary or "No document summary available",
149 procedural_title=metadata.consolidated_header or "Untitled Table",
150 headers=metadata.headers or [],
151 descriptive_text=metadata.descriptive_text,
152 table_preview=metadata.table_markdown or "", # Send full table for accurate analysis
153 preceding_text=metadata.preceding_text,
154 page_context=page_context,
155 page_number=metadata.page_number,
156 )
158 # Get LLM inference with telemetry
159 try:
160 prompt = self.format_prompt(context)
162 if self.settings.obs.enable_logfire:
163 with logfire.span("analyze_table", table_id=metadata.table_index, page=metadata.page_number):
164 result = await self.agent.run(prompt)
165 llm_info = result.output
166 else:
167 result = await self.agent.run(prompt)
168 llm_info = result.output
170 logger.debug(f"Successfully analyzed table {metadata.table_index}")
172 except Exception as e:
173 logger.error(f"LLM inference failed for table {metadata.table_index}: {e}")
174 llm_info = TableInfo(
175 title=metadata.consolidated_header or f"Table {metadata.table_index + 1}",
176 description="Unable to generate description",
177 )
179 # Combine all information
180 return {
181 "table_id": metadata.table_index,
182 "page": metadata.page_number,
183 "procedural_title": metadata.consolidated_header,
184 "headers": metadata.headers,
185 "descriptive_text": metadata.descriptive_text,
186 "preceding_text": metadata.preceding_text,
187 "llm_title": llm_info.title,
188 "llm_description": llm_info.description,
189 "table_markdown": metadata.table_markdown,
190 "metadata": {
191 "rows": metadata.rows,
192 "cols": metadata.cols,
193 "first_cell": metadata.first_cell_content,
194 },
195 }
197 async def analyze_tables_batch(
198 self,
199 metadata_list: List[TableMetadata],
200 page_contexts: Optional[Dict[int, str]] = None,
201 executive_summary: Optional[str] = None,
202 max_concurrent: Optional[int] = None,
203 ) -> List[Dict[str, Any]]:
204 """
205 Analyze multiple tables concurrently with rate limiting.
207 Args:
208 metadata_list: List of table metadata to analyze
209 page_contexts: Optional dict mapping table indices to page contexts
210 executive_summary: Document-level executive summary
211 max_concurrent: Maximum concurrent analyses (uses config default if None)
213 Returns:
214 List of analysis results
215 """
216 max_concurrent = max_concurrent or self.settings.max_concurrent_tables
217 page_contexts = page_contexts or {}
219 # Create semaphore for rate limiting
220 semaphore = asyncio.Semaphore(max_concurrent)
222 async def analyze_with_limit(metadata: TableMetadata):
223 async with semaphore:
224 page_context = page_contexts.get(metadata.table_index)
225 return await self.analyze_table(metadata, page_context, executive_summary)
227 # Process all tables concurrently with rate limiting
228 tasks = [analyze_with_limit(metadata) for metadata in metadata_list]
230 with logfire.span(
231 "analyze_tables_batch", total_tables=len(metadata_list), max_concurrent=max_concurrent
232 ):
233 results = await asyncio.gather(*tasks)
235 return results
237 def format_prompt(self, context: TableContext) -> str:
238 """Format the context as a prompt for the agent."""
239 return f"""
240Analyze this table from page {context.page_number}:
242=== DOCUMENT EXECUTIVE SUMMARY ===
243{context.executive_summary}
245=== EXTENDED PAGE CONTEXT ===
246{context.page_context or "No extended page context available"}
248=== IMMEDIATE TABLE CONTEXT ===
249Procedural Title: {context.procedural_title}
250Headers Found: {", ".join(context.headers) if context.headers else "None"}
252Preceding Text (chronological order, furthest to closest):
253{context.preceding_text or "None"}
255Note: The LAST element in preceding text may be a footnote from the previous table.
256[Table] markers indicate preceding tables.
258=== TABLE DATA (may be truncated at page boundary) ===
259{context.table_preview}
261Based on ALL context levels (document summary, page context, procedural headers, and table data),
262extract the most appropriate title and provide a functional description for what this COMPLETE table represents.
263"""
265 def get_table_preview(self, markdown: Optional[str], max_rows: int = 10) -> str:
266 """Get preview of table (first N rows)."""
267 if not markdown:
268 return "No table content available"
270 lines = markdown.strip().split("\n")
271 preview_lines = lines[: min(max_rows + 2, len(lines))] # +2 for header and separator
272 return "\n".join(preview_lines)
274 def extract_page_context(
275 self, document: "DoclingDocument", page_number: int, char_limit: int = 2000
276 ) -> str:
277 """
278 Extract extended context from the page containing the table.
280 Args:
281 document: The Docling document object
282 page_number: The page number (1-indexed) where the table appears
283 char_limit: Maximum characters to extract for context
285 Returns:
286 String containing page-level context including headers and surrounding text
287 """
288 context_parts: List[str] = []
289 char_count = 0
291 # Text on the table's own page. iterate_items() yields (item, level)
292 # tuples; the page of each item comes from its provenance, and headers
293 # carry a level. Text and level are read with getattr because only some
294 # item types (text, headers) have them.
295 current_page_text: List[str] = []
296 for item, _level in document.iterate_items():
297 if provenance_page(item) != page_number:
298 continue
299 text = (getattr(item, "text", "") or "").strip()
300 if not text:
301 continue
302 level = getattr(item, "level", None)
303 if level is not None:
304 current_page_text.append(f"[Header L{level}] {text}")
305 else:
306 current_page_text.append(text)
307 char_count += len(text)
308 if char_count >= char_limit:
309 break
311 # If we have room and the table is not on the first page, add a tail of
312 # the previous page for lead-in context.
313 if char_count < char_limit * 0.7 and page_number > 1: # Use 70% to leave room
314 prev_page = page_number - 1
315 prev_page_text: List[str] = []
316 for item, _level in document.iterate_items():
317 if provenance_page(item) != prev_page:
318 continue
319 text = (getattr(item, "text", "") or "").strip()
320 if text:
321 prev_page_text.append(text)
323 # Take last ~500 chars from previous page
324 if prev_page_text:
325 prev_text = "\n".join(prev_page_text[-5:]) # Last 5 text elements
326 if len(prev_text) > 500:
327 prev_text = "..." + prev_text[-500:]
328 context_parts.append("=== Previous Page Context ===\n" + prev_text)
330 # Add current page context
331 if current_page_text:
332 context_parts.append(
333 "=== Current Page Content ===\n" + "\n".join(current_page_text[:20])
334 ) # First 20 elements
336 return "\n\n".join(context_parts) if context_parts else "No page context available"
338 async def generate_executive_summary(self, document: "DoclingDocument") -> str:
339 """
340 Generate a 3-4 sentence executive summary of the entire document.
342 Args:
343 document: The Docling document object
345 Returns:
346 Executive summary string
347 """
348 # Export full document to markdown
349 document_markdown = document.export_to_markdown()
351 # Create a specialized agent for executive summary generation
352 summary_agent = self.agent_factory.create_agent(
353 output_type=ExecutiveSummary,
354 system_prompt="You are analyzing business documents to create executive summaries.",
355 provider=self.provider,
356 model=self.model,
357 )
359 # Create prompt for executive summary
360 prompt = f"""
361Analyze this complete document and provide a 3-4 sentence executive summary.
363The summary should describe:
3641. What this document is (type of document, purpose)
3652. Who it is about (company, organization, or subject)
3663. What are the major points being made (key metrics, themes, or findings)
368=== DOCUMENT ===
369{document_markdown}
371Provide a concise executive summary that captures the essential nature and content of this document.
372"""
374 try:
375 if self.settings.obs.enable_logfire:
376 with logfire.span("generate_executive_summary"):
377 result = await summary_agent.run(prompt)
378 return result.output.summary
379 else:
380 result = await summary_agent.run(prompt)
381 return result.output.summary
383 except Exception as e:
384 logger.error(f"Failed to generate executive summary: {e}")
385 return "Unable to generate document summary"
387 async def process_document(self, pdf_path: "str | os.PathLike[str]") -> Dict[str, Any]:
388 """
389 Process all tables in a document.
391 Args:
392 pdf_path: Path to PDF document. Accepts a local path string or any
393 os.PathLike (e.g. a cloudpathlib S3Path); Path(pdf_path)
394 materializes the latter to a local cached file before parsing.
396 Returns:
397 Dictionary with document and table information
398 """
399 # str() keeps the label cheap for S3Path (its str is the s3:// URI, not
400 # a download); os.path.basename(S3Path) would force materialization.
401 doc_name = os.path.basename(str(pdf_path))
403 # Log start of document processing
404 logger.info(f"Starting document processing: {doc_name}")
406 # Parse into the canonical DoclingDocument, then consolidate headers
407 # as a separate post-processing step (plan §4).
408 parser = parser_for_preset(self.preset)
409 consolidator = HeaderConsolidator(separator=" ^ ", enable_consolidation=True, debug=False)
411 with logfire.span("parse_document", document=str(pdf_path)):
412 document = parser.parse(Path(pdf_path)).document
413 table_metadata = consolidator.consolidate(document)
415 # Log completion of resource-intensive document parsing
416 logger.info(f"Document parsed successfully: {doc_name}")
418 # Generate executive summary of the document
419 logger.info(f"Generating executive summary for {doc_name}")
420 executive_summary = await self.generate_executive_summary(document)
421 logger.debug(f"Executive summary: {executive_summary[:100]}...")
423 # Log table analysis details
424 if len(table_metadata) > 0:
425 model_info = self.model or "default"
426 if self.provider:
427 model_info = f"{self.provider.value}:{model_info}"
428 logger.info(f"Analyzing {len(table_metadata)} tables using {model_info}")
430 # Extract page contexts for all tables
431 page_contexts = {}
432 for metadata in table_metadata:
433 page_contexts[metadata.table_index] = self.extract_page_context(document, metadata.page_number)
435 # Process all tables concurrently
436 tables = await self.analyze_tables_batch(table_metadata, page_contexts, executive_summary)
438 if len(tables) > 0:
439 logger.info(f"Analysis complete: {len(tables)} tables processed")
441 # Build final output
442 return {
443 "document": doc_name,
444 "extraction_date": datetime.now().isoformat(),
445 "total_pages": len(document.pages),
446 "total_tables": len(tables),
447 "model_provider": self.provider.value if self.provider else "auto-detected",
448 "model": self.model or "default",
449 "tables": tables,
450 }
452 def save_json(self, data: Dict[str, Any], output_path: str):
453 """Save results as JSON."""
454 with open(output_path, "w", encoding="utf-8") as f:
455 json.dump(data, f, indent=2, ensure_ascii=False)
456 logger.debug(f"JSON results saved to: {output_path}")
458 def save_summary(self, data: Dict[str, Any], output_path: str):
459 """Save human-readable summary."""
460 lines = []
461 lines.append("# Table Analysis Report")
462 lines.append("")
463 lines.append(f"**Document:** {data['document']}")
464 lines.append("")
465 lines.append(f"**Date:** {data['extraction_date']}")
466 lines.append("")
467 lines.append(f"**Model:** {data.get('model_provider', 'unknown')}/{data.get('model', 'unknown')}")
468 lines.append("")
469 lines.append(f"**Total Tables:** {data['total_tables']}")
470 lines.append("")
472 for table in data["tables"]:
473 lines.append(f"\n## Table {table['table_id'] + 1} (Page {table['page']})")
475 # Show procedural information
476 lines.append("\n### Procedural Extraction:")
477 lines.append(f"- **Consolidated Headers:** {table['procedural_title'] or 'None'}")
478 if table.get("headers"):
479 lines.append(f"- **Individual Headers:** {', '.join(table['headers'])}")
480 if table.get("descriptive_text"):
481 lines.append(f"- **Descriptive Text:** {table['descriptive_text'][:200]}...")
483 # Show LLM analysis
484 lines.append("\n### LLM Analysis:")
485 lines.append(f"- **Title:** {table['llm_title']}")
486 lines.append(f"- **Description:** {table['llm_description']}")
488 # Show metadata
489 lines.append("\n### Metadata:")
490 lines.append(
491 f"- **Size:** {table['metadata']['rows']} rows × {table['metadata']['cols']} columns"
492 )
493 if table["metadata"].get("first_cell"):
494 lines.append(f"- **First Cell:** {table['metadata']['first_cell'][:50]}...")
496 # Show table preview (first few rows)
497 if table.get("table_markdown"):
498 lines.append("\n### Table Preview:")
499 preview_lines = table["table_markdown"].split("\n")[:7] # First 7 lines
500 for line in preview_lines:
501 lines.append(line)
502 if len(table["table_markdown"].split("\n")) > 7:
503 lines.append("| ... | ... | (additional rows omitted) |")
505 lines.append("\n---")
507 with open(output_path, "w", encoding="utf-8") as f:
508 f.write("\n".join(lines))
509 logger.debug(f"Summary saved to: {output_path}")