lambdagent/agentexample/qaagent67lambda/scripts/extract_docs.py

#!/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()