baodan/api/insurance/ppt/extraction.py

"""PDF 提取服务 — 从 PDF 计划书中提取结构化数据。"""
import os
import re
import json
import time
import hashlib
import logging
from dataclasses import dataclass, field
from typing import Optional

logger = logging.getLogger(__name__)

CACHE_VERSION = 3


@dataclass
class ExtractionResult:
    pdf_path: str
    product_name: str
    plan_type: str  # savings/ci/iul
    status: str  # success/cached/error
    data: Optional[dict] = None
    usage: Optional[dict] = None
    error: Optional[str] = None
    duration_ms: float = 0


def infer_plan_type(raw: dict) -> str:
    """从 LLM 输出推断产品类型。"""
    t = str(raw.get("product_type", "")).lower()
    if "ci" in t or "critical" in t:
        return "ci"
    if "iul" in t or "universal" in t:
        return "iul"

    rows = raw.get("benefit_illustration", [])
    if not isinstance(rows, list):
        rows = []

    has_savings = any(
        r.get("total_surrender_value") is not None
        or r.get("guaranteed_cash_value") is not None
        or r.get("reversionary_bonus") is not None
        for r in rows if isinstance(r, dict)
    )
    has_ci = any(
        r.get("surrender_value_total") is not None
        or r.get("death_benefit_total") is not None
        for r in rows if isinstance(r, dict)
    )
    has_iul = any(
        r.get("cash_value") is not None
        or r.get("account_value") is not None
        for r in rows if isinstance(r, dict)
    )

    if has_iul:
        return "iul"
    if has_savings:
        return "savings"
    if has_ci:
        return "ci"

    policy = raw.get("policy", {})
    if isinstance(policy, dict):
        if policy.get("index_account_rate") is not None or policy.get("capital_partition") is not None:
            return "iul"
        if policy.get("sum_insured") is not None:
            return "ci"
    return "savings"


def _hash_file(file_path: str) -> str:
    """计算文件 SHA-256 哈希。"""
    h = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def _get_cache_path(pdf_path: str, cache_dir: str) -> str:
    """获取缓存文件路径。"""
    file_hash = _hash_file(pdf_path)
    return os.path.join(cache_dir, f"{file_hash}.json")


def _extract_pdf_text(pdf_path: str, max_chars: int = 30000) -> str:
    """使用 PyMuPDF 提取 PDF 文本。"""
    try:
        import fitz  # PyMuPDF
        doc = fitz.open(pdf_path)
        text_parts = []
        for page in doc:
            text_parts.append(page.get_text())
        doc.close()
        full_text = "\n".join(text_parts)
        if len(full_text) > max_chars:
            full_text = full_text[:max_chars]
        return full_text
    except ImportError:
        logger.warning("PyMuPDF 未安装，无法提取 PDF 文本")
        return ""


def _looks_corrupted(text: str) -> bool:
    """检测 PDF 文本是否乱码。"""
    if not text or len(text) < 50:
        return True
    # 计算坏字符比例
    bad_chars = sum(1 for c in text if ord(c) < 32 and c not in "\n\r\t")
    return bad_chars / len(text) > 0.1


class ExtractionOrchestrator:
    """PDF 提取编排器。"""

    def __init__(self, use_cache: bool = True, cache_dir: str = ".cache/insurance-ppt"):
        self.use_cache = use_cache
        self.cache_dir = cache_dir

    async def extract_plan(self, pdf_path: str, plan_type: str = "savings") -> ExtractionResult:
        """从 PDF 提取结构化数据。"""
        from insurance.ppt.llm_client import llm_client
        from insurance.ppt.prompts import (
            SAVINGS_PLAN_SYSTEM_PROMPT, CI_PLAN_SYSTEM_PROMPT, IUL_SYSTEM_PROMPT,
            build_savings_prompt,
        )

        start = time.time()
        abs_path = os.path.abspath(pdf_path)

        if not os.path.exists(abs_path):
            return ExtractionResult(
                pdf_path=abs_path, product_name="unknown", plan_type=plan_type,
                status="error", error="文件不存在",
                duration_ms=(time.time() - start) * 1000,
            )

        # 检查缓存
        if self.use_cache:
            cached = self._load_from_cache(abs_path)
            if cached:
                cached.duration_ms = (time.time() - start) * 1000
                return cached

        # 提取 PDF 文本
        pdf_text = _extract_pdf_text(abs_path)
        if not pdf_text:
            return ExtractionResult(
                pdf_path=abs_path, product_name="unknown", plan_type=plan_type,
                status="error", error="无法提取 PDF 文本",
                duration_ms=(time.time() - start) * 1000,
            )

        # 选择 prompt
        prompts = {
            "savings": SAVINGS_PLAN_SYSTEM_PROMPT,
            "ci": CI_PLAN_SYSTEM_PROMPT,
            "iul": IUL_SYSTEM_PROMPT,
        }
        system_prompt = prompts.get(plan_type, SAVINGS_PLAN_SYSTEM_PROMPT)

        # 调用 LLM
        try:
            data, response = await llm_client.structured_output(
                prompt=f"请从以下PDF文本中提取保险计划书数据：\n\n{pdf_text[:20000]}",
                system_prompt=system_prompt,
            )
        except Exception as e:
            return ExtractionResult(
                pdf_path=abs_path, product_name="unknown", plan_type=plan_type,
                status="error", error=f"LLM 调用失败: {e}",
                duration_ms=(time.time() - start) * 1000,
            )

        # 储蓄险数据修复：确保 total >= gcv
        if plan_type == "savings" and isinstance(data.get("benefit_illustration"), list):
            for row in data["benefit_illustration"]:
                if not isinstance(row, dict):
                    continue
                gcv = float(row.get("guaranteed_cash_value") or 0)
                rev = float(row.get("reversionary_bonus") or 0)
                term = float(row.get("terminal_dividend") or 0)
                total = float(row.get("total_surrender_value") or 0)
                if total < gcv:
                    row["total_surrender_value"] = gcv + rev + term

        # 推断产品类型
        detected_type = infer_plan_type(data)

        # 写入缓存
        if self.use_cache:
            self._save_to_cache(abs_path, data)

        product_name = data.get("product_name", "unknown")
        return ExtractionResult(
            pdf_path=abs_path, product_name=product_name,
            plan_type=detected_type, status="success", data=data,
            usage={"input": response.tokens.get("input", 0), "output": response.tokens.get("output", 0)} if response.tokens else None,
            duration_ms=(time.time() - start) * 1000,
        )

    async def extract_multiple(self, pdf_paths: list[str], plan_type: str = "savings") -> list[ExtractionResult]:
        """顺序提取多个 PDF。"""
        results = []
        for pdf_path in pdf_paths:
            logger.info(f"  📄 {os.path.basename(pdf_path)}...")
            results.append(await self.extract_plan(pdf_path, plan_type))
        return results

    def _load_from_cache(self, pdf_path: str) -> Optional[ExtractionResult]:
        """从缓存加载。"""
        try:
            cache_path = _get_cache_path(pdf_path, self.cache_dir)
            if not os.path.exists(cache_path):
                return None
            with open(cache_path, "r", encoding="utf-8") as f:
                raw = json.load(f)
            meta = raw.get("_meta", {})
            if meta.get("cacheVersion") != CACHE_VERSION:
                return None
            data = raw.get("_data", raw)
            product_name = data.get("product_name", "unknown")
            plan_type = infer_plan_type(data)
            return ExtractionResult(
                pdf_path=pdf_path, product_name=product_name,
                plan_type=plan_type, status="cached", data=data,
            )
        except Exception:
            return None

    def _save_to_cache(self, pdf_path: str, data: dict):
        """写入缓存。"""
        try:
            os.makedirs(self.cache_dir, exist_ok=True)
            file_hash = _hash_file(pdf_path)
            cache_data = {
                "_data": data,
                "_meta": {
                    "cacheVersion": CACHE_VERSION,
                    "originalFile": os.path.basename(pdf_path),
                    "extractedAt": __import__("datetime").datetime.now().isoformat(),
                    "fileHash": file_hash,
                },
            }
            cache_path = os.path.join(self.cache_dir, f"{file_hash}.json")
            with open(cache_path, "w", encoding="utf-8") as f:
                json.dump(cache_data, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.warning(f"缓存写入失败: {e}")

    async def extract_for_poster(self, filepath: str) -> dict:
        """提取海报所需的关键字段（精简 prompt，降低 LLM 成本）。

        返回:
            {age, gender, currency, sum_assured, premium_term,
             annual_premium, coverage_period, key_benefits}
        """
        from insurance.ppt.llm_client import llm_client

        abs_path = os.path.abspath(filepath)
        text = _extract_pdf_text(abs_path)
        if not text:
            raise ValueError("无法提取 PDF 文本")

        text = text[:6000]

        system_prompt = (
            "你是一位保险计划书解析专家。请从以下计划书内容中提取海报所需的关键字段。\n"
            "输出 JSON 格式（不要包含 markdown 代码块标记）：\n"
            '{"age": 35, "gender": "男", "currency": "USD", "sum_assured": 500000, '
            '"premium_term": 5, "annual_premium": 100000, "coverage_period": "终身", '
            '"key_benefits": ["身故赔偿", "全残保障"]}\n'
            "注意：数值用数字，不要带货币符号。"
        )

        result, _response = await llm_client.structured_output(
            text, system_prompt,
            schema={
                "type": "object",
                "properties": {
                    "age": {"type": "number"},
                    "gender": {"type": "string"},
                    "currency": {"type": "string"},
                    "sum_assured": {"type": "number"},
                    "premium_term": {"type": "number"},
                    "annual_premium": {"type": "number"},
                    "coverage_period": {"type": "string"},
                    "key_benefits": {"type": "array", "items": {"type": "string"}},
                },
                "required": ["age", "gender", "currency", "sum_assured", "annual_premium"],
            },
        )
        return result