#!/usr/bin/env python3 """Batch extraction script for processing multiple documents with V2 pipeline. This script: 1. Scans a folder for all supported documents: - Excel: .xlsx, .xls, .xlsm - Word: .docx, .doc, .docm - PDF: .pdf - PowerPoint: .pptx, .ppt, .pptm 2. Tracks all files in a JSON status file 3. Processes each file through ExtractionPipelineV2 4. Updates status (pending/processing/completed/failed) in real-time 5. Saves detailed results and error messages 6. Generates comprehensive metrics by document type and format """ import json import sys import asyncio from pathlib import Path from datetime import datetime from typing import List, Dict, Any, Optional, Tuple from enum import Enum import traceback import zipfile from src.db.session import get_sync_session from src.db import models as db_models from src.extraction_v2.pipeline import ExtractionPipelineV2 from src.workers.extraction_v2_tasks import _save_v2_records, _mark_document_completed, _index_vectors_sync # Default folder to scan DEFAULT_FOLDER = "CustomerDocument/" STATUS_FILE = "batch_extraction_status.json" class ErrorCategory(str, Enum): """Categorize extraction errors for better diagnostics.""" # Transient errors (retry possible) NETWORK_ERROR = "network_error" TIMEOUT = "timeout" RATE_LIMIT = "rate_limit" SERVICE_UNAVAILABLE = "service_unavailable" # File-related errors (permanent) FILE_CORRUPTED = "file_corrupted" FILE_TOO_LARGE = "file_too_large" UNSUPPORTED_FORMAT = "unsupported_format" PASSWORD_PROTECTED = "password_protected" # Pipeline errors (investigate) CONVERSION_FAILED = "conversion_failed" PARSING_FAILED = "parsing_failed" LLM_ERROR = "llm_error" VALIDATION_FAILED = "validation_failed" # System errors OUT_OF_MEMORY = "out_of_memory" DISK_FULL = "disk_full" # Unknown UNKNOWN = "unknown" def categorize_error(error: Exception) -> ErrorCategory: """Categorize an exception into error category. Args: error: The exception to categorize Returns: ErrorCategory: The categorized error type """ error_msg = str(error).lower() # Network/transient if "timeout" in error_msg or "timed out" in error_msg: return ErrorCategory.TIMEOUT if "rate limit" in error_msg or "429" in error_msg: return ErrorCategory.RATE_LIMIT if "connection" in error_msg or "network" in error_msg: return ErrorCategory.NETWORK_ERROR if "503" in error_msg or "unavailable" in error_msg: return ErrorCategory.SERVICE_UNAVAILABLE # File errors if "corrupted" in error_msg or "damaged" in error_msg or "badzipfile" in error_msg: return ErrorCategory.FILE_CORRUPTED if "too large" in error_msg or "file size" in error_msg: return ErrorCategory.FILE_TOO_LARGE if "password" in error_msg or "encrypted" in error_msg: return ErrorCategory.PASSWORD_PROTECTED if "unsupported" in error_msg or "invalid format" in error_msg: return ErrorCategory.UNSUPPORTED_FORMAT # Pipeline errors if "conversion" in error_msg or "docling" in error_msg: return ErrorCategory.CONVERSION_FAILED if "parsing" in error_msg or "html" in error_msg: return ErrorCategory.PARSING_FAILED if "openai" in error_msg or "llm" in error_msg or "gpt" in error_msg: return ErrorCategory.LLM_ERROR # System errors if "memory" in error_msg or "memoryerror" in error_msg: return ErrorCategory.OUT_OF_MEMORY if "disk" in error_msg or "no space" in error_msg: return ErrorCategory.DISK_FULL return ErrorCategory.UNKNOWN class RetryConfig: """Configuration for retry logic.""" max_retries: int = 3 base_delay: float = 1.0 # seconds max_delay: float = 30.0 exponential_backoff: bool = True # Which error categories should be retried retryable_categories = { ErrorCategory.NETWORK_ERROR, ErrorCategory.TIMEOUT, ErrorCategory.RATE_LIMIT, ErrorCategory.SERVICE_UNAVAILABLE, } def process_with_retry(pipeline, file_path: Path, doc_id: str, config: RetryConfig = None) -> tuple[bool, Any, Optional[str], Optional[str], int]: """Process file with automatic retry for transient errors. Args: pipeline: ExtractionPipelineV2 instance file_path: Path to file doc_id: Document ID config: Retry configuration (uses defaults if None) Returns: Tuple of (success, result, error_msg, error_category, attempts) """ if config is None: config = RetryConfig() last_error = None last_category = None attempts = 0 for attempt in range(config.max_retries + 1): attempts = attempt + 1 try: # Use large table detection for Excel files result = pipeline.process_document_with_large_table_detection_sync(doc_id, str(file_path.absolute())) return True, result, None, None, attempts except Exception as e: last_error = e last_category = categorize_error(e) print(f"[RETRY] Attempt {attempts} failed: {last_category.value} - {str(e)[:100]}") # Don't retry if not retryable if last_category not in config.retryable_categories: print(f"[RETRY] Error not retryable: {last_category.value}") break # Don't retry if max attempts reached if attempt >= config.max_retries: print(f"[RETRY] Max retry attempts ({config.max_retries}) reached") break # Calculate backoff delay if config.exponential_backoff: delay = min(config.base_delay * (2 ** attempt), config.max_delay) else: delay = config.base_delay print(f"[RETRY] Retrying after {delay:.1f}s delay...") import time time.sleep(delay) # All retries exhausted error_msg = f"{type(last_error).__name__}: {str(last_error)}" return False, None, error_msg, last_category.value if last_category else None, attempts def validate_excel_file(file_path: Path) -> Tuple[bool, Optional[str]]: """Validate file before processing. Args: file_path: Path to Excel file to validate Returns: Tuple of (is_valid, error_message). error_message is None if valid. """ # Check exists and readable if not file_path.exists(): return False, "File not found" # Check size (< 100MB) size_mb = file_path.stat().st_size / (1024 * 1024) if size_mb > 100: return False, f"File too large: {size_mb:.1f}MB" # Verify valid Excel format ext = file_path.suffix.lower() if ext not in ['.xlsx', '.xls', '.xlsm']: return False, f"Unsupported format: {ext}" # For .xlsx/.xlsm: verify valid ZIP structure if ext in ['.xlsx', '.xlsm']: try: with zipfile.ZipFile(file_path, 'r') as zf: if 'xl/workbook.xml' not in zf.namelist(): return False, "Invalid Excel structure" except zipfile.BadZipFile: return False, "Corrupted file" return True, None class BatchProcessor: """Process multiple Excel and Word files with V2 pipeline and track status.""" def __init__(self, folder_path: str, status_file: str = STATUS_FILE): self.folder_path = Path(folder_path) self.status_file = Path(status_file) self.status_data: Dict[str, Any] = { "folder": str(self.folder_path), "started_at": None, "completed_at": None, "files": {} } def find_all_files(self) -> List[Path]: """Find all supported document files in the folder (recursively).""" all_files = [] # Excel patterns excel_patterns = ["**/*.xlsx", "**/*.xlsm", "**/*.xls"] # Word patterns word_patterns = ["**/*.docx", "**/*.doc", "**/*.docm"] # PDF patterns pdf_patterns = ["**/*.pdf"] # PowerPoint patterns pptx_patterns = ["**/*.pptx", "**/*.ppt", "**/*.pptm"] patterns = excel_patterns + word_patterns + pdf_patterns + pptx_patterns for pattern in patterns: matches = list(self.folder_path.glob(pattern)) print(f"[DEBUG] Pattern {pattern}: found {len(matches)} files") all_files.extend(matches) return sorted(all_files) def find_excel_files(self) -> List[Path]: """Find all Excel files in the folder (recursively). DEPRECATED: Use find_all_files() instead for Excel + Word support. """ excel_files = [] # Support both .xlsx and .xls for pattern in ["**/*.xlsx", "**/*.xlsm", "**/*.xls"]: excel_files.extend(self.folder_path.glob(pattern)) return sorted(excel_files) def _categorize_files(self, files: List[Path]) -> Dict[str, Dict[str, List[Path]]]: """Categorize files by document type and format. Returns: Dict with structure: { "excel": {".xlsx": [...], ".xls": [...], ".xlsm": [...]}, "word": {".docx": [...], ".doc": [...], ".docm": [...]}, "pdf": {".pdf": [...]}, "powerpoint": {".pptx": [...], ".ppt": [...], ".pptm": [...]} } """ categorized = { "excel": { ".xlsx": [], ".xls": [], ".xlsm": [] }, "word": { ".docx": [], ".doc": [], ".docm": [] }, "pdf": { ".pdf": [] }, "powerpoint": { ".pptx": [], ".ppt": [], ".pptm": [] } } for file in files: ext = file.suffix.lower() if ext in categorized["excel"]: categorized["excel"][ext].append(file) elif ext in categorized["word"]: categorized["word"][ext].append(file) elif ext in categorized["pdf"]: categorized["pdf"][ext].append(file) elif ext in categorized["powerpoint"]: categorized["powerpoint"][ext].append(file) return categorized def generate_metrics_summary(self) -> Dict[str, Any]: """Generate comprehensive metrics from status data. Returns metrics broken down by: - Overall (total, completed, failed, success rate) - By type (Excel vs Word) - By format (.xlsx, .xls, .xlsm, .docx, .doc, .docm) """ files = list(self.status_data["files"].values()) # Overall metrics total = len(files) completed = sum(1 for f in files if f["status"] == "completed") failed = sum(1 for f in files if f["status"] == "failed") success_rate = (completed / total * 100) if total > 0 else 0 # By format by_format = {} for file_info in files: filename = file_info["filename"] ext = Path(filename).suffix.lower() if ext not in by_format: by_format[ext] = {"total": 0, "success": 0, "failed": 0} by_format[ext]["total"] += 1 if file_info["status"] == "completed": by_format[ext]["success"] += 1 elif file_info["status"] == "failed": by_format[ext]["failed"] += 1 # By document type (Excel, Word, PDF, PowerPoint) excel_formats = [".xlsx", ".xls", ".xlsm"] word_formats = [".docx", ".doc", ".docm"] pdf_formats = [".pdf"] pptx_formats = [".pptx", ".ppt", ".pptm"] excel_files = [f for f in files if Path(f["filename"]).suffix.lower() in excel_formats] word_files = [f for f in files if Path(f["filename"]).suffix.lower() in word_formats] pdf_files = [f for f in files if Path(f["filename"]).suffix.lower() in pdf_formats] pptx_files = [f for f in files if Path(f["filename"]).suffix.lower() in pptx_formats] excel_success = sum(1 for f in excel_files if f["status"] == "completed") word_success = sum(1 for f in word_files if f["status"] == "completed") pdf_success = sum(1 for f in pdf_files if f["status"] == "completed") pptx_success = sum(1 for f in pptx_files if f["status"] == "completed") excel_rate = (excel_success / len(excel_files) * 100) if excel_files else 0 word_rate = (word_success / len(word_files) * 100) if word_files else 0 pdf_rate = (pdf_success / len(pdf_files) * 100) if pdf_files else 0 pptx_rate = (pptx_success / len(pptx_files) * 100) if pptx_files else 0 return { "summary": { "total_files": total, "completed": completed, "failed": failed, "success_rate": round(success_rate, 2) }, "by_type": { "excel": { "total": len(excel_files), "success": excel_success, "success_rate": round(excel_rate, 2) }, "word": { "total": len(word_files), "success": word_success, "success_rate": round(word_rate, 2) }, "pdf": { "total": len(pdf_files), "success": pdf_success, "success_rate": round(pdf_rate, 2) }, "powerpoint": { "total": len(pptx_files), "success": pptx_success, "success_rate": round(pptx_rate, 2) } }, "by_format": by_format } def load_status(self): """Load existing status file if it exists.""" if self.status_file.exists(): with open(self.status_file, 'r', encoding='utf-8') as f: self.status_data = json.load(f) print(f"[INFO] Loaded existing status from {self.status_file}") else: print(f"[INFO] Starting fresh (no existing status file)") def save_status(self): """Save current status to JSON file.""" with open(self.status_file, 'w', encoding='utf-8') as f: json.dump(self.status_data, f, indent=2, ensure_ascii=False) def initialize_files(self, all_files: List[Path]): """Initialize or update file list in status.""" for file_path in all_files: file_key = str(file_path) # Only add if not already tracked if file_key not in self.status_data["files"]: # Determine file type for appropriate metrics file_ext = file_path.suffix.lower() # Categorize file type if file_ext in [".xlsx", ".xls", ".xlsm"]: file_type = "excel" elif file_ext in [".docx", ".doc", ".docm"]: file_type = "word" elif file_ext == ".pdf": file_type = "pdf" elif file_ext in [".pptx", ".ppt", ".pptm"]: file_type = "powerpoint" else: file_type = "unknown" file_record = { "status": "pending", "filename": file_path.name, "file_type": file_type, "format": file_ext, "size_bytes": file_path.stat().st_size if file_path.exists() else 0, "discovered_at": datetime.now().isoformat(), "processed_at": None, "document_id": None, "record_count": 0, "table_count": 0, "vector_count": 0, "duration_seconds": None, "error": None, "error_category": None, "attempts": 0, "pipeline_version": None } # Add chunking-specific fields for text documents (Word, PDF, PowerPoint) if file_type in ["word", "pdf", "powerpoint"]: file_record["chunks"] = 0 file_record["avg_chunk_size"] = 0 self.status_data["files"][file_key] = file_record self.save_status() print(f"[INFO] Initialized {len(all_files)} files in status tracker") def update_file_status(self, file_path: Path, updates: Dict[str, Any]): """Update status for a specific file.""" file_key = str(file_path) if file_key in self.status_data["files"]: self.status_data["files"][file_key].update(updates) self.save_status() def calculate_and_save_metrics(self): """Calculate comprehensive metrics and add to status data.""" files = self.status_data.get("files", {}) # Count by status completed = sum(1 for f in files.values() if f.get("status") == "completed") processing = sum(1 for f in files.values() if f.get("status") == "processing") pending = sum(1 for f in files.values() if f.get("status") == "pending") failed = sum(1 for f in files.values() if f.get("status") == "failed") total = len(files) # Success rate success_rate = (completed / total * 100) if total > 0 else 0.0 # Count by pipeline version v2_success = sum(1 for f in files.values() if f.get("pipeline_version") == "v2") v1_fallback = sum(1 for f in files.values() if f.get("pipeline_version") == "v1") # Count by format by_format = {} for file_key, file_data in files.items(): ext = Path(file_key).suffix.lower() if ext not in by_format: by_format[ext] = {"total": 0, "success": 0, "failed": 0} by_format[ext]["total"] += 1 if file_data.get("status") == "completed": by_format[ext]["success"] += 1 elif file_data.get("status") == "failed": by_format[ext]["failed"] += 1 # Update status data with summary and metrics self.status_data["summary"] = { "completed": completed, "processing": processing, "pending": pending, "failed": failed, "success_rate": round(success_rate, 2) } self.status_data["metrics"] = { "v2_success": v2_success, "v1_fallback": v1_fallback, "by_format": by_format } self.save_status() def process_file(self, file_path: Path) -> bool: """Process a single file through V2 pipeline. Returns: bool: True if successful, False if failed """ file_key = str(file_path) file_status = self.status_data["files"][file_key] # Skip if already completed if file_status["status"] == "completed": print(f"[SKIP] {file_path.name} - already completed") return True print(f"\n{'='*70}") print(f"[PROCESSING] {file_path.name}") print(f"{'='*70}") # Pre-flight validation is_valid, validation_error = validate_excel_file(file_path) if not is_valid: print(f"[VALIDATION] ✗ File validation failed: {validation_error}") self.update_file_status(file_path, { "status": "failed", "error": f"Validation failed: {validation_error}" }) return False print(f"[VALIDATION] ✓ File validation passed") # Update status to processing self.update_file_status(file_path, { "status": "processing", "processed_at": datetime.now().isoformat() }) start_time = datetime.now() try: with get_sync_session() as session: # Create document in database doc = db_models.Document( filename=file_path.name, file_path=str(file_path.absolute()), status="pending" ) session.add(doc) session.commit() doc_id = doc.id print(f"[OK] Created document: {doc_id}") # Update status with document ID self.update_file_status(file_path, {"document_id": str(doc_id)}) # Run V2 pipeline (route based on file type) file_ext = file_path.suffix.lower() v2_pipeline = ExtractionPipelineV2() if file_ext in [".xlsx", ".xls", ".xlsm"]: print(f"[V2] Running Excel extraction with large table detection...") v2_result = v2_pipeline.process_document_with_large_table_detection_sync(doc_id, str(file_path.absolute())) elif file_ext in [".docx", ".doc", ".docm"]: print(f"[V2] Running Docling-based Word extraction with semantic chunking...") # Use async wrapper for Word documents import asyncio try: loop = asyncio.get_event_loop() v2_result = asyncio.run(v2_pipeline.process_word_document(doc_id, str(file_path.absolute()))) except RuntimeError: v2_result = asyncio.run(v2_pipeline.process_word_document(doc_id, str(file_path.absolute()))) elif file_ext == ".pdf": print(f"[V2] Running PDF extraction with semantic chunking (OCR disabled)...") v2_result = v2_pipeline.process_pdf_document_sync(doc_id, str(file_path.absolute())) elif file_ext in [".pptx", ".ppt", ".pptm"]: print(f"[V2] Running PowerPoint extraction with semantic chunking...") v2_result = v2_pipeline.process_pptx_document_sync(doc_id, str(file_path.absolute())) else: raise ValueError(f"Unsupported file format: {file_ext}") # Use retry wrapper success, v2_result, error_msg, error_category, attempts = process_with_retry( v2_pipeline, file_path, doc_id ) if success and v2_result.success: print(f"[V2] ✓ Extraction complete: {v2_result.record_count} records") # Save records to database print(f"[V2] Saving records to database...") _save_v2_records(session, v2_result) _mark_document_completed(session, str(doc_id), v2_result) print(f"[V2] ✓ Saved to database") # Index vectors in Milvus vector_count = 0 try: print(f"[V2] Indexing vectors in Milvus...") vector_count = _index_vectors_sync(v2_result.records, doc_id, str(file_path.absolute())) print(f"[V2] ✓ Indexed {vector_count} vectors") except Exception as e: print(f"[V2] ✗ Vector indexing failed: {e}") # Calculate duration duration = (datetime.now() - start_time).total_seconds() # Build status update status_update = { "status": "completed", "pipeline_version": "v2", "record_count": v2_result.record_count, "table_count": v2_result.table_count, "vector_count": vector_count, "duration_seconds": duration, "attempts": attempts, "error": None, "error_category": None } # Add chunk-specific metrics for text documents (Word, PDF, PowerPoint) if file_ext in [".docx", ".doc", ".docm", ".pdf", ".pptx", ".ppt", ".pptm"]: status_update["chunks"] = v2_result.record_count # For text docs, records are chunks if v2_result.records: # Try different possible text field names avg_size = 0 for r in v2_result.records: text_content = r.get("text") or r.get("content") or r.get("chunk_text") or "" avg_size += len(text_content) if v2_result.records: avg_size = avg_size / len(v2_result.records) status_update["avg_chunk_size"] = int(avg_size) # Update status - success self.update_file_status(file_path, status_update) print(f"[OK] {file_path.name} completed in {duration:.2f}s ({attempts} attempts)") return True else: # V2 failed final_error = error_msg or (v2_result.error if v2_result else "Unknown error") final_category = error_category or categorize_error(Exception(final_error)).value print(f"[V2] ✗ Extraction failed: {final_error}") duration = (datetime.now() - start_time).total_seconds() self.update_file_status(file_path, { "status": "failed", "pipeline_version": None, "duration_seconds": duration, "attempts": attempts, "error": final_error, "error_category": final_category }) return False except Exception as e: error_msg = f"{type(e).__name__}: {str(e)}" error_cat = categorize_error(e) print(f"[ERROR] {error_msg}") print(f"[ERROR] Category: {error_cat.value}") # Print traceback for debugging traceback.print_exc() duration = (datetime.now() - start_time).total_seconds() self.update_file_status(file_path, { "status": "failed", "pipeline_version": None, "duration_seconds": duration, "attempts": 1, "error": error_msg, "error_category": error_cat.value }) return False def run(self, resume: bool = False): """Run the batch processing. Args: resume: If True, resume from existing status file. If False, start fresh. """ print(f"{'='*70}") print(f"BATCH EXTRACTION - V2 PIPELINE") print(f"{'='*70}") print(f"Folder: {self.folder_path}") print(f"Status file: {self.status_file}") print(f"Mode: {'RESUME' if resume else 'FRESH START'}") print(f"{'='*70}\n") # Load existing status if resuming if resume: self.load_status() # Find all supported document files print(f"[SCAN] Finding document files in {self.folder_path}...") all_files = self.find_all_files() if not all_files: print(f"[WARN] No supported document files found in {self.folder_path}") return # Categorize files for reporting categorized = self._categorize_files(all_files) excel_count = sum(len(files) for files in categorized["excel"].values()) word_count = sum(len(files) for files in categorized["word"].values()) pdf_count = sum(len(files) for files in categorized["pdf"].values()) pptx_count = sum(len(files) for files in categorized["powerpoint"].values()) print(f"[SCAN] Found {len(all_files)} files total:") print(f" Excel: {excel_count} files") for ext, files in categorized["excel"].items(): if files: print(f" {ext}: {len(files)}") print(f" Word: {word_count} files") for ext, files in categorized["word"].items(): if files: print(f" {ext}: {len(files)}") print(f" PDF: {pdf_count} files") for ext, files in categorized["pdf"].items(): if files: print(f" {ext}: {len(files)}") print(f" PowerPoint: {pptx_count} files") for ext, files in categorized["powerpoint"].items(): if files: print(f" {ext}: {len(files)}") print() # Initialize file tracking self.initialize_files(all_files) # Start batch processing self.status_data["started_at"] = datetime.now().isoformat() self.save_status() # Process each file total = len(all_files) success_count = 0 failed_count = 0 skipped_count = 0 for i, file_path in enumerate(all_files, 1): print(f"\n[{i}/{total}] Processing: {file_path.name}") # Check if already completed file_key = str(file_path) if self.status_data["files"][file_key]["status"] == "completed": skipped_count += 1 success_count += 1 # Count as success since it's already done continue # Process the file success = self.process_file(file_path) if success: success_count += 1 else: failed_count += 1 # Print progress print(f"\n[PROGRESS] {i}/{total} files processed") print(f" ✓ Success: {success_count}") print(f" ✗ Failed: {failed_count}") if skipped_count > 0: print(f" → Skipped: {skipped_count} (already completed)") # Mark completion self.status_data["completed_at"] = datetime.now().isoformat() # Calculate metrics and add summary self.calculate_and_save_metrics() # Print final summary print(f"\n{'='*70}") print(f"BATCH PROCESSING COMPLETE!") print(f"{'='*70}") print(f"Total files: {total}") print(f" ✓ Success: {success_count}") print(f" ✗ Failed: {failed_count}") if skipped_count > 0: print(f" → Skipped: {skipped_count} (already completed)") if "summary" in self.status_data: summary = self.status_data["summary"] print(f"\nSuccess Rate: {summary.get('success_rate', 0):.1f}%") if "metrics" in self.status_data: metrics = self.status_data["metrics"] print(f" V2 Success: {metrics.get('v2_success', 0)}") if metrics.get('v1_fallback', 0) > 0: print(f" V1 Fallback: {metrics.get('v1_fallback', 0)}") print(f"\nStatus saved to: {self.status_file}") print(f"{'='*70}\n") def main(): """Main entry point.""" import argparse parser = argparse.ArgumentParser( description="Batch process documents (Excel, Word, PDF, PowerPoint) with V2 extraction pipeline" ) parser.add_argument( "--folder", default=DEFAULT_FOLDER, help=f"Folder to scan for supported document files (default: {DEFAULT_FOLDER})" ) parser.add_argument( "--status-file", default=STATUS_FILE, help=f"JSON file to track status (default: {STATUS_FILE})" ) parser.add_argument( "--resume", action="store_true", help="Resume from existing status file (skip completed files)" ) args = parser.parse_args() # Validate folder exists folder_path = Path(args.folder) if not folder_path.exists(): print(f"[ERROR] Folder not found: {folder_path}") sys.exit(1) # Create processor and run processor = BatchProcessor( folder_path=str(folder_path), status_file=args.status_file ) try: processor.run(resume=args.resume) except KeyboardInterrupt: print("\n\n[INTERRUPTED] Batch processing interrupted by user") print(f"[INFO] Progress saved to: {processor.status_file}") print(f"[INFO] Resume with: python {sys.argv[0]} --resume") sys.exit(1) if __name__ == "__main__": main()