""" Record Builder for creating structured JSON records from tables with headers. This module maps hierarchical headers to cell values and generates JSON records with source metadata for tracking. """ from typing import Any from uuid import UUID import structlog from src.extraction_v2.html_table_extractor import HtmlTable from src.extraction_v2.llm_header_detector import HeaderStructure from src.llm.factory import get_token_counter logger = structlog.get_logger(__name__) # Default max tokens for embedding (8191 is the common 8K-context ceiling). DEFAULT_MAX_TOKENS = 8000 # Minimum tokens threshold - tables smaller than this may be merged with consecutive tables MIN_TOKENS_THRESHOLD = 100 class RecordBuilder: """Builds structured records from HTML tables with detected headers.""" def __init__(self): """Initialize the record builder.""" pass def build_records( self, table: HtmlTable, headers: HeaderStructure, document_id: UUID, sheet_name: str | None = None, color_map: dict | None = None ) -> list[dict[str, Any]]: """ Convert table rows to JSON records. Args: table: HtmlTable with extracted table data headers: HeaderStructure with detected headers document_id: Document UUID for source tracking sheet_name: Optional sheet name (overrides table.sheet_name) Returns: List of record dictionaries with hierarchical header keys """ try: logger.info( f"Building records from table: {table.row_count} rows, " f"{table.col_count} cols" ) # Use provided sheet_name or fall back to table's sheet_name final_sheet_name = sheet_name or table.sheet_name or "Unknown" # Determine which rows are data (skip header rows) if not headers.header_rows: # No headers detected, treat all rows as data data_start_row = 0 else: # Start after the last header row data_start_row = max(headers.header_rows) + 1 logger.debug( f"Data starts at row {data_start_row} " f"(header rows: {headers.header_rows})" ) # Get flattened header strings flat_headers = headers.get_flat_headers() # Skip tables with generic column-based headers (table fragments without real headers) if self._has_only_column_based_headers(flat_headers): logger.info( "Skipping table with only generic column headers " "(table fragment without proper headers)" ) return [] # Build records from data rows records = [] for row_idx in range(data_start_row, len(table.rows)): record = self._build_record( row_data=table.rows[row_idx], flat_headers=flat_headers, document_id=document_id, sheet_name=final_sheet_name, row_number=row_idx, table_index=table.table_index, color_map=color_map ) # Only include non-empty records if self._is_record_non_empty(record): records.append(record) logger.info( f"Built {len(records)} records " f"(skipped {table.row_count - data_start_row - len(records)} " f"empty rows)" ) return records except Exception as e: logger.error(f"Error building records: {e}") return [] def _build_record( self, row_data: list[str], flat_headers: list[str], document_id: UUID, sheet_name: str, row_number: int, table_index: int, color_map: dict | None = None ) -> dict[str, Any]: """ Build a single record from a row of data. Normalized merged cell handling: - Only stores non-empty values from the first column of merged cells - Empty strings from colspan/rowspan expansion are skipped - This creates cleaner records without redundant empty values Args: row_data: List of cell values for this row (normalized from merged cells) flat_headers: List of flattened header strings document_id: Document UUID sheet_name: Sheet name row_number: Row number (0-indexed) table_index: Table index in document Returns: Record dictionary with content and metadata """ # Create content dictionary mapping headers to values content = {} for col_idx, cell_value in enumerate(row_data): if col_idx < len(flat_headers): header_key = flat_headers[col_idx] # NORMALIZATION: Only store non-empty values # This skips empty strings from merged cell colspan/rowspan expansion # Benefits: # - Cleaner records without redundant empty columns # - Reduced storage size # - Easier to identify actual data vs merged cell artifacts if cell_value: # Skip empty strings and None content[header_key] = cell_value else: # More columns than headers (shouldn't happen, but handle it) if cell_value: # Also skip empties for overflow columns content[f"Column {col_idx}"] = cell_value # Create record with content and source metadata record = { "content": content, "_source": { "document_id": str(document_id), "sheet": sheet_name, "row": row_number, "table_id": table_index } } # Add cell colors if available (memory optimized - only this row's colors) if color_map: cell_colors = {} sheet_colors = color_map.get(sheet_name, {}) for col_idx, cell_value in enumerate(row_data): if col_idx < len(flat_headers): header_key = flat_headers[col_idx] color_key = f"{row_number},{col_idx}" if color_key in sheet_colors: # Copy color dict (not reference) to avoid shared state cell_colors[header_key] = dict(sheet_colors[color_key]) # Only add if non-empty if cell_colors: record["_source"]["cell_colors"] = cell_colors return record def build_raw_text_records( self, table: HtmlTable, document_id: UUID, sheet_name: str | None = None ) -> list[dict[str, Any]]: """ Build records from single-column raw text tables. Concatenates all row text into a single record with 'text' field. Sets row_number = -1 to distinguish from tabular data records. No LLM header detection is needed for raw text. Args: table: HtmlTable with single-column raw text content document_id: Document UUID for source tracking sheet_name: Optional sheet name (overrides table.sheet_name) Returns: List containing single record with concatenated text, or empty list if no content """ try: logger.info( f"Building raw text record from table {table.table_index}: " f"{table.row_count} rows" ) # Use provided sheet_name or fall back to table's sheet_name final_sheet_name = sheet_name or table.sheet_name or "Unknown" # Concatenate all row text (filter empty rows) text_lines = [] for row in table.rows: # Get first column text (single-column table) if row and row[0]: text_content = row[0].strip() if text_content: text_lines.append(text_content) # Skip empty tables if not text_lines: logger.debug( f"Skipping empty raw text table {table.table_index}" ) return [] # Concatenate with newlines concatenated_text = "\n".join(text_lines) # Create single record with concatenated text record = { "content": { "text": concatenated_text }, "_source": { "document_id": str(document_id), "sheet": final_sheet_name, "row": -1, # Distinguish raw text from table rows "table_id": table.table_index } } logger.info( f"Built raw text record: {len(text_lines)} lines, " f"{len(concatenated_text)} chars" ) return [record] except Exception as e: logger.error(f"Error building raw text record: {e}") return [] def build_combined_raw_text_records( self, tables: list[HtmlTable], document_id: UUID, sheet_name: str, max_tokens: int = DEFAULT_MAX_TOKENS, min_tokens: int = MIN_TOKENS_THRESHOLD, color_map: dict | None = None ) -> list[dict[str, Any]]: """ Build records from raw text tables, merging small consecutive tables. Logic: 1. Each raw text table becomes one record by default 2. If a table has < min_tokens, try to merge with next consecutive table 3. Stop merging when combined exceeds max_tokens Args: tables: List of HtmlTable objects (single-column raw text tables) document_id: Document UUID for source tracking sheet_name: Sheet name for all tables max_tokens: Maximum tokens per record (default: 8000 for embedding) min_tokens: Minimum tokens threshold for merging (default: 100) Returns: List of records, with small consecutive tables merged """ try: if not tables: return [] # Sort tables by table_index to maintain document order sorted_tables = sorted(tables, key=lambda t: t.table_index) logger.info( f"Building raw text records from {len(sorted_tables)} tables " f"on sheet '{sheet_name}' (min_tokens={min_tokens}, max_tokens={max_tokens})" ) # Provider-aware tokenizer (Azure -> tiktoken cl100k_base; vLLM -> HF tokenizer) count_tokens = get_token_counter() def extract_table_text(table: HtmlTable) -> tuple[list[str], int]: """Extract text lines and token count from a table.""" texts = [] for row in table.rows: if row and row[0]: text_content = row[0].strip() if text_content: texts.append(text_content) combined = "\n".join(texts) tokens = count_tokens(combined) if texts else 0 return texts, tokens def split_texts_by_tokens(texts: list[str], max_tok: int) -> list[list[str]]: """Split a list of text lines into chunks that fit within max_tok.""" if not texts: return [] chunks = [] current_chunk = [] current_tokens = 0 for line in texts: line_tokens = count_tokens(line) # If single line exceeds max_tok, split the line itself if line_tokens > max_tok: # Flush current chunk first if current_chunk: chunks.append(current_chunk) current_chunk = [] current_tokens = 0 # Split long line by words words = line.split() line_chunk = [] line_chunk_tokens = 0 for word in words: word_tokens = count_tokens(word + ' ') if line_chunk_tokens + word_tokens > max_tok: if line_chunk: chunks.append([' '.join(line_chunk)]) line_chunk = [word] line_chunk_tokens = word_tokens else: line_chunk.append(word) line_chunk_tokens += word_tokens if line_chunk: chunks.append([' '.join(line_chunk)]) continue # Check if adding this line would exceed limit if current_tokens + line_tokens + 1 > max_tok: # +1 for newline if current_chunk: chunks.append(current_chunk) current_chunk = [line] current_tokens = line_tokens else: current_chunk.append(line) current_tokens += line_tokens + 1 # Add remaining lines if current_chunk: chunks.append(current_chunk) return chunks def create_record(texts: list[str], table_ids: list[int]) -> dict[str, Any]: """Create a record from accumulated text.""" concatenated = "\n".join(texts) return { "content": { "text": concatenated }, "_source": { "document_id": str(document_id), "sheet": sheet_name, "row": -1, "table_ids": table_ids } } records = [] current_texts = [] current_table_ids = [] current_tokens = 0 last_table_index = None for table in sorted_tables: table_texts, table_tokens = extract_table_text(table) if not table_texts: last_table_index = table.table_index continue # If this single table exceeds max_tokens, split it into multiple records if table_tokens > max_tokens: # Flush any accumulated text first if current_texts: records.append(create_record(current_texts, current_table_ids)) logger.debug( f"Created raw text record: {current_tokens} tokens, " f"tables {current_table_ids}" ) current_texts = [] current_table_ids = [] current_tokens = 0 # Split this large table into chunks text_chunks = split_texts_by_tokens(table_texts, max_tokens) logger.info( f"Split large table {table.table_index} ({table_tokens} tokens) " f"into {len(text_chunks)} chunks" ) for chunk_idx, chunk_texts in enumerate(text_chunks): chunk_record = { "content": { "text": "\n".join(chunk_texts) }, "_source": { "document_id": str(document_id), "sheet": sheet_name, "row": -1, "table_id": table.table_index, "chunk_index": chunk_idx, "total_chunks": len(text_chunks), } } records.append(chunk_record) last_table_index = table.table_index continue # Check if this table is consecutive with the previous is_consecutive = ( last_table_index is not None and table.table_index - last_table_index == 1 ) # Determine if we should try to merge # Merge only if: # 1. Current accumulated text is small (< min_tokens), AND # 2. Tables are consecutive, AND # 3. Combined won't exceed max_tokens should_merge = ( current_texts and is_consecutive and current_tokens < min_tokens and current_tokens + table_tokens + 1 <= max_tokens # +1 for newline ) if current_texts and not should_merge: # Flush current record before starting new one records.append(create_record(current_texts, current_table_ids)) logger.debug( f"Created raw text record: {current_tokens} tokens, " f"tables {current_table_ids}" ) current_texts = [] current_table_ids = [] current_tokens = 0 # Add this table's text current_texts.extend(table_texts) current_table_ids.append(table.table_index) current_tokens = count_tokens("\n".join(current_texts)) last_table_index = table.table_index # Flush any remaining text if current_texts: records.append(create_record(current_texts, current_table_ids)) logger.debug( f"Created raw text record: {current_tokens} tokens, " f"tables {current_table_ids}" ) logger.info( f"Built {len(records)} raw text record(s) from " f"{len(sorted_tables)} tables on sheet '{sheet_name}'" ) return records except Exception as e: logger.error(f"Error building raw text records: {e}") return [] def _has_only_column_based_headers(self, flat_headers: list[str]) -> bool: """ Check if all headers are generic column-based (Column 0, Column 1, etc.). These are table fragments without proper headers and should be skipped. Args: flat_headers: List of flattened header strings Returns: True if all headers match "Column N" pattern """ if not flat_headers: return False # Check if all headers match "Column N" pattern column_pattern_count = sum( 1 for header in flat_headers if header.startswith("Column ") and header[7:].isdigit() ) # If more than 80% are "Column N" headers, skip this table threshold = len(flat_headers) * 0.8 return column_pattern_count > threshold def _is_record_non_empty(self, record: dict[str, Any]) -> bool: """ Check if a record has any non-empty content. Args: record: Record dictionary Returns: True if record has at least one non-empty cell """ content = record.get("content", {}) # Check if any value is non-empty after stripping whitespace for value in content.values(): if isinstance(value, str) and value.strip(): return True elif value is not None and not isinstance(value, str): # Non-string, non-None value (e.g., number) return True return False # Convenience function for easy imports def build_records_from_table( table: HtmlTable, headers: HeaderStructure, document_id: UUID, sheet_name: str | None = None ) -> list[dict[str, Any]]: """ Convenience function to build records from a table. Args: table: HtmlTable with extracted data headers: HeaderStructure with detected headers document_id: Document UUID sheet_name: Optional sheet name Returns: List of record dictionaries """ builder = RecordBuilder() return builder.build_records(table, headers, document_id, sheet_name) def build_raw_text_records_from_table( table: HtmlTable, document_id: UUID, sheet_name: str | None = None ) -> list[dict[str, Any]]: """ Convenience function to build raw text records from a single-column table. Args: table: HtmlTable with single-column raw text content document_id: Document UUID sheet_name: Optional sheet name Returns: List containing single record with concatenated text """ builder = RecordBuilder() return builder.build_raw_text_records(table, document_id, sheet_name)