kgrag/mulfilesparse.py

#!/usr/bin/env python3
"""
批量 OCR 解析工具
- 扫描输入目录下所有 PDF / DOCX / DOC / PPTX / PPT
- DOCX / DOC / PPTX / PPT 先用 LibreOffice 转换为 PDF，再上传解析
- 结果 content_list 以 JSON 格式保存到输出目录
- 已处理的文件自动跳过，支持多线程并发
"""
import json
import os
import shutil
import sys
import subprocess
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import requests

# ============ 配置（按需修改）============
API_URL      = "http://192.168.0.46:59988/analyze-pdf"
INPUT_DIR    = "/app/testfiles/input"   # 输入文件夹
OUTPUT_DIR   = "/app/testfiles/output"   # 输出文件夹
MAX_WORKERS  = 3     # 并发请求数（避免打垮 API 服务）
API_TIMEOUT  = 300   # API 请求超时（秒）
CONV_TIMEOUT = 120   # LibreOffice 转换超时（秒）
SKIP_EXISTING = True # 跳过已存在的输出文件

# Updated SUPPORTED extensions to include ppt and pptx
SUPPORTED = {".pdf", ".docx", ".doc", ".pptx", ".ppt"}


# ============ DOCX/DOC/PPTX/PPT → PDF 转换 ============
def convert_to_pdf(input_path: str, output_dir: str) -> str:
    """使用 LibreOffice 将 DOCX/DOC/PPTX/PPT 转换为 PDF，返回生成的 PDF 路径。

    每次调用使用独立的 profile 目录，避免多线程并发时 LibreOffice 互相抢锁导致静默失败。
    """
    profile_dir = tempfile.mkdtemp(prefix="lo_profile_")
    try:
        cmd = [
            "libreoffice",
            f"-env:UserInstallation=file://{profile_dir}",
            "--headless", "--convert-to", "pdf",
            "--outdir", output_dir, input_path,
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=CONV_TIMEOUT)
        if result.returncode != 0:
            detail = result.stderr.strip() or result.stdout.strip() or f"退出码 {result.returncode}"
            raise RuntimeError(f"LibreOffice 转换失败: {detail}")
    finally:
        shutil.rmtree(profile_dir, ignore_errors=True)

    stem = Path(input_path).stem
    pdf_path = os.path.join(output_dir, f"{stem}.pdf")
    if not os.path.exists(pdf_path):
        raise RuntimeError(f"转换后找不到 PDF: {pdf_path}")
    return pdf_path


# ============ API 调用 ============
def analyze_pdf(pdf_path: str) -> list:
    """上传 PDF 到 OCR API，返回 content_list。"""
    with open(pdf_path, "rb") as f:
        files = {"file": (os.path.basename(pdf_path), f, "application/pdf")}
        resp = requests.post(API_URL, files=files, timeout=API_TIMEOUT)

    if resp.status_code != 200:
        try:
            detail = resp.json()
        except Exception:
            detail = resp.text[:300]
        raise RuntimeError(f"HTTP {resp.status_code}: {detail}")

    try:
        result = resp.json()
    except Exception:
        raise RuntimeError(f"响应非 JSON: {resp.text[:300]}")

    return result.get("data", {}).get("content_list", [])


# ============ 单文件处理 ============
def process_file(input_path: str, output_dir: str, temp_dir: str) -> dict:
    """
    处理单个文件，返回状态字典。
    status: "success" | "skipped" | "empty" | "error"
    """
    stem = Path(input_path).stem
    ext  = Path(input_path).suffix.lower()
    output_path = os.path.join(output_dir, f"{stem}_content_list.json")

    if SKIP_EXISTING and os.path.exists(output_path):
        return {"file": stem, "status": "skipped"}

    pdf_path  = input_path
    converted = False
    try:
        # DOCX / DOC / PPTX / PPT 先转 PDF
        # Updated condition to include ppt and pptx
        if ext in (".docx", ".doc", ".pptx", ".ppt"):
            pdf_path  = convert_to_pdf(input_path, temp_dir)
            converted = True

        content_list = analyze_pdf(pdf_path)

        if not content_list:
            return {"file": stem, "status": "empty"}

        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(content_list, f, ensure_ascii=False, indent=2)

        return {"file": stem, "status": "success", "items": len(content_list)}

    except Exception as e:
        return {"file": stem, "status": "error", "reason": str(e)}

    finally:
        if converted and os.path.exists(pdf_path):
            os.remove(pdf_path)


# ============ 主程序 ============
def main():
    input_dir  = sys.argv[1] if len(sys.argv) > 1 else INPUT_DIR
    output_dir = sys.argv[2] if len(sys.argv) > 2 else OUTPUT_DIR

    if not os.path.isdir(input_dir):
        print(f"❌ 输入目录不存在: {input_dir}")
        sys.exit(1)

    os.makedirs(output_dir, exist_ok=True)

    # 收集并分类文件
    all_files = [p for p in Path(input_dir).iterdir() if p.is_file()]
    files     = [str(p) for p in all_files if p.suffix.lower() in SUPPORTED]
    skipped   = [p.name for p in all_files if p.suffix.lower() not in SUPPORTED]

    if skipped:
        print(f"⏭️  跳过 {len(skipped)} 个不支持的文件:")
        for name in skipped:
            print(f"     · {name}")
        print()

    if not files:
        print(f"⚠️  目录下没有 PDF/DOCX/DOC/PPTX/PPT 文件: {input_dir}")
        return

    print(f"📂 输入目录: {input_dir}")
    print(f"💾 输出目录: {output_dir}")
    print(f"📄 共 {len(files)} 个文件，并发数 {MAX_WORKERS}\n")

    stats = {"success": 0, "skipped": 0, "empty": 0, "error": 0}

    with tempfile.TemporaryDirectory() as temp_dir:
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            futures = {
                executor.submit(process_file, fp, output_dir, temp_dir): fp
                for fp in files
            }
            for i, future in enumerate(as_completed(futures), 1):
                res = future.result()
                status = res["status"]
                stats[status] += 1

                icon = {"success": "✅", "skipped": "⏭️", "empty": "⚠️", "error": "❌"}[status]
                detail = ""
                if status == "success":
                    detail = f"({res['items']} items)"
                elif status == "error":
                    detail = f"- {res['reason']}"

                print(f"[{i:>3}/{len(files)}] {icon} {res['file']} {detail}")

    print(f"\n🎉 完成！成功 {stats['success']} | 跳过 {stats['skipped']} "
          f"| 空结果 {stats['empty']} | 失败 {stats['error']}")


if __name__ == "__main__":
    main()