""" Document Chunker Implementation Creates semantic chunks from document content with bounding box tracking. Includes TableAwareChunker for preserving table structure in markdown format. """ import uuid import time import re from typing import List, Optional, Dict, Any, Tuple from dataclasses import dataclass from pydantic import BaseModel, Field from loguru import logger from collections import defaultdict from ..schemas.core import ( BoundingBox, DocumentChunk, ChunkType, LayoutRegion, LayoutType, OCRRegion, ) class ChunkerConfig(BaseModel): """Configuration for document chunking.""" # Chunk size limits max_chunk_chars: int = Field( default=1000, ge=100, description="Maximum characters per chunk" ) min_chunk_chars: int = Field( default=50, ge=10, description="Minimum characters per chunk" ) overlap_chars: int = Field( default=100, ge=0, description="Character overlap between chunks" ) # Chunking strategy strategy: str = Field( default="semantic", description="Chunking strategy: semantic, fixed, or layout" ) respect_layout: bool = Field( default=True, description="Respect layout region boundaries" ) merge_small_regions: bool = Field( default=True, description="Merge small adjacent regions" ) # Special element handling chunk_tables: bool = Field( default=True, description="Create separate chunks for tables" ) chunk_figures: bool = Field( default=True, description="Create separate chunks for figures" ) include_captions: bool = Field( default=True, description="Include captions with figures/tables" ) # Sentence handling split_on_sentences: bool = Field( default=True, description="Split on sentence boundaries when possible" ) # Table-aware chunking (FG-002) preserve_table_structure: bool = Field( default=True, description="Preserve table structure as markdown with structured data" ) table_row_threshold: float = Field( default=10.0, description="Y-coordinate threshold for grouping cells into rows" ) table_col_threshold: float = Field( default=20.0, description="X-coordinate threshold for grouping cells into columns" ) detect_table_headers: bool = Field( default=True, description="Attempt to detect and mark header rows" ) # Map layout types to chunk types LAYOUT_TO_CHUNK_TYPE = { LayoutType.TEXT: ChunkType.TEXT, LayoutType.TITLE: ChunkType.TITLE, LayoutType.HEADING: ChunkType.HEADING, LayoutType.PARAGRAPH: ChunkType.PARAGRAPH, LayoutType.LIST: ChunkType.LIST_ITEM, LayoutType.TABLE: ChunkType.TABLE, LayoutType.FIGURE: ChunkType.FIGURE, LayoutType.CHART: ChunkType.CHART, LayoutType.FORMULA: ChunkType.FORMULA, LayoutType.CAPTION: ChunkType.CAPTION, LayoutType.FOOTNOTE: ChunkType.FOOTNOTE, LayoutType.HEADER: ChunkType.HEADER, LayoutType.FOOTER: ChunkType.FOOTER, } class DocumentChunker: """Base class for document chunkers.""" def __init__(self, config: Optional[ChunkerConfig] = None): self.config = config or ChunkerConfig() def create_chunks( self, ocr_regions: List[OCRRegion], layout_regions: Optional[List[LayoutRegion]] = None, document_id: str = "", source_path: Optional[str] = None, ) -> List[DocumentChunk]: """ Create chunks from OCR and layout regions. Args: ocr_regions: OCR text regions layout_regions: Optional layout regions document_id: Parent document ID source_path: Source file path Returns: List of DocumentChunk """ raise NotImplementedError class SemanticChunker(DocumentChunker): """ Semantic chunker that respects document structure. Creates chunks based on: - Layout region boundaries - Semantic coherence (paragraphs, sections) - Size constraints with overlap """ def create_chunks( self, ocr_regions: List[OCRRegion], layout_regions: Optional[List[LayoutRegion]] = None, document_id: str = "", source_path: Optional[str] = None, ) -> List[DocumentChunk]: """Create semantic chunks from document content.""" if not ocr_regions: return [] start_time = time.time() chunks = [] chunk_index = 0 if layout_regions and self.config.respect_layout: # Use layout regions to guide chunking chunks = self._chunk_by_layout( ocr_regions, layout_regions, document_id, source_path ) else: # Fall back to text-based chunking chunks = self._chunk_by_text( ocr_regions, document_id, source_path ) # Assign sequence indices for i, chunk in enumerate(chunks): chunk.sequence_index = i logger.debug( f"Created {len(chunks)} chunks in " f"{(time.time() - start_time) * 1000:.1f}ms" ) return chunks def _chunk_by_layout( self, ocr_regions: List[OCRRegion], layout_regions: List[LayoutRegion], document_id: str, source_path: Optional[str], ) -> List[DocumentChunk]: """Create chunks based on layout regions.""" chunks = [] # Sort layout regions by reading order sorted_layouts = sorted( layout_regions, key=lambda r: (r.reading_order or 0, r.bbox.y_min, r.bbox.x_min) ) for layout in sorted_layouts: # Get OCR regions within this layout region contained_ocr = self._get_contained_ocr(ocr_regions, layout) if not contained_ocr: continue # Determine chunk type chunk_type = LAYOUT_TO_CHUNK_TYPE.get(layout.type, ChunkType.TEXT) # Handle special types differently if layout.type == LayoutType.TABLE and self.config.chunk_tables: chunk = self._create_table_chunk( contained_ocr, layout, document_id, source_path ) chunks.append(chunk) elif layout.type in (LayoutType.FIGURE, LayoutType.CHART) and self.config.chunk_figures: chunk = self._create_figure_chunk( contained_ocr, layout, document_id, source_path ) chunks.append(chunk) else: # Regular text chunk - may need splitting text_chunks = self._create_text_chunks( contained_ocr, layout, chunk_type, document_id, source_path ) chunks.extend(text_chunks) return chunks def _chunk_by_text( self, ocr_regions: List[OCRRegion], document_id: str, source_path: Optional[str], ) -> List[DocumentChunk]: """Create chunks from text without layout guidance.""" chunks = [] # Sort by reading order (y then x) sorted_regions = sorted( ocr_regions, key=lambda r: (r.page, r.bbox.y_min, r.bbox.x_min) ) # Group by page pages: Dict[int, List[OCRRegion]] = {} for r in sorted_regions: if r.page not in pages: pages[r.page] = [] pages[r.page].append(r) # Process each page for page_num in sorted(pages.keys()): page_regions = pages[page_num] page_chunks = self._split_text_regions( page_regions, document_id, source_path, page_num ) chunks.extend(page_chunks) return chunks def _get_contained_ocr( self, ocr_regions: List[OCRRegion], layout: LayoutRegion, ) -> List[OCRRegion]: """Get OCR regions contained within a layout region.""" contained = [] for ocr in ocr_regions: if ocr.page == layout.page: # Check if OCR region overlaps significantly with layout iou = layout.bbox.iou(ocr.bbox) if iou > 0.3 or layout.bbox.contains(ocr.bbox): contained.append(ocr) return contained def _create_text_chunks( self, ocr_regions: List[OCRRegion], layout: LayoutRegion, chunk_type: ChunkType, document_id: str, source_path: Optional[str], ) -> List[DocumentChunk]: """Create text chunks from OCR regions, splitting if needed.""" chunks = [] # Combine text text = " ".join(r.text for r in ocr_regions) # Calculate average confidence avg_conf = sum(r.confidence for r in ocr_regions) / len(ocr_regions) # Check if splitting is needed if len(text) <= self.config.max_chunk_chars: # Single chunk chunk = DocumentChunk( chunk_id=f"{document_id}_{uuid.uuid4().hex[:8]}", chunk_type=chunk_type, text=text, bbox=layout.bbox, page=layout.page, document_id=document_id, source_path=source_path, sequence_index=0, confidence=avg_conf, ) chunks.append(chunk) else: # Split into multiple chunks split_chunks = self._split_text( text, layout.bbox, layout.page, chunk_type, document_id, source_path, avg_conf ) chunks.extend(split_chunks) return chunks def _split_text( self, text: str, bbox: BoundingBox, page: int, chunk_type: ChunkType, document_id: str, source_path: Optional[str], confidence: float, ) -> List[DocumentChunk]: """Split long text into multiple chunks with overlap.""" chunks = [] max_chars = self.config.max_chunk_chars overlap = self.config.overlap_chars # Split on sentences if enabled if self.config.split_on_sentences: sentences = self._split_sentences(text) else: sentences = [text] current_text = "" for sentence in sentences: if len(current_text) + len(sentence) > max_chars and current_text: # Create chunk chunk = DocumentChunk( chunk_id=f"{document_id}_{uuid.uuid4().hex[:8]}", chunk_type=chunk_type, text=current_text.strip(), bbox=bbox, page=page, document_id=document_id, source_path=source_path, sequence_index=len(chunks), confidence=confidence, ) chunks.append(chunk) # Start new chunk with overlap if overlap > 0: overlap_text = current_text[-overlap:] if len(current_text) > overlap else current_text current_text = overlap_text + " " + sentence else: current_text = sentence else: current_text += " " + sentence if current_text else sentence # Don't forget the last chunk if current_text.strip(): chunk = DocumentChunk( chunk_id=f"{document_id}_{uuid.uuid4().hex[:8]}", chunk_type=chunk_type, text=current_text.strip(), bbox=bbox, page=page, document_id=document_id, source_path=source_path, sequence_index=len(chunks), confidence=confidence, ) chunks.append(chunk) return chunks def _split_sentences(self, text: str) -> List[str]: """Split text into sentences.""" # Simple sentence splitting import re sentences = re.split(r'(?<=[.!?])\s+', text) return [s.strip() for s in sentences if s.strip()] def _create_table_chunk( self, ocr_regions: List[OCRRegion], layout: LayoutRegion, document_id: str, source_path: Optional[str], ) -> DocumentChunk: """ Create a chunk for table content with structure preservation. Enhanced table handling (FG-002): - Reconstructs table structure from OCR regions - Generates markdown table representation - Stores structured data for SQL-like queries - Detects and marks header rows """ if not ocr_regions: return DocumentChunk( chunk_id=f"{document_id}_table_{uuid.uuid4().hex[:8]}", chunk_type=ChunkType.TABLE, text="[Empty Table]", bbox=layout.bbox, page=layout.page, document_id=document_id, source_path=source_path, sequence_index=0, confidence=0.0, extra=layout.extra or {}, ) avg_conf = sum(r.confidence for r in ocr_regions) / len(ocr_regions) # Check if we should preserve table structure if not self.config.preserve_table_structure: # Fall back to simple pipe-separated format text = " | ".join(r.text for r in ocr_regions) return DocumentChunk( chunk_id=f"{document_id}_table_{uuid.uuid4().hex[:8]}", chunk_type=ChunkType.TABLE, text=text, bbox=layout.bbox, page=layout.page, document_id=document_id, source_path=source_path, sequence_index=0, confidence=avg_conf, extra=layout.extra or {}, ) # Reconstruct table structure from spatial positions table_data = self._reconstruct_table_structure(ocr_regions) # Generate markdown representation markdown_table = self._table_to_markdown( table_data["rows"], table_data["headers"], table_data["has_header"] ) # Create rich metadata for structured queries table_extra = { **(layout.extra or {}), "table_structure": { "row_count": table_data["row_count"], "col_count": table_data["col_count"], "has_header": table_data["has_header"], "headers": table_data["headers"], "cells": table_data["cells"], # 2D list of cell values "cell_positions": table_data["cell_positions"], # For highlighting }, "format": "markdown", "searchable_text": table_data["searchable_text"], } return DocumentChunk( chunk_id=f"{document_id}_table_{uuid.uuid4().hex[:8]}", chunk_type=ChunkType.TABLE, text=markdown_table, bbox=layout.bbox, page=layout.page, document_id=document_id, source_path=source_path, sequence_index=0, confidence=avg_conf, extra=table_extra, ) def _reconstruct_table_structure( self, ocr_regions: List[OCRRegion], ) -> Dict[str, Any]: """ Reconstruct table structure from OCR regions based on spatial positions. Groups OCR regions into rows and columns by analyzing their bounding boxes. Returns structured table data for markdown generation and queries. """ if not ocr_regions: return { "rows": [], "headers": [], "has_header": False, "row_count": 0, "col_count": 0, "cells": [], "cell_positions": [], "searchable_text": "", } # Sort regions by vertical position (y_min) then horizontal (x_min) sorted_regions = sorted( ocr_regions, key=lambda r: (r.bbox.y_min, r.bbox.x_min) ) # Group into rows based on y-coordinate proximity row_threshold = self.config.table_row_threshold rows: List[List[OCRRegion]] = [] current_row: List[OCRRegion] = [] current_y = None for region in sorted_regions: if current_y is None: current_y = region.bbox.y_min current_row.append(region) elif abs(region.bbox.y_min - current_y) <= row_threshold: current_row.append(region) else: if current_row: # Sort row by x position current_row.sort(key=lambda r: r.bbox.x_min) rows.append(current_row) current_row = [region] current_y = region.bbox.y_min # Don't forget the last row if current_row: current_row.sort(key=lambda r: r.bbox.x_min) rows.append(current_row) # Determine column structure # Find consistent column boundaries across all rows col_positions = self._detect_column_positions(rows) num_cols = len(col_positions) if col_positions else max(len(row) for row in rows) # Build structured cell data cells: List[List[str]] = [] cell_positions: List[List[Dict[str, Any]]] = [] for row in rows: row_cells = self._assign_cells_to_columns(row, col_positions, num_cols) cells.append([cell["text"] for cell in row_cells]) cell_positions.append([{ "text": cell["text"], "bbox": cell["bbox"], "confidence": cell["confidence"] } for cell in row_cells]) # Detect header row has_header = False headers: List[str] = [] if self.config.detect_table_headers and len(cells) > 0: has_header, headers = self._detect_header_row(cells, rows) # Build searchable text (for vector embedding) searchable_parts = [] for i, row in enumerate(cells): if has_header and i == 0: searchable_parts.append("Headers: " + ", ".join(row)) else: if has_header and headers: # Include header context for each value for j, cell in enumerate(row): if j < len(headers) and headers[j]: searchable_parts.append(f"{headers[j]}: {cell}") else: searchable_parts.append(cell) else: searchable_parts.extend(row) return { "rows": cells, "headers": headers, "has_header": has_header, "row_count": len(cells), "col_count": num_cols, "cells": cells, "cell_positions": cell_positions, "searchable_text": " | ".join(searchable_parts), } def _detect_column_positions( self, rows: List[List[OCRRegion]], ) -> List[Tuple[float, float]]: """ Detect consistent column boundaries from table rows. Returns list of (x_start, x_end) tuples for each column. """ if not rows: return [] col_threshold = self.config.table_col_threshold # Collect all x positions all_x_starts = [] for row in rows: for region in row: all_x_starts.append(region.bbox.x_min) if not all_x_starts: return [] # Cluster x positions into columns all_x_starts.sort() columns = [] current_col_start = all_x_starts[0] current_col_regions = [all_x_starts[0]] for x in all_x_starts[1:]: if x - current_col_regions[-1] <= col_threshold: current_col_regions.append(x) else: # Calculate column boundary col_center = sum(current_col_regions) / len(current_col_regions) columns.append(col_center) current_col_regions = [x] # Last column if current_col_regions: col_center = sum(current_col_regions) / len(current_col_regions) columns.append(col_center) # Convert to column ranges col_ranges = [] for i, col_x in enumerate(columns): x_start = col_x - col_threshold if i < len(columns) - 1: x_end = (col_x + columns[i + 1]) / 2 else: x_end = col_x + col_threshold * 3 # Extend last column col_ranges.append((x_start, x_end)) return col_ranges def _assign_cells_to_columns( self, row_regions: List[OCRRegion], col_positions: List[Tuple[float, float]], num_cols: int, ) -> List[Dict[str, Any]]: """ Assign OCR regions in a row to their respective columns. Handles merged cells and missing cells. """ # Initialize empty cells for each column row_cells = [ {"text": "", "bbox": None, "confidence": 0.0} for _ in range(num_cols) ] if not col_positions: # No column positions detected, just use order for i, region in enumerate(row_regions): if i < num_cols: row_cells[i] = { "text": region.text.strip(), "bbox": region.bbox.to_xyxy(), "confidence": region.confidence, } return row_cells # Assign regions to columns based on x position for region in row_regions: region_x = region.bbox.x_min assigned = False for col_idx, (x_start, x_end) in enumerate(col_positions): if x_start <= region_x <= x_end: # Append to existing cell (handle multi-line cells) if row_cells[col_idx]["text"]: row_cells[col_idx]["text"] += " " + region.text.strip() else: row_cells[col_idx]["text"] = region.text.strip() row_cells[col_idx]["bbox"] = region.bbox.to_xyxy() row_cells[col_idx]["confidence"] = max( row_cells[col_idx]["confidence"], region.confidence ) assigned = True break # If not assigned, put in nearest column if not assigned: min_dist = float("inf") nearest_col = 0 for col_idx, (x_start, x_end) in enumerate(col_positions): col_center = (x_start + x_end) / 2 dist = abs(region_x - col_center) if dist < min_dist: min_dist = dist nearest_col = col_idx if row_cells[nearest_col]["text"]: row_cells[nearest_col]["text"] += " " + region.text.strip() else: row_cells[nearest_col]["text"] = region.text.strip() row_cells[nearest_col]["bbox"] = region.bbox.to_xyxy() row_cells[nearest_col]["confidence"] = region.confidence return row_cells def _detect_header_row( self, cells: List[List[str]], rows: List[List[OCRRegion]], ) -> Tuple[bool, List[str]]: """ Detect if the first row is a header row. Heuristics used: - First row contains non-numeric text - First row text is shorter (labels vs data) - First row has distinct formatting (if available) """ if not cells or len(cells) < 2: return False, [] first_row = cells[0] other_rows = cells[1:] # Check if first row is mostly non-numeric first_row_numeric_count = sum( 1 for cell in first_row if cell and self._is_numeric(cell) ) first_row_text_ratio = (len(first_row) - first_row_numeric_count) / max(len(first_row), 1) # Check if other rows are more numeric other_numeric_ratios = [] for row in other_rows: if row: numeric_count = sum(1 for cell in row if cell and self._is_numeric(cell)) other_numeric_ratios.append(numeric_count / max(len(row), 1)) avg_other_numeric = sum(other_numeric_ratios) / max(len(other_numeric_ratios), 1) # Header detection: first row is text-heavy, others are more numeric is_header = ( first_row_text_ratio > 0.5 and (avg_other_numeric > first_row_text_ratio * 0.5 or first_row_text_ratio > 0.8) ) # Also consider: shorter cell lengths in first row (labels are usually shorter) first_row_avg_len = sum(len(cell) for cell in first_row) / max(len(first_row), 1) other_avg_lens = [ sum(len(cell) for cell in row) / max(len(row), 1) for row in other_rows ] avg_other_len = sum(other_avg_lens) / max(len(other_avg_lens), 1) if first_row_avg_len < avg_other_len * 0.8: is_header = True return is_header, first_row if is_header else [] def _is_numeric(self, text: str) -> bool: """Check if text is primarily numeric (including currency, percentages).""" cleaned = re.sub(r'[$€£¥%,.\s\-+()]', '', text) return cleaned.isdigit() if cleaned else False def _table_to_markdown( self, rows: List[List[str]], headers: List[str], has_header: bool, ) -> str: """ Convert table data to markdown format. Creates a properly formatted markdown table with: - Header row (if detected) - Separator row - Data rows """ if not rows: return "[Empty Table]" # Determine column count num_cols = max(len(row) for row in rows) if rows else 0 if num_cols == 0: return "[Empty Table]" # Normalize all rows to same column count normalized_rows = [] for row in rows: normalized = row + [""] * (num_cols - len(row)) normalized_rows.append(normalized) # Build markdown lines md_lines = [] if has_header and headers: # Use detected headers header_line = "| " + " | ".join(headers + [""] * (num_cols - len(headers))) + " |" separator = "| " + " | ".join(["---"] * num_cols) + " |" md_lines.append(header_line) md_lines.append(separator) data_rows = normalized_rows[1:] else: # No header - create generic headers generic_headers = [f"Col{i+1}" for i in range(num_cols)] header_line = "| " + " | ".join(generic_headers) + " |" separator = "| " + " | ".join(["---"] * num_cols) + " |" md_lines.append(header_line) md_lines.append(separator) data_rows = normalized_rows # Add data rows for row in data_rows: # Escape pipe characters in cell content escaped_row = [cell.replace("|", "\\|") for cell in row] row_line = "| " + " | ".join(escaped_row) + " |" md_lines.append(row_line) return "\n".join(md_lines) def _create_figure_chunk( self, ocr_regions: List[OCRRegion], layout: LayoutRegion, document_id: str, source_path: Optional[str], ) -> DocumentChunk: """Create a chunk for figure/chart content.""" # For figures, text is usually caption text = " ".join(r.text for r in ocr_regions) if ocr_regions else "[Figure]" avg_conf = sum(r.confidence for r in ocr_regions) / len(ocr_regions) if ocr_regions else 0.5 chunk_type = ChunkType.CHART if layout.type == LayoutType.CHART else ChunkType.FIGURE return DocumentChunk( chunk_id=f"{document_id}_{chunk_type.value}_{uuid.uuid4().hex[:8]}", chunk_type=chunk_type, text=text, bbox=layout.bbox, page=layout.page, document_id=document_id, source_path=source_path, sequence_index=0, confidence=avg_conf, caption=text if ocr_regions else None, ) def _split_text_regions( self, ocr_regions: List[OCRRegion], document_id: str, source_path: Optional[str], page_num: int, ) -> List[DocumentChunk]: """Split OCR regions into chunks without layout guidance.""" if not ocr_regions: return [] chunks = [] current_text = "" current_regions = [] for region in ocr_regions: if len(current_text) + len(region.text) > self.config.max_chunk_chars: if current_regions: # Create chunk from accumulated regions chunk = self._create_chunk_from_regions( current_regions, document_id, source_path, page_num, len(chunks) ) chunks.append(chunk) current_text = region.text current_regions = [region] else: current_text += " " + region.text current_regions.append(region) # Final chunk if current_regions: chunk = self._create_chunk_from_regions( current_regions, document_id, source_path, page_num, len(chunks) ) chunks.append(chunk) return chunks def _create_chunk_from_regions( self, regions: List[OCRRegion], document_id: str, source_path: Optional[str], page_num: int, sequence_index: int, ) -> DocumentChunk: """Create a chunk from a list of OCR regions.""" text = " ".join(r.text for r in regions) avg_conf = sum(r.confidence for r in regions) / len(regions) # Compute bounding box x_min = min(r.bbox.x_min for r in regions) y_min = min(r.bbox.y_min for r in regions) x_max = max(r.bbox.x_max for r in regions) y_max = max(r.bbox.y_max for r in regions) bbox = BoundingBox( x_min=x_min, y_min=y_min, x_max=x_max, y_max=y_max, normalized=False, ) return DocumentChunk( chunk_id=f"{document_id}_{uuid.uuid4().hex[:8]}", chunk_type=ChunkType.TEXT, text=text, bbox=bbox, page=page_num, document_id=document_id, source_path=source_path, sequence_index=sequence_index, confidence=avg_conf, ) # Factory _document_chunker: Optional[DocumentChunker] = None def get_document_chunker( config: Optional[ChunkerConfig] = None, ) -> DocumentChunker: """Get or create singleton document chunker.""" global _document_chunker if _document_chunker is None: config = config or ChunkerConfig() _document_chunker = SemanticChunker(config) return _document_chunker