vent_deep_agent/tools/pdf_tools.py

# -*- coding: utf-8 -*-
"""
PDF 文本提取工具模块

功能：
- 可选 MinerU 云端解析（高精度，需网络，由 MINERU_ENABLED 环境变量控制）
- 使用 PyPDF2 (pypdf) 提取 PDF 文本内容
- 若 PyPDF2 提取失败或内容为空，回退到 PaddleOCR 进行 OCR 识别
- 支持上传文件保存到临时目录
- 基于文件 MD5 哈希的缓存：同一 PDF 不重复提取
- MinerU 与本地解析使用独立缓存，互不干扰

解析优先级（当 MINERU_ENABLED=true 时）：
  MinerU 云端解析 → PyPDF2 本地提取 → PaddleOCR 回退

工具列表：
- extract_pdf_text: 从 PDF 文件提取文本（含缓存，MinerU → PyPDF2 → PaddleOCR）
- save_upload_file: 保存上传文件到临时目录
- calculate_file_hash: 计算文件 MD5 哈希
- clear_cache: 清除 PDF 提取缓存
"""

import hashlib
import os
import re
import time
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Optional

from dotenv import load_dotenv

load_dotenv()


# ============================================================
# 缓存目录
# ============================================================

# 缓存目录：项目根目录下的 data/pdf_cache/
_CACHE_DIR = Path(__file__).parent.parent / "data" / "pdf_cache"

# MinerU 缓存文件前缀（与本地解析缓存隔离）
_MINERU_CACHE_PREFIX = "mineru_"


# ============================================================
# 文件哈希计算
# ============================================================

def calculate_file_hash(file_path: str) -> str:
    """计算文件的 MD5 哈希值，作为缓存 key。

    同一文件（相同内容）产生相同哈希，实现"同一 PDF 不重复提取"。

    Args:
        file_path: 文件路径

    Returns:
        str: 32位 MD5 十六进制字符串
    """
    try:
        with open(file_path, 'rb') as f:
            file_data = f.read()
            return hashlib.md5(file_data).hexdigest()
    except Exception as e:
        print(f"[PDF] 计算文件哈希失败: {e}")
        # 失败时用时间戳生成唯一 key，确保不误命中缓存
        return hashlib.md5(str(datetime.now()).encode()).hexdigest()


def _get_cache_path(file_hash: str) -> Path:
    """获取缓存文件路径。

    Args:
        file_hash: 文件 MD5 哈希值

    Returns:
        Path: 缓存文件路径，如 data/pdf_cache/abc123.txt
    """
    return _CACHE_DIR / f"{file_hash}.txt"


def _read_cache(file_hash: str) -> Optional[str]:
    """从缓存中读取已提取的文本。

    Args:
        file_hash: 文件 MD5 哈希值

    Returns:
        str | None: 缓存文本，不存在或读取失败返回 None
    """
    cache_path = _get_cache_path(file_hash)
    if not cache_path.exists():
        return None
    try:
        text = cache_path.read_text(encoding='utf-8')
        if text.strip():
            print(f"[PDF] 命中缓存: {cache_path.name} ({len(text)} 字符)")
            return text
        return None
    except Exception as e:
        print(f"[PDF] 缓存读取失败: {e}")
        return None


def _write_cache(file_hash: str, text: str):
    """将提取的文本写入缓存。

    Args:
        file_hash: 文件 MD5 哈希值
        text: 提取的文本内容
    """
    try:
        _CACHE_DIR.mkdir(parents=True, exist_ok=True)
        cache_path = _get_cache_path(file_hash)
        cache_path.write_text(text, encoding='utf-8')
        print(f"[PDF] 已缓存: {cache_path.name} ({len(text)} 字符)")
    except Exception as e:
        print(f"[PDF] 缓存写入失败: {e}")


def clear_cache(file_hash: Optional[str] = None):
    """清除 PDF 提取缓存（同时清除本地解析缓存和 MinerU 缓存）。

    Args:
        file_hash: 指定要清除的哈希（可选），不传则清除全部缓存
    """
    if file_hash:
        # 清除本地缓存
        local_path = _get_cache_path(file_hash)
        if local_path.exists():
            local_path.unlink()
            print(f"[PDF] 已清除本地缓存: {local_path.name}")
        # 清除 MinerU 缓存
        mineru_path = _get_mineru_cache_path(file_hash)
        if mineru_path.exists():
            mineru_path.unlink()
            print(f"[PDF] 已清除 MinerU 缓存: {mineru_path.name}")
    else:
        if _CACHE_DIR.exists():
            count = 0
            for f in _CACHE_DIR.glob("*.txt"):
                f.unlink()
                count += 1
            print(f"[PDF] 已清除全部缓存 ({count} 个文件)")


# ============================================================
# 上传文件保存
# ============================================================

async def save_upload_file(upload_file, upload_dir: Optional[str] = None) -> str:
    """保存 FastAPI UploadFile 到临时目录，返回文件路径。

    Args:
        upload_file: FastAPI UploadFile 对象
        upload_dir: 保存目录，默认为系统临时目录下的 vent_review_pdfs

    Returns:
        str: 保存后的文件绝对路径
    """
    if upload_dir is None:
        upload_dir = os.path.join(tempfile.gettempdir(), "vent_review_pdfs")

    os.makedirs(upload_dir, exist_ok=True)

    # 保留原始文件名，避免冲突加时间戳
    import time
    original_name = upload_file.filename or "upload.pdf"
    safe_name = f"{int(time.time() * 1000)}_{original_name}"
    file_path = os.path.join(upload_dir, safe_name)

    content = await upload_file.read()
    with open(file_path, "wb") as f:
        f.write(content)

    return file_path


# ============================================================
# PDF 文本提取 - PyPDF2 方式
# ============================================================

def _extract_with_pypdf(file_path: str) -> Optional[str]:
    """使用 pypdf (PyPDF2) 提取 PDF 文本。

    Args:
        file_path: PDF 文件路径

    Returns:
        str | None: 提取的文本内容，失败返回 None
    """
    try:
        from pypdf import PdfReader

        reader = PdfReader(file_path)
        pages_text = []
        for i, page in enumerate(reader.pages):
            text = page.extract_text()
            if text:
                pages_text.append(f"--- 第 {i + 1} 页 ---\n{text.strip()}")

        if pages_text:
            return "\n\n".join(pages_text)
        return None

    except ImportError:
        print("[PDF] pypdf 未安装，跳过 PyPDF2 提取")
        return None
    except Exception as e:
        print(f"[PDF] PyPDF2 提取失败: {e}")
        return None


# ============================================================
# PDF 文本提取 - PaddleOCR 回退方式
# ============================================================

def _pdf_to_images(file_path: str, dpi: int = 150) -> list:
    """将 PDF 每页转换为 PIL Image 列表。

    Args:
        file_path: PDF 文件路径
        dpi: 图片分辨率，默认 150（平衡速度与精度）

    Returns:
        list[PIL.Image]: 图片列表
    """
    try:
        from pdf2image import convert_from_path
        return convert_from_path(file_path, dpi=dpi)
    except ImportError:
        print("[PDF] pdf2image 未安装，无法转换 PDF 为图片")
        return []
    except Exception as e:
        print(f"[PDF] PDF 转图片失败: {e}")
        return []


def _extract_with_paddleocr(file_path: str,
                            progress_callback: callable = None) -> Optional[str]:
    """使用 PaddleOCR 对 PDF 进行 OCR 识别。

    兼容 PaddleOCR v2.x (ocr/cls API) 和 v3.x (predict API)。

    先将 PDF 每页转为图片，再逐页 OCR。

    Args:
        file_path: PDF 文件路径
        progress_callback: 进度回调，签名 callback(page_num: int, total_pages: int)，每页扫描前调用

    Returns:
        str | None: OCR 识别结果，失败返回 None
    """
    try:
        from paddleocr import PaddleOCR

        # 检测 API 版本：v3.x 主要使用 predict()，ocr() 已废弃（deprecated）
        # 判断依据：v3.x 的 predict 方法签名中不含 cls 参数
        import inspect
        _is_v3 = False
        if hasattr(PaddleOCR, 'predict'):
            try:
                sig = inspect.signature(PaddleOCR.predict)
                _is_v3 = 'cls' not in sig.parameters
            except Exception:
                pass
        # 如果 predict 不存在，则为 v2.x（只有 ocr 方法）
        if not hasattr(PaddleOCR, 'predict'):
            _is_v3 = False

        # 初始化（兼容两种版本 + 多级加速策略）
        # 加速优先级：onnxruntime → paddle_dynamic+mkldnn → paddle_static+mkldnn → 纯CPU
        try:
            if _is_v3:
                # v3.x 加速参数说明：
                # - text_det_limit_side_len=960：检测时缩图边长上限，越小越快
                # - enable_mkldnn=True：Intel oneDNN 加速（paddle 引擎时有效）
                # - engine="onnxruntime"：ONNX Runtime 推理，CPU 上比 paddle 快 2-5x
                _v3_base = dict(
                    lang="ch",
                    use_doc_orientation_classify=False,
                    use_doc_unwarping=False,
                    use_textline_orientation=False,
                    text_det_limit_side_len=960,
                    cpu_threads=10,
                )
                # 策略0: onnxruntime（最快，绕过 PIR/onednn 问题）
                try:
                    ocr = PaddleOCR(**_v3_base, engine="onnxruntime")
                    print("[PDF] PaddleOCR 初始化: onnxruntime")
                except Exception:
                    # 策略1: paddle_dynamic + mkldnn
                    try:
                        ocr = PaddleOCR(**_v3_base, enable_mkldnn=True, engine="paddle_dynamic")
                        print("[PDF] PaddleOCR 初始化: paddle_dynamic + mkldnn")
                    except Exception:
                        # 策略2: paddle_static + mkldnn（可能触发 PIR 错误）
                        try:
                            ocr = PaddleOCR(**_v3_base, enable_mkldnn=True)
                            print("[PDF] PaddleOCR 初始化: paddle_static + mkldnn")
                        except Exception:
                            # 策略3: 回退到纯 CPU（稳定但慢）
                            try:
                                ocr = PaddleOCR(**_v3_base, enable_mkldnn=False)
                                print("[PDF] PaddleOCR 初始化: paddle_static (无加速)")
                            except Exception:
                                ocr = PaddleOCR(**_v3_base)
            else:
                # v2.x: 传统参数
                try:
                    ocr = PaddleOCR(use_angle_cls=True, lang="ch", show_log=False)
                except Exception:
                    ocr = PaddleOCR(use_angle_cls=True, lang="ch")
        except Exception:
            # 最后兜底：最少参数
            ocr = PaddleOCR(lang="ch")

        # 将 PDF 转为图片
        images = _pdf_to_images(file_path)
        if not images:
            print("[PDF] PDF 转图片失败，无法进行 OCR")
            return None

        # 逐页 OCR
        total = len(images)
        pages_text = []
        for i, img in enumerate(images):
            page_num = i + 1

            # 进度回调
            if progress_callback:
                try:
                    progress_callback(page_num, total)
                except Exception:
                    pass  # 回调异常不影响主流程

            import tempfile
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
                img.save(tmp.name)
                tmp_path = tmp.name

            try:
                if _is_v3:
                    # v3.x: predict() → list of OCRResult，传入加速参数
                    result = ocr.predict(
                        tmp_path,
                        text_det_limit_side_len=960,
                        text_rec_score_thresh=0.5,
                    )
                    page_text = _extract_text_from_v3_result(result)
                else:
                    # v2.x: ocr(img, cls=True) → nested list
                    result = ocr.ocr(tmp_path, cls=True)
                    page_text = _extract_text_from_v2_result(result)

                if page_text:
                    pages_text.append(f"--- 第 {i + 1} 页 (OCR) ---\n{page_text}")
            finally:
                try:
                    os.unlink(tmp_path)
                except OSError:
                    pass

        if pages_text:
            return "\n\n".join(pages_text)
        return None

    except ImportError:
        print("[PDF] PaddleOCR 未安装，跳过 OCR 提取")
        return None
    except Exception as e:
        print(f"[PDF] PaddleOCR 提取失败: {e}")
        import traceback
        traceback.print_exc()
        return None


def _extract_text_from_v2_result(result) -> str:
    """从 PaddleOCR v2.x 的 ocr() 返回结果中提取文本。

    v2 格式: [[[bbox, (text, confidence)], ...], ...]
    """
    if not result or not result[0]:
        return ""
    lines = []
    for line in result[0]:
        if line and len(line) >= 2:
            text = line[1][0] if isinstance(line[1], (list, tuple)) else str(line[1])
            lines.append(text)
    return "\n".join(lines)


def _extract_text_from_v3_result(result) -> str:
    """从 PaddleOCR v3.x 的 predict() 返回结果中提取文本。

    v3 格式: list of OCRResult，每个 OCRResult 可通过 ["rec_texts"] 或 str() 获取文本。
    """
    if not result:
        return ""
    lines = []
    for res in result:
        # 尝试多种方式提取文本
        if hasattr(res, '__getitem__'):
            try:
                rec_texts = res["rec_texts"]
                if rec_texts:
                    lines.extend(rec_texts if isinstance(rec_texts, list) else [str(rec_texts)])
                    continue
            except (KeyError, TypeError):
                pass
        # 尝试 str()
        text = str(res)
        if text:
            lines.append(text)
    return "\n".join(lines)


# ============================================================
# PDF 页数统计
# ============================================================

def _count_pdf_pages(file_path: str) -> int:
    """统计 PDF 文件的总页数。

    Args:
        file_path: PDF 文件路径

    Returns:
        int: 总页数，失败返回 0
    """
    try:
        from pypdf import PdfReader
        reader = PdfReader(file_path)
        return len(reader.pages)
    except Exception as e:
        print(f"[PDF] 无法获取 PDF 页数: {e}")
        return 0


# ============================================================
# MinerU 缓存（独立于本地解析缓存）
# ============================================================

def _get_mineru_cache_path(file_hash: str) -> Path:
    """获取 MinerU 专用缓存文件路径。

    Args:
        file_hash: 文件 MD5 哈希值

    Returns:
        Path: 缓存文件路径，如 data/pdf_cache/mineru_abc123.txt
    """
    return _CACHE_DIR / f"{_MINERU_CACHE_PREFIX}{file_hash}.txt"


def _read_mineru_cache(file_hash: str) -> Optional[str]:
    """从 MinerU 专用缓存中读取已提取的文本。

    Args:
        file_hash: 文件 MD5 哈希值

    Returns:
        str | None: 缓存文本，不存在或读取失败返回 None
    """
    cache_path = _get_mineru_cache_path(file_hash)
    if not cache_path.exists():
        return None
    try:
        text = cache_path.read_text(encoding='utf-8')
        if text.strip():
            print(f"[PDF] MinerU 缓存命中: {cache_path.name} ({len(text)} 字符)")
            return text
        return None
    except Exception as e:
        print(f"[PDF] MinerU 缓存读取失败: {e}")
        return None


def _write_mineru_cache(file_hash: str, text: str):
    """将 MinerU 提取的文本写入专用缓存。

    Args:
        file_hash: 文件 MD5 哈希值
        text: 提取的文本内容
    """
    try:
        _CACHE_DIR.mkdir(parents=True, exist_ok=True)
        cache_path = _get_mineru_cache_path(file_hash)
        cache_path.write_text(text, encoding='utf-8')
        print(f"[PDF] MinerU 已缓存: {cache_path.name} ({len(text)} 字符)")
    except Exception as e:
        print(f"[PDF] MinerU 缓存写入失败: {e}")


# ============================================================
# MinerU 文本清洗
# ============================================================

def _clean_mineru_text(text: str) -> str:
    """清洗 MinerU 输出的文本。

    - 去除每行首尾空白
    - 合并多余空行（连续 3 个以上空行压缩为 2 个）

    Args:
        text: MinerU 原始输出文本

    Returns:
        str: 清洗后的文本
    """
    if not text:
        return ""
    text = "\n".join(line.strip() for line in text.splitlines())
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


# ============================================================
# MinerU 云端 PDF 解析
# ============================================================

def _extract_with_mineru(file_path: str,
                          force_refresh: bool = False,
                          use_cache: bool = True) -> Optional[str]:
    """使用 MinerU 云端服务解析 PDF（高精度，需要网络）。

    解析流程：
    1. 检查 MinerU 专用缓存（mineru_{hash}.txt）
    2. 若缓存命中且非强制刷新，直接返回
    3. 通过 langchain_mineru.MinerULoader 分片加载 PDF
    4. 清洗每页文本并拼接
    5. 写入 MinerU 专用缓存

    Args:
        file_path: PDF 文件路径
        force_refresh: 是否强制重新提取（忽略缓存）
        use_cache: 是否使用缓存

    Returns:
        str | None: 提取的文本内容，失败返回 None
    """
    mineru_token = os.getenv("MINERU_TOKEN", "").strip()
    mineru_max_pages = int(os.getenv("MINERU_MAX_PAGES", "50"))

    if not mineru_token:
        print("[PDF] MINERU_TOKEN 未配置，跳过 MinerU 解析")
        return None

    file_hash = calculate_file_hash(file_path)

    # ── 检查 MinerU 专用缓存 ──
    if use_cache and not force_refresh:
        cached = _read_mineru_cache(file_hash)
        if cached is not None:
            return cached

    # ── 统计总页数 ──
    total_pages = _count_pdf_pages(file_path)
    if total_pages == 0:
        print("[PDF] 无法获取 PDF 页数，跳过 MinerU 解析")
        return None

    print(f"[PDF] PDF总页数：{total_pages}（MinerU 云端解析）")

    start_time = time.time()
    all_text_parts = []

    try:
        from langchain_mineru import MinerULoader

        for start in range(1, total_pages + 1, mineru_max_pages):
            end = min(start + mineru_max_pages - 1, total_pages)
            print(f"[PDF] MinerU 正在加载分片：{start} ~ {end} 页")

            loader = MinerULoader(
                source=file_path,
                mode="precision",
                token=mineru_token,
                pages=f"{start}-{end}",
            )
            docs = loader.load()
            for idx, doc in enumerate(docs):
                cleaned = _clean_mineru_text(doc.page_content)
                if cleaned:
                    all_text_parts.append(
                        f"--- 第 {start + idx} 页 (MinerU) ---\n{cleaned}"
                    )

        if not all_text_parts:
            print("[PDF] MinerU 未提取到任何文本内容")
            return None

        full_text = "\n\n".join(all_text_parts)

        # ── 写入 MinerU 专用缓存 ──
        if use_cache:
            _write_mineru_cache(file_hash, full_text)

        elapsed = round(time.time() - start_time, 2)
        print(f"[PDF] MinerU 解析完成 | 总页数：{total_pages} | 总耗时：{elapsed} 秒")
        return full_text

    except ImportError:
        print("[PDF] langchain-mineru 未安装，跳过 MinerU 解析")
        return None
    except Exception as e:
        print(f"[PDF] MinerU 解析失败: {e}")
        import traceback
        traceback.print_exc()
        return None


# ============================================================
# 主提取函数
# ============================================================

def extract_pdf_text(file_path: str, min_text_length: int = 50,
                     use_cache: bool = True, force_refresh: bool = False,
                     progress_callback: callable = None) -> str:
    """从 PDF 文件提取文本。

    提取策略（按优先级）：
    0. 若 MINERU_ENABLED=true，优先使用 MinerU 云端解析（高精度，需网络）
       - 失败或内容不足则自动回退到本地解析
    1. 计算文件 MD5 哈希，若本地缓存命中且非强制刷新，直接返回
    2. 尝试 PyPDF2 (pypdf) 直接提取文本
    3. 若提取内容为空或长度不足 min_text_length，回退到 PaddleOCR
    4. 提取成功后将结果写入缓存
    5. 若全部失败，返回错误信息

    Args:
        file_path: PDF 文件路径
        min_text_length: 最小有效文本长度，低于此值认为提取失败
        use_cache: 是否使用缓存，默认 True
        force_refresh: 是否强制重新提取（忽略缓存），默认 False
        progress_callback: PaddleOCR 进度回调（MinerU 不使用此回调）

    Returns:
        str: 提取的文本内容
    """
    # 验证文件存在
    if not os.path.exists(file_path):
        return f"[错误] 文件不存在: {file_path}"

    # 计算文件哈希
    file_hash = calculate_file_hash(file_path)

    # ── 步骤0: MinerU 云端解析（优先，如果启用）──
    use_mineru = os.getenv("MINERU_ENABLED", "false").strip().lower() == "true"
    if use_mineru:
        print("[PDF] MinerU 已启用，优先使用云端解析...")
        mineru_text = _extract_with_mineru(
            file_path,
            force_refresh=force_refresh,
            use_cache=use_cache,
        )
        if mineru_text and len(mineru_text.strip()) >= min_text_length:
            print(f"[PDF] MinerU 解析成功，共 {len(mineru_text)} 字符")
            return mineru_text
        print("[PDF] MinerU 解析失败或内容不足，回退到本地解析...")

    # ── 步骤1: 检查本地缓存 ──
    if use_cache and not force_refresh:
        cached = _read_cache(file_hash)
        if cached:
            return cached

    # ── 步骤2: PyPDF2 提取 ──
    print(f"[PDF] 正在使用 PyPDF2 提取: {file_path}")
    text = _extract_with_pypdf(file_path)

    if text and len(text.strip()) >= min_text_length:
        print(f"[PDF] PyPDF2 提取成功，共 {len(text)} 字符")
        if use_cache:
            _write_cache(file_hash, text)
        return text

    # ── 步骤3: PaddleOCR 回退 ──
    print(f"[PDF] PyPDF2 提取不足 ({len(text or '')} 字符)，回退到 PaddleOCR...")
    ocr_text = _extract_with_paddleocr(file_path, progress_callback=progress_callback)

    if ocr_text and len(ocr_text.strip()) >= min_text_length:
        print(f"[PDF] PaddleOCR 提取成功，共 {len(ocr_text)} 字符")
        if use_cache:
            _write_cache(file_hash, ocr_text)
        return ocr_text

    # ── 步骤4: 全部失败 ──
    if text and text.strip():
        print(f"[PDF] 警告: 提取内容较短 ({len(text)} 字符)，但已是全部可用内容")
        if use_cache:
            _write_cache(file_hash, text)
        return text

    return "[错误] 无法从 PDF 中提取文本内容。文件可能为扫描件且 OCR 不可用，或文件已损坏。"


# ============================================================
# 快速提取（仅 PyPDF2，不触发 OCR）
# ============================================================

def extract_pdf_text_fast(file_path: str) -> str:
    """仅使用 PyPDF2 快速提取 PDF 文本（不触发 OCR）。

    用于 LLM 预览场景，避免 OCR 耗时。

    Args:
        file_path: PDF 文件路径

    Returns:
        str: 提取的文本内容
    """
    text = _extract_with_pypdf(file_path)
    return text or "[提示] 该 PDF 可能为扫描件，PyPDF2 无法提取文本，请使用完整 OCR 模式。"