#!/usr/bin/env python3 """Manual test script for the multi-pass extraction pipeline. This script demonstrates how to use the ExtractionPipeline directly to process an Excel document and extract structured records with: PASS 1: Symbol Detection PASS 2: Table Extraction PASS 3: Symbol Resolution PASS 4: Cross-Reference Detection """ from pathlib import Path import redis as sync_redis from src.db.session import get_sync_session from src.db import models as db_models from src.extraction.pipeline import ExtractionPipeline # V1 pipeline from src.extraction_v2.pipeline import ExtractionPipelineV2 # V2 pipeline from src.core.config import settings def main(): """Run the extraction pipeline on a test file.""" # Choose a test file (you can change this to any Excel file) test_files = [ # Test fixtures ("tests/fixtures/sample_excel/test_multi_table_crossref.xlsx", "Multi-table + cross-refs (recommended)"), ("tests/fixtures/sample_excel/test_multipass_comprehensive.xlsx", "Multi-pass test (symbols + cross-refs)"), ("tests/fixtures/sample_excel/test_standard_table.xlsx", "Standard table (simple)"), ("tests/fixtures/sample_excel/test_multiple_tables.xlsx", "Multiple tables"), ("tests/fixtures/sample_excel/test_mixed_depth_headers.xlsx", "Mixed depth headers"), ("tests/fixtures/sample_excel/test_complex_merges.xlsx", "Complex merged cells"), # Customer documents (paths from filesystem) ("CustomerDocument/IFф╗ХцзШцЫ╕(BB-CASTAR)/focus_1.xlsx", "Customer: BB-CASTAR"), ("CustomerDocument/уВ╖уВЩуГзуГХуВЩф╕Ашжз/1.2 JP1шинхоЪщаЕчЫо.xlsx", "Customer: JP1 Settings"), ("CustomerDocument/уВ╖уВЩуГзуГХуВЩф╕Ашжз/1.1 JP1уВ╣уВ▒уВ╖уВЩуГеуГ╝уГлф╕Ашжз.xlsx", "Customer: JP1 Schedule"), ("CustomerDocument/чФ╗щЭвшиншиИч╖и/02_чФ╗щЭвф╗ХцзШ/01_щА▓цНЧчобчРЖцйЯшГ╜/хИеч┤Щ_01.02.01.03.05_х╖еф║ЛщА▓цНЧщБЕх╗╢уГБуВзуГГуВп.xlsx", "Customer: Screen Design"), ("CustomerDocument/ц╡БщАЪщаЕчЫоф╗ХцзШцЫ╕/01.01.01_ц╡БщАЪщаЕчЫоф╗ХцзШцЫ╕уАРуВвуГ│уГПуВЩуГ│уГИуВЩуГлцКЬч▓ЛчЙИуАС.xlsx", "Customer: Distribution Spec"), ("CustomerDocument/IFф╗ХцзШцЫ╕(уВвуГ│уГПуВЩуГ│уГИуВЩуГл)/ч╡▒хРИHHC/01.03.03-06_SOAP FaultуВдуГ│уВ┐уГХуВзуГ╝уВ╣щЫ╗цЦЗ.xlsx", "Customer: SOAP Fault IF"), ('/Users/nguyen/Work/DSOL/CustomerDocument/IFф╗ХцзШцЫ╕(уВвуГ│уГПуВЩуГ│уГИуВЩуГл)/уВкуГ╝уВ┐уВЩхИ╢х╛б/1 хЕЙуВвуГ│уГПуВЩуГ│уГИуВЩуГлценхЛЩцФпцП┤уВ╖уВ╣уГЖуГая╝ИхИЖцЮРч│╗я╝Й/1.12 уГХуВбуВдуГлчХкхП╖уБиуГлуГ╝уГИуВ┐уВпуВЩхп╛х┐Ьшби.xls', "Customer: Complex Document"), ] print("=" * 70) print("MULTI-PASS EXTRACTION PIPELINE TEST") print("=" * 70) print("\nAvailable test files:") for i, (file, desc) in enumerate(test_files, 1): exists = "[OK]" if Path(file).exists() else "[--]" print(f" {i}. {exists} {desc}") print(f" {file}") choice = input(f"\nSelect a file (1-{len(test_files)}) or enter a custom path: ").strip() try: choice_num = int(choice) if 1 <= choice_num <= len(test_files): test_file = Path(test_files[choice_num - 1][0]) else: print("Invalid choice, using default file") test_file = Path(test_files[0][0]) except ValueError: # User entered a custom path test_file = Path(choice) if not test_file.exists(): print(f"Error: File not found: {test_file}") return print(f"\nUsing file: {test_file}") print("=" * 70) # Select pipeline version print("\nSelect extraction pipeline:") print(" 1. V1 Pipeline (openpyxl-based, original)") print(" 2. V2 Pipeline (Docling-based, with LLM headers)") print(" 3. Both (compare side-by-side)") pipeline_choice = input("Pipeline (1-3, default=2): ").strip() or "2" if pipeline_choice not in ["1", "2", "3"]: print("Invalid choice, using V2 pipeline") pipeline_choice = "2" print("=" * 70) # Create a test document in database with get_sync_session() as session: # Create test document doc = db_models.Document( filename=test_file.name, file_path=str(test_file.absolute()), status="pending" ) session.add(doc) session.commit() doc_id = doc.id print(f"\n[OK] Created test document: {doc_id}") # Optional: Connect to Redis for progress tracking redis_client = None try: redis_client = sync_redis.from_url(settings.redis_url, decode_responses=False) print("[OK] Redis connected (progress tracking enabled)") except Exception as e: print(f"[WARN] Redis not available: {e}") print(" (Proceeding without progress tracking)") print(f"\n[START] Running extraction pipeline(s)...") print("=" * 70) try: v1_summary = None v2_result = None # Run V1 pipeline if pipeline_choice in ["1", "3"]: print("\n[V1] Running openpyxl-based extraction pipeline...") print("-" * 70) v1_pipeline = ExtractionPipeline(db_session=session, redis_client=redis_client) v1_summary = v1_pipeline.process_document(doc_id) print(f"[V1] ✓ Complete: {v1_summary.records_extracted} records extracted") # Run V2 pipeline if pipeline_choice in ["2", "3"]: print("\n[V2] Running Docling-based extraction pipeline...") print("-" * 70) v2_pipeline = ExtractionPipelineV2() v2_result = v2_pipeline.process_document_with_large_table_detection_sync(doc_id, str(test_file.absolute())) if v2_result.success: print(f"[V2] ✓ Complete: {v2_result.record_count} records extracted") # Save V2 records to database from src.workers.extraction_v2_tasks import _save_v2_records, _mark_document_completed, _index_vectors_sync _save_v2_records(session, v2_result) _mark_document_completed(session, str(doc_id), v2_result) # Index vectors in Milvus print("[V2] Indexing vectors in Milvus...") try: vector_count = _index_vectors_sync(v2_result.records, doc_id, str(test_file.absolute())) print(f"[V2] ✓ Indexed {vector_count} vectors in Milvus") except Exception as e: print(f"[V2] ✗ Vector indexing failed: {e}") else: print(f"[V2] ✗ Failed: {v2_result.error}") print("[V2] Note: In production, this would trigger V1 fallback") # Print results print("\n" + "=" * 70) print("EXTRACTION COMPLETE!") print("=" * 70) # Show comparison if both were run if pipeline_choice == "3" and v1_summary and v2_result and v2_result.success: print("\n--- V1 vs V2 Comparison ---") print(f" Records extracted:") print(f" V1 (openpyxl): {v1_summary.records_extracted}") print(f" V2 (Docling): {v2_result.record_count}") print(f" Tables detected:") print(f" V1: {v1_summary.tables_detected}") print(f" V2: {v2_result.table_count}") print(f" Processing time:") print(f" V1: {v1_summary.duration_ms}ms ({v1_summary.duration_ms/1000:.2f}s)") print(f" V2: (not tracked separately)") # Determine which to show detailed results for summary = v1_summary # Default to V1 for detailed output print("\n Showing detailed V1 results below (V2 records also saved to DB)") elif v2_result and v2_result.success: # V2 only - create a summary-like object for display summary = type('obj', (object,), { 'symbols_detected': 0, 'sheets_processed': v2_result.table_count, 'tables_detected': v2_result.table_count, 'records_extracted': v2_result.record_count, 'symbols_resolved': 0, 'unresolved_symbols': 0, 'cross_references_found': 0, 'cross_references_resolved': 0, 'cross_references_unresolved': 0, 'duration_ms': 0, 'warnings': [] })() else: # V1 only or V2 failed summary = v1_summary # Display detailed results if pipeline_choice == "1" or (pipeline_choice == "3" and v1_summary): # Full V1 multi-pass results print("\n--- PASS 1: Symbol Detection (V1 only) ---") print(f" Symbols detected: {summary.symbols_detected}") print("\n--- PASS 2: Table Extraction ---") print(f" Sheets processed: {summary.sheets_processed}") print(f" Tables detected: {summary.tables_detected}") print(f" Records extracted: {summary.records_extracted}") print("\n--- PASS 3: Symbol Resolution (V1 only) ---") print(f" Symbols resolved: {summary.symbols_resolved}") print(f" Unresolved symbols: {summary.unresolved_symbols}") print("\n--- PASS 4: Cross-Reference Detection (V1 only) ---") print(f" Cross-references found: {summary.cross_references_found}") print(f" Cross-references resolved: {summary.cross_references_resolved}") print(f" Cross-references unresolved: {summary.cross_references_unresolved}") print("\n--- Performance ---") print(f" Duration: {summary.duration_ms}ms ({summary.duration_ms/1000:.2f}s)") else: # V2 results (simpler - just table extraction) print("\n--- V2 Extraction Results ---") print(f" Tables extracted: {summary.tables_detected}") print(f" Records extracted: {summary.records_extracted}") print(" Note: V2 pipeline focuses on table extraction with LLM-based headers") print(" Symbol resolution and cross-refs would run in subsequent passes") if summary.warnings: print(f"\n--- Warnings ({len(summary.warnings)}) ---") for i, warning in enumerate(summary.warnings, 1): print(f" {i}. {warning}") # Check database print(f"\n--- Database Status ---") print(f" Document status: {doc.status}") print(f" Record count: {doc.record_count}") print(f" Sheet count: {doc.sheet_count}") print(f" Processed at: {doc.processed_at}") # Query symbol dictionaries symbols = session.query(db_models.SymbolDictionary).filter_by( document_id=doc_id ).all() print(f"\n--- Symbol Dictionaries in Database ---") print(f" Total symbols: {len(symbols)}") if symbols: print(f" All symbols:") for i, sym in enumerate(symbols, 1): print(f" {i}. '{sym.symbol}' = '{sym.meaning}' (sheet: {sym.sheet_name})") # Query extracted records records = session.query(db_models.ExtractedRecord).filter_by( document_id=doc_id ).order_by(db_models.ExtractedRecord.sheet_name, db_models.ExtractedRecord.row_number).all() print(f"\n--- Extracted Records in Database ---") print(f" Total records: {len(records)}") if records: # Group by sheet sheets = {} for record in records: if record.sheet_name not in sheets: sheets[record.sheet_name] = [] sheets[record.sheet_name].append(record) print(f" Records by sheet:") for sheet_name, sheet_records in sheets.items(): print(f" - {sheet_name}: {len(sheet_records)} records") # Show all records print(f"\n--- All Extracted Records ---") for i, record in enumerate(records, 1): title_str = f", Table: {record.table_title}" if record.table_title else "" print(f"\n Record {i} (Sheet: {record.sheet_name}{title_str}, Row: {record.row_number}):") print(f" Headers: {record.headers}") print(f" Content: {record.content}") if record.resolved_content: print(f" Resolved: {record.resolved_content}") # Query cross-references (through source_record relationship) cross_refs = session.query(db_models.CrossReference).join( db_models.ExtractedRecord, db_models.CrossReference.source_record_id == db_models.ExtractedRecord.id ).filter( db_models.ExtractedRecord.document_id == doc_id ).all() print(f"\n--- Cross-References in Database ---") print(f" Total cross-references: {len(cross_refs)}") if cross_refs: resolved = [cr for cr in cross_refs if cr.resolved] unresolved = [cr for cr in cross_refs if not cr.resolved] print(f" Resolved: {len(resolved)}") print(f" Unresolved: {len(unresolved)}") print(f"\n All cross-references:") for i, cr in enumerate(cross_refs, 1): status = "[RESOLVED]" if cr.resolved else "[UNRESOLVED]" print(f" {i}. {status} '{cr.reference_text}'") print(f" Type: {cr.reference_type}") # Show Redis progress (if available) if redis_client: try: key = f"processing:documents:{doc_id}" progress = redis_client.get(key) if progress: import json progress_data = json.loads(progress) print(f"\n--- Redis Progress Data ---") print(f" {progress_data}") else: print(f"\n--- Redis: No progress data (expired or cleared) ---") except Exception as e: print(f"\n--- Redis: Could not fetch progress: {e} ---") print("\n" + "=" * 70) print("Test completed successfully!") print("=" * 70) except Exception as e: print("\n" + "=" * 70) print("EXTRACTION FAILED!") print("=" * 70) print(f"\nError: {e}") print(f"\nDocument status: {doc.status}") if doc.error_message: print(f"Error message: {doc.error_message}") import traceback print("\nFull traceback:") traceback.print_exc() if __name__ == "__main__": main()