#!/usr/bin/env python3 """ Multi-format document extractor. Supports: PDF (with tables), DOCX, PPTX, XLSX, MD, TXT Recursively scans subdirectories. Usage: python extract_docs.py # Reads paths from config.py (agent-config.yml) """ import os import re import sys from pathlib import Path from config import cfg RAW_DIR = cfg.raw_dir OUT_DIR = cfg.processed_dir OUT_DIR.mkdir(parents=True, exist_ok=True) SKIP_EXTENSIONS = {'.mp4', '.zip', '.rar', '.7z', '.avi', '.mov', '.wav', '.mp3', '.DS_Store'} success = 0 failed = 0 skipped = 0 tables_found = 0 def safe_filename(name, max_len=200): """Convert path to safe flat filename""" name = name.replace('/', '_').replace('\\', '_') name = re.sub(r'[<>:"|?*]', '_', name) return name[:max_len] # ── PDF extraction (with table support) ── def extract_pdf(pdf_path): global tables_found import fitz doc = fitz.open(str(pdf_path)) text_parts = [] for page_num in range(len(doc)): page = doc[page_num] # Detect tables page_tables = page.find_tables() table_rects = [] if page_tables.tables: for tab in page_tables.tables: table_data = tab.extract() if table_data and len(table_data) > 1: md_table = _table_to_markdown(table_data) if md_table: text_parts.append(f'[TABLE_START]\n{md_table}\n[TABLE_END]') tables_found += 1 table_rects.append(tab.bbox) # Get text excluding table areas if not table_rects: page_text = page.get_text() if page_text.strip(): text_parts.append(page_text) else: blocks = page.get_text("blocks") for block in blocks: bx0, by0, bx1, by1, block_text = block[:5] if not block_text.strip(): continue in_table = False for trect in table_rects: tx0, ty0, tx1, ty1 = trect bcx, bcy = (bx0 + bx1) / 2, (by0 + by1) / 2 if tx0 <= bcx <= tx1 and ty0 <= bcy <= ty1: in_table = True break if not in_table: text_parts.append(block_text) doc.close() return '\n'.join(text_parts) def _table_to_markdown(table_data): if not table_data or not table_data[0]: return '' rows = [] for row in table_data: cells = [str(c).replace('\n', ' ').strip() if c else '' for c in row] rows.append(cells) if not rows: return '' max_cols = max(len(r) for r in rows) for r in rows: while len(r) < max_cols: r.append('') lines = ['| ' + ' | '.join(rows[0]) + ' |'] lines.append('| ' + ' | '.join(['---'] * max_cols) + ' |') for row in rows[1:]: lines.append('| ' + ' | '.join(row) + ' |') return '\n'.join(lines) # ── DOCX extraction ── def extract_docx(docx_path): from docx import Document doc = Document(str(docx_path)) parts = [] for para in doc.paragraphs: text = para.text.strip() if text: # Preserve heading structure if para.style and para.style.name.startswith('Heading'): level = para.style.name.replace('Heading ', '').replace('Heading', '1') try: level = int(level) except: level = 1 parts.append('#' * level + ' ' + text) else: parts.append(text) # Extract tables for table in doc.tables: rows = [] for row in table.rows: cells = [cell.text.strip().replace('\n', ' ') for cell in row.cells] rows.append(cells) if rows: md = _table_to_markdown(rows) if md: parts.append(f'[TABLE_START]\n{md}\n[TABLE_END]') return '\n\n'.join(parts) # ── PPTX extraction ── def extract_pptx(pptx_path): from pptx import Presentation prs = Presentation(str(pptx_path)) parts = [] for slide_num, slide in enumerate(prs.slides, 1): slide_texts = [f'--- Slide {slide_num} ---'] for shape in slide.shapes: if shape.has_text_frame: for para in shape.text_frame.paragraphs: text = para.text.strip() if text: slide_texts.append(text) if shape.has_table: rows = [] for row in shape.table.rows: cells = [cell.text.strip().replace('\n', ' ') for cell in row.cells] rows.append(cells) if rows: md = _table_to_markdown(rows) if md: slide_texts.append(f'[TABLE_START]\n{md}\n[TABLE_END]') if len(slide_texts) > 1: parts.append('\n'.join(slide_texts)) return '\n\n'.join(parts) # ── XLSX extraction ── def extract_xlsx(xlsx_path): from openpyxl import load_workbook wb = load_workbook(str(xlsx_path), read_only=True, data_only=True) parts = [] for sheet_name in wb.sheetnames: ws = wb[sheet_name] rows = [] for row in ws.iter_rows(values_only=True): cells = [str(c).strip() if c is not None else '' for c in row] if any(cells): rows.append(cells) if rows: parts.append(f'## Sheet: {sheet_name}') md = _table_to_markdown(rows) if md: parts.append(f'[TABLE_START]\n{md}\n[TABLE_END]') wb.close() return '\n\n'.join(parts) # ── TXT/MD extraction ── def extract_text(txt_path): with open(txt_path, 'r', encoding='utf-8', errors='ignore') as f: return f.read() # ── Main ── def main(): global success, failed, skipped # Recursively find all files all_files = sorted(RAW_DIR.rglob('*')) files = [f for f in all_files if f.is_file()] total = len(files) print(f"Raw dir: {RAW_DIR}") print(f"Output dir: {OUT_DIR}") print(f"Total files: {total}") print("=" * 60) for i, file_path in enumerate(files): ext = file_path.suffix.lower() # Relative path for output filename rel_path = file_path.relative_to(RAW_DIR) out_name = safe_filename(str(rel_path).replace(ext, '')) + '.txt' out_path = OUT_DIR / out_name # Skip already processed if out_path.exists() and out_path.stat().st_size > 0: skipped += 1 continue # Skip unsupported if ext in SKIP_EXTENSIONS or file_path.name.startswith('.'): skipped += 1 continue try: if ext == '.pdf': text = extract_pdf(file_path) elif ext == '.docx': text = extract_docx(file_path) elif ext == '.pptx': text = extract_pptx(file_path) elif ext == '.xlsx': text = extract_xlsx(file_path) elif ext in ('.txt', '.md'): text = extract_text(file_path) else: skipped += 1 continue if text and text.strip(): with open(out_path, 'w', encoding='utf-8') as f: f.write(f'# Source: {rel_path}\n\n') f.write(text) success += 1 else: failed += 1 except Exception as e: failed += 1 print(f" ERROR [{i+1}/{total}] {rel_path}: {e}", file=sys.stderr) if (i + 1) % 50 == 0: print(f" Progress: {i+1}/{total} (ok={success} skip={skipped} fail={failed})") print(f"\nDone! success={success} skipped={skipped} failed={failed} tables={tables_found}") if __name__ == '__main__': main()