alert-contract-expiring-hr/scripts/service/material_parser.py

"""材料包 PDF / Word 文档文本提取。

策略（自动降级）：
- .pdf : pypdf 直接抽 → 若 < 200 字视为扫描件，pdf2image + pytesseract OCR 兜底
- .docx: python-docx 抽取段落
- .doc : 调 antiword 兜底（不在 win32 默认安装，仅 best-effort）
- 其他 : 返回 ("", "failed")

依赖均为可选：导入失败返回空串，让上层降级到规则版。
"""

from __future__ import annotations

from pathlib import Path

from util.logging_setup import get_skill_logger

logger = get_skill_logger()

# 单个材料文件 OCR 后保留的最大长度，避免 LLM context 爆栈
MAX_TEXT_PER_FILE = 8000
# LLM Prompt 里允许的总材料摘要上限
MAX_TOTAL_SUMMARY = 16000


def _pdf_pypdf(path: Path) -> str:
    try:
        from pypdf import PdfReader
    except ImportError:
        return ""
    try:
        reader = PdfReader(str(path))
        out = []
        for page in reader.pages:
            t = page.extract_text() or ""
            if t:
                out.append(t)
        return "\n".join(out)
    except Exception as e:
        logger.warning("pypdf 解析失败 %s: %s", path.name, e)
        return ""


def _pdf_ocr(path: Path) -> str:
    try:
        from pdf2image import convert_from_path
        import pytesseract
    except ImportError:
        logger.warning("OCR 依赖未安装（pdf2image / pytesseract）")
        return ""
    try:
        images = convert_from_path(str(path), dpi=200)
        return "\n".join(
            pytesseract.image_to_string(img, lang="chi_sim+eng") for img in images
        )
    except Exception as e:
        logger.warning("OCR 解析失败 %s: %s", path.name, e)
        return ""


def _docx_text(path: Path) -> str:
    try:
        from docx import Document
    except ImportError:
        return ""
    try:
        doc = Document(str(path))
        return "\n".join(p.text for p in doc.paragraphs if p.text)
    except Exception as e:
        logger.warning("docx 解析失败 %s: %s", path.name, e)
        return ""


def _doc_antiword(path: Path) -> str:
    """legacy .doc：调系统 antiword 命令行。"""
    import shutil, subprocess

    if not shutil.which("antiword"):
        return ""
    try:
        r = subprocess.run(
            ["antiword", str(path)], capture_output=True, text=True, timeout=30
        )
        return r.stdout if r.returncode == 0 else ""
    except Exception as e:
        logger.warning("antiword 解析失败 %s: %s", path.name, e)
        return ""


def extract_text(path: Path) -> tuple[str, str]:
    """返回 (text, method)。method ∈ {'pypdf','docx','ocr','antiword','empty'}。

    text 为空时 method='empty'，调用方应据此判断是否降级到规则版。
    """
    if not path.exists():
        logger.warning("材料文件不存在: %s", path)
        return ("", "empty")

    suffix = path.suffix.lower()
    text = ""
    method = "empty"

    if suffix == ".pdf":
        text = _pdf_pypdf(path)
        method = "pypdf" if text else "empty"
        if len(text.strip()) < 200:
            ocr_text = _pdf_ocr(path)
            if ocr_text:
                text = ocr_text
                method = "ocr"
    elif suffix == ".docx":
        text = _docx_text(path)
        method = "docx" if text else "empty"
    elif suffix == ".doc":
        text = _doc_antiword(path)
        method = "antiword" if text else "empty"

    if text:
        text = text[:MAX_TEXT_PER_FILE]

    return (text, method)


def summarize(texts: list[tuple[str, str, str]]) -> tuple[str, list[str]]:
    """把多份材料的 (file_name, text, method) 汇总为 LLM 摘要。

    返回 (summary, methods)。summary 截断到 MAX_TOTAL_SUMMARY 字以内。
    """
    if not texts:
        return ("", [])

    parts: list[str] = []
    methods: list[str] = []
    for file_name, text, method in texts:
        if not text:
            continue
        methods.append(f"{file_name} ({method})")
        parts.append(f"### {file_name}\n{text}")

    joined = "\n\n".join(parts)
    if len(joined) > MAX_TOTAL_SUMMARY:
        joined = joined[:MAX_TOTAL_SUMMARY] + "\n...(以下材料内容过长已截断)"
    return (joined, methods)